1use std::future::Future;
4use std::pin::Pin;
5use std::sync::{Arc, Mutex};
6use std::time::Duration;
7
8use serde::{Deserialize, Serialize};
9use serde_json::json;
10use tracing::{info, info_span};
11
12use crate::error::Error;
13use crate::llm::types::{TokenUsage, ToolDefinition};
14use crate::llm::{BoxedProvider, LlmProvider};
15use crate::tool::{Tool, ToolOutput};
16use crate::types::DispatchMode;
17
18use crate::memory::Memory;
19
20use crate::knowledge::KnowledgeBase;
21
22use crate::tool::builtins::OnQuestion;
23
24use super::blackboard::{Blackboard, InMemoryBlackboard};
25use super::blackboard_tools::blackboard_tools;
26use super::context::ContextStrategy;
27use super::events::{AgentEvent, OnEvent};
28use super::guardrail::Guardrail;
29use super::{AgentOutput, AgentRunner};
30
31#[derive(Clone)]
33pub(crate) struct SubAgentDef {
34 pub(crate) name: String,
35 pub(crate) description: String,
36 pub(crate) system_prompt: String,
37 pub(crate) tools: Vec<Arc<dyn Tool>>,
38 pub(crate) context_strategy: Option<ContextStrategy>,
39 pub(crate) summarize_threshold: Option<u32>,
40 pub(crate) tool_timeout: Option<Duration>,
41 pub(crate) max_tool_output_bytes: Option<usize>,
42 pub(crate) max_turns: Option<usize>,
44 pub(crate) max_tokens: Option<u32>,
46 pub(crate) response_schema: Option<serde_json::Value>,
48 pub(crate) guardrails: Vec<Arc<dyn Guardrail>>,
50 pub(crate) run_timeout: Option<Duration>,
52 pub(crate) provider_override: Option<Arc<BoxedProvider>>,
55 pub(crate) reasoning_effort: Option<crate::llm::types::ReasoningEffort>,
57 pub(crate) enable_reflection: Option<bool>,
59 pub(crate) tool_output_compression_threshold: Option<usize>,
61 pub(crate) max_tools_per_turn: Option<usize>,
63 pub(crate) tool_profile: Option<super::tool_filter::ToolProfile>,
65 pub(crate) max_identical_tool_calls: Option<u32>,
67 pub(crate) max_fuzzy_identical_tool_calls: Option<u32>,
69 pub(crate) max_tool_calls_per_turn: Option<u32>,
71 pub(crate) session_prune_config: Option<crate::agent::pruner::SessionPruneConfig>,
73 pub(crate) enable_recursive_summarization: Option<bool>,
75 pub(crate) reflection_threshold: Option<u32>,
77 pub(crate) consolidate_on_exit: Option<bool>,
79 pub(crate) workspace: Option<std::path::PathBuf>,
81 pub(crate) max_total_tokens: Option<u64>,
83 pub(crate) audit_trail: Option<Arc<dyn super::audit::AuditTrail>>,
85 pub(crate) audit_user_id: Option<String>,
87 pub(crate) audit_tenant_id: Option<String>,
89 pub(crate) audit_delegation_chain: Vec<String>,
91}
92
93impl std::fmt::Debug for SubAgentDef {
94 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
95 f.debug_struct("SubAgentDef")
96 .field("name", &self.name)
97 .field("description", &self.description)
98 .field("tools_count", &self.tools.len())
99 .finish()
100 }
101}
102
103impl SubAgentDef {
104 pub(crate) fn new(
107 name: impl Into<String>,
108 description: impl Into<String>,
109 system_prompt: impl Into<String>,
110 ) -> Self {
111 Self {
112 name: name.into(),
113 description: description.into(),
114 system_prompt: system_prompt.into(),
115 tools: vec![],
116 context_strategy: None,
117 summarize_threshold: None,
118 tool_timeout: None,
119 max_tool_output_bytes: None,
120 max_turns: None,
121 max_tokens: None,
122 response_schema: None,
123 run_timeout: None,
124 guardrails: vec![],
125 provider_override: None,
126 reasoning_effort: None,
127 enable_reflection: None,
128 tool_output_compression_threshold: None,
129 max_tools_per_turn: None,
130 tool_profile: None,
131 max_identical_tool_calls: None,
132 max_fuzzy_identical_tool_calls: None,
133 max_tool_calls_per_turn: None,
134 session_prune_config: None,
135 enable_recursive_summarization: None,
136 reflection_threshold: None,
137 consolidate_on_exit: None,
138 workspace: None,
139 max_total_tokens: None,
140 audit_trail: None,
141 audit_user_id: None,
142 audit_tenant_id: None,
143 audit_delegation_chain: Vec::new(),
144 }
145 }
146}
147
148impl From<SubAgentConfig> for SubAgentDef {
149 fn from(def: SubAgentConfig) -> Self {
150 Self {
151 name: def.name,
152 description: def.description,
153 system_prompt: def.system_prompt,
154 tools: def.tools,
155 context_strategy: def.context_strategy,
156 summarize_threshold: def.summarize_threshold,
157 tool_timeout: def.tool_timeout,
158 max_tool_output_bytes: def.max_tool_output_bytes,
159 max_turns: def.max_turns,
160 max_tokens: def.max_tokens,
161 response_schema: def.response_schema,
162 run_timeout: def.run_timeout,
163 guardrails: def.guardrails,
164 provider_override: def.provider,
165 reasoning_effort: def.reasoning_effort,
166 enable_reflection: def.enable_reflection,
167 tool_output_compression_threshold: def.tool_output_compression_threshold,
168 max_tools_per_turn: def.max_tools_per_turn,
169 tool_profile: def.tool_profile,
170 max_identical_tool_calls: def.max_identical_tool_calls,
171 max_fuzzy_identical_tool_calls: def.max_fuzzy_identical_tool_calls,
172 max_tool_calls_per_turn: def.max_tool_calls_per_turn,
173 session_prune_config: def.session_prune_config,
174 enable_recursive_summarization: def.enable_recursive_summarization,
175 reflection_threshold: def.reflection_threshold,
176 consolidate_on_exit: def.consolidate_on_exit,
177 workspace: def.workspace,
178 max_total_tokens: def.max_total_tokens,
179 audit_trail: def.audit_trail,
180 audit_user_id: def.audit_user_id,
181 audit_tenant_id: def.audit_tenant_id,
182 audit_delegation_chain: def.audit_delegation_chain,
183 }
184 }
185}
186
187#[derive(Debug, Clone, Serialize, Deserialize)]
189pub(crate) struct DelegatedTask {
190 pub(crate) agent: String,
191 pub(crate) task: String,
192}
193
194#[derive(Debug, Clone)]
196pub(crate) struct SubAgentResult {
197 pub(crate) agent: String,
198 pub(crate) result: String,
199 pub(crate) tokens_used: TokenUsage,
200 pub(crate) success: bool,
201}
202
203pub struct Orchestrator<P: LlmProvider> {
213 runner: AgentRunner<P>,
214 sub_agent_tokens: Arc<Mutex<TokenUsage>>,
216}
217
218impl<P: LlmProvider + 'static> Orchestrator<P> {
219 pub fn builder(provider: Arc<P>) -> OrchestratorBuilder<P> {
221 OrchestratorBuilder {
222 provider,
223 sub_agents: vec![],
224 max_turns: 10,
225 max_tokens: 4096,
226 context_strategy: None,
227 summarize_threshold: None,
228 tool_timeout: None,
229 max_tool_output_bytes: None,
230 shared_memory: None,
231 memory_namespace_prefix: None,
232 blackboard: None,
233 knowledge_base: None,
234 on_text: None,
235 on_approval: None,
236 on_event: None,
237 guardrails: Vec::new(),
238 on_question: None,
239 run_timeout: None,
240 enable_squads: None,
241 reasoning_effort: None,
242 enable_reflection: false,
243 tool_output_compression_threshold: None,
244 max_tools_per_turn: None,
245 max_identical_tool_calls: None,
246 max_fuzzy_identical_tool_calls: None,
247 max_tool_calls_per_turn: None,
248 permission_rules: super::permission::PermissionRuleset::default(),
249 instruction_text: None,
250 learned_permissions: None,
251 lsp_manager: None,
252 observability_mode: None,
253 dispatch_mode: DispatchMode::Parallel,
254 workspace: None,
255 audit_trail: None,
256 audit_user_id: None,
257 audit_tenant_id: None,
258 audit_delegation_chain: Vec::new(),
259 allow_shared_write: true,
260 multi_agent_prompt: true,
261 spawn_config: None,
262 spawn_builtin_tools: Vec::new(),
263 tenant_tracker: None,
264 }
265 }
266
267 pub async fn run(&mut self, task: &str) -> Result<AgentOutput, Error> {
278 {
280 let mut acc = self.sub_agent_tokens.lock().expect("token lock poisoned");
281 *acc = TokenUsage::default();
282 }
283 match self.runner.execute(task).await {
284 Ok(mut output) => {
285 let sub_tokens = *self.sub_agent_tokens.lock().expect("token lock poisoned");
287 output.tokens_used += sub_tokens;
288 Ok(output)
289 }
290 Err(e) => {
291 let sub_tokens = *self.sub_agent_tokens.lock().expect("token lock poisoned");
294 let mut usage = e.partial_usage();
295 usage += sub_tokens;
296 Err(e.with_partial_usage(usage))
297 }
298 }
299 }
300}
301
302struct DelegateTaskTool {
310 shared_provider: Arc<BoxedProvider>,
311 sub_agents: Vec<SubAgentDef>,
312 max_turns: usize,
313 max_tokens: u32,
314 permission_rules: super::permission::PermissionRuleset,
316 accumulated_tokens: Arc<Mutex<TokenUsage>>,
318 shared_memory: Option<Arc<dyn Memory>>,
320 memory_namespace_prefix: Option<String>,
323 blackboard: Option<Arc<dyn Blackboard>>,
325 knowledge_base: Option<Arc<dyn KnowledgeBase>>,
327 cached_definition: ToolDefinition,
330 on_event: Option<Arc<OnEvent>>,
332 on_text: Option<Arc<crate::llm::OnText>>,
334 lsp_manager: Option<Arc<crate::lsp::LspManager>>,
336 observability_mode: super::observability::ObservabilityMode,
338 allow_shared_write: bool,
340 tenant_tracker: Option<Arc<crate::agent::tenant_tracker::TenantTokenTracker>>,
342 guardrails: Vec<Arc<dyn Guardrail>>,
347}
348
349impl DelegateTaskTool {
350 async fn delegate(&self, tasks: Vec<DelegatedTask>) -> Result<String, Error> {
351 if tasks.is_empty() {
352 return Err(Error::Agent(
353 "delegate_task requires at least one task".into(),
354 ));
355 }
356 let task_count = tasks.len();
357 let agent_names: Vec<String> = tasks.iter().map(|t| t.agent.clone()).collect();
358 let _delegate_span = info_span!(
359 "heartbit.orchestrator.delegate",
360 agent_count = task_count,
361 agents = ?agent_names,
362 );
363
364 if let Some(ref cb) = self.on_event {
365 cb(AgentEvent::SubAgentsDispatched {
366 agent: "orchestrator".into(),
367 agents: agent_names.clone(),
368 });
369 }
370
371 let mut join_set = tokio::task::JoinSet::new();
372
373 for (idx, task) in tasks.into_iter().enumerate() {
374 let agent_def = match self.sub_agents.iter().find(|a| a.name == task.agent) {
375 Some(def) => def.clone(),
376 None => {
377 let agent_name = task.agent.clone();
379 join_set.spawn(async move {
380 (
381 idx,
382 SubAgentResult {
383 agent: agent_name.clone(),
384 result: format!("Error: unknown agent '{agent_name}'"),
385 tokens_used: TokenUsage::default(),
386 success: false,
387 },
388 )
389 });
390 continue;
391 }
392 };
393
394 let provider = agent_def
395 .provider_override
396 .clone()
397 .unwrap_or_else(|| self.shared_provider.clone());
398 let max_turns = agent_def.max_turns.unwrap_or(self.max_turns);
399 let max_tokens = agent_def.max_tokens.unwrap_or(self.max_tokens);
400 let shared_memory = self.shared_memory.clone();
401 let ns_prefix = self.memory_namespace_prefix.clone();
402 let blackboard = self.blackboard.clone();
403 let knowledge_base = self.knowledge_base.clone();
404 let on_event = self.on_event.clone();
405 let on_text = self.on_text.clone();
406 let lsp_manager = self.lsp_manager.clone();
407 let permission_rules = self.permission_rules.clone();
408 let observability_mode = self.observability_mode;
409 let allow_shared_write = self.allow_shared_write;
410 let tenant_tracker = self.tenant_tracker.clone();
411 let orchestrator_guardrails = self.guardrails.clone();
417
418 info!(agent = %agent_def.name, task = %task.task, "spawning sub-agent");
419
420 join_set.spawn(async move {
421 let mut builder = AgentRunner::builder(provider)
422 .name(&agent_def.name)
423 .system_prompt(&agent_def.system_prompt)
424 .tools(agent_def.tools)
425 .max_turns(max_turns)
426 .max_tokens(max_tokens);
427
428 if let Some(strategy) = agent_def.context_strategy {
429 builder = builder.context_strategy(strategy);
430 }
431 if let Some(threshold) = agent_def.summarize_threshold {
432 builder = builder.summarize_threshold(threshold);
433 }
434 if let Some(timeout) = agent_def.tool_timeout {
435 builder = builder.tool_timeout(timeout);
436 }
437 if let Some(max) = agent_def.max_tool_output_bytes {
438 builder = builder.max_tool_output_bytes(max);
439 }
440 if let Some(schema) = agent_def.response_schema {
441 builder = builder.structured_schema(schema);
442 }
443 let mut combined_guardrails = orchestrator_guardrails;
447 combined_guardrails.extend(agent_def.guardrails);
448 if !combined_guardrails.is_empty() {
449 builder = builder.guardrails(combined_guardrails);
450 }
451 if let Some(timeout) = agent_def.run_timeout {
452 builder = builder.run_timeout(timeout);
453 }
454 if let Some(effort) = agent_def.reasoning_effort {
455 builder = builder.reasoning_effort(effort);
456 }
457 if let Some(true) = agent_def.enable_reflection {
458 builder = builder.enable_reflection(true);
459 }
460 if let Some(threshold) = agent_def.tool_output_compression_threshold {
461 builder = builder.tool_output_compression_threshold(threshold);
462 }
463 if let Some(max) = agent_def.max_tools_per_turn {
464 builder = builder.max_tools_per_turn(max);
465 }
466 if let Some(profile) = agent_def.tool_profile {
467 builder = builder.tool_profile(profile);
468 }
469 if let Some(max) = agent_def.max_identical_tool_calls {
470 builder = builder.max_identical_tool_calls(max);
471 }
472 if let Some(max) = agent_def.max_fuzzy_identical_tool_calls {
473 builder = builder.max_fuzzy_identical_tool_calls(max);
474 }
475 if let Some(cap) = agent_def.max_tool_calls_per_turn {
476 builder = builder.max_tool_calls_per_turn(cap);
477 }
478 if let Some(ref config) = agent_def.session_prune_config {
479 builder = builder.session_prune_config(config.clone());
480 }
481 if let Some(true) = agent_def.enable_recursive_summarization {
482 builder = builder.enable_recursive_summarization(true);
483 }
484 if let Some(threshold) = agent_def.reflection_threshold {
485 builder = builder.reflection_threshold(threshold);
486 }
487 if let Some(true) = agent_def.consolidate_on_exit {
488 builder = builder.consolidate_on_exit(true);
489 }
490 if let Some(ref ws) = agent_def.workspace {
491 builder = builder.workspace(ws.clone());
492 }
493 if let Some(max) = agent_def.max_total_tokens {
494 builder = builder.max_total_tokens(max);
495 }
496 if let Some(trail) = agent_def.audit_trail {
497 builder = builder.audit_trail(trail);
498 }
499 if let Some(uid) = &agent_def.audit_user_id
500 && let Some(tid) = &agent_def.audit_tenant_id
501 {
502 builder = builder.audit_user_context(uid.clone(), tid.clone());
503 }
504 if !agent_def.audit_delegation_chain.is_empty() {
505 builder =
506 builder.audit_delegation_chain(agent_def.audit_delegation_chain.clone());
507 }
508
509 if !permission_rules.is_empty() {
511 builder = builder.permission_rules(permission_rules);
512 }
513
514 builder = builder.observability_mode(observability_mode);
516
517 if let Some(ref tracker) = tenant_tracker {
520 builder = builder.tenant_tracker(tracker.clone());
521 }
522
523 if let Some(ref lsp) = lsp_manager {
525 builder = builder.lsp_manager(lsp.clone());
526 }
527
528 if let Some(ref on_event) = on_event {
530 builder = builder.on_event(on_event.clone());
531 }
532 if let Some(ref on_text) = on_text {
534 builder = builder.on_text(on_text.clone());
535 }
536
537 if let Some(ref memory) = shared_memory {
539 let agent_ns = match &ns_prefix {
540 Some(prefix) => format!("{prefix}:{}", agent_def.name),
541 None => agent_def.name.clone(),
542 };
543 let ns = Arc::new(crate::memory::namespaced::NamespacedMemory::new(
544 memory.clone(),
545 &agent_ns,
546 ));
547 builder = builder.memory(ns);
548 let mem_scope = crate::auth::TenantScope::from_audit_fields(
549 agent_def.audit_tenant_id.as_deref(),
550 agent_def.audit_user_id.as_deref(),
551 );
552 builder = builder.tools(crate::memory::shared_tools::shared_memory_tools(
553 memory.clone(),
554 &agent_ns,
555 mem_scope,
556 allow_shared_write,
557 ));
558 }
559
560 if let Some(ref bb) = blackboard {
565 builder = builder.tools(blackboard_tools(bb.clone(), &agent_def.name));
566 }
567
568 if let Some(ref kb) = knowledge_base {
570 builder = builder.knowledge(kb.clone());
571 }
572
573 let runner = match builder.build() {
574 Ok(r) => r,
575 Err(e) => {
576 return (
577 idx,
578 SubAgentResult {
579 agent: agent_def.name,
580 result: format!("Error building agent: {e}"),
581 tokens_used: TokenUsage::default(),
582 success: false,
583 },
584 );
585 }
586 };
587
588 let result = match runner.execute(&task.task).await {
589 Ok(output) => {
590 if let Some(ref bb) = blackboard {
592 let key = format!("agent:{}", agent_def.name);
593 if let Err(e) = bb
594 .write(&key, serde_json::Value::String(output.result.clone()))
595 .await
596 {
597 tracing::warn!(
598 agent = %agent_def.name,
599 error = %e,
600 "failed to write result to blackboard"
601 );
602 }
603 }
604 SubAgentResult {
605 agent: agent_def.name,
606 result: output.result,
607 tokens_used: output.tokens_used,
608 success: true,
609 }
610 }
611 Err(e) => SubAgentResult {
612 agent: agent_def.name,
613 result: format!("Error: {e}"),
614 tokens_used: e.partial_usage(),
615 success: false,
616 },
617 };
618
619 (idx, result)
620 });
621 }
622
623 let mut results: Vec<Option<(usize, SubAgentResult)>> = vec![None; task_count];
624 while let Some(result) = join_set.join_next().await {
625 match result {
626 Ok((idx, sub_result)) => {
627 results[idx] = Some((idx, sub_result));
628 }
629 Err(e) => {
630 tracing::error!(error = %e, "sub-agent task panicked");
631 }
632 }
633 }
634
635 let mut results: Vec<(usize, SubAgentResult)> = results
637 .into_iter()
638 .enumerate()
639 .map(|(idx, r)| {
640 r.unwrap_or_else(|| {
641 (
642 idx,
643 SubAgentResult {
644 agent: agent_names[idx].clone(),
645 result: "Error: sub-agent task panicked".into(),
646 tokens_used: TokenUsage::default(),
647 success: false,
648 },
649 )
650 })
651 })
652 .collect();
653 results.sort_by_key(|(idx, _)| *idx);
654
655 {
657 let mut acc = self.accumulated_tokens.lock().expect("token lock poisoned");
658 for (_, r) in &results {
659 *acc += r.tokens_used;
660 }
661 }
662
663 if let Some(ref cb) = self.on_event {
665 for (_, r) in &results {
666 cb(AgentEvent::SubAgentCompleted {
667 agent: r.agent.clone(),
668 success: r.success,
669 usage: r.tokens_used,
670 });
671 }
672 }
673
674 let formatted = results
675 .iter()
676 .map(|(_, r)| format!("=== Agent: {} ===\n{}", r.agent, r.result))
677 .collect::<Vec<_>>()
678 .join("\n\n");
679
680 Ok(formatted)
681 }
682}
683
684impl Tool for DelegateTaskTool {
685 fn definition(&self) -> ToolDefinition {
686 self.cached_definition.clone()
687 }
688
689 fn execute(
690 &self,
691 _ctx: &crate::ExecutionContext,
692 input: serde_json::Value,
693 ) -> Pin<Box<dyn Future<Output = Result<ToolOutput, Error>> + Send + '_>> {
694 Box::pin(async move {
695 let delegate_input: DelegateInput = serde_json::from_value(input)
696 .map_err(|e| Error::Agent(format!("Invalid delegate_task input: {e}")))?;
697
698 let result = self.delegate(delegate_input.tasks).await?;
699 Ok(ToolOutput::success(result))
700 })
701 }
702}
703
704#[derive(Deserialize)]
705struct DelegateInput {
706 tasks: Vec<DelegatedTask>,
707}
708
709struct FormSquadTool {
719 shared_provider: Arc<BoxedProvider>,
720 agent_pool: Vec<SubAgentDef>,
721 default_max_turns: usize,
722 default_max_tokens: u32,
723 permission_rules: super::permission::PermissionRuleset,
725 accumulated_tokens: Arc<Mutex<TokenUsage>>,
726 shared_memory: Option<Arc<dyn Memory>>,
727 memory_namespace_prefix: Option<String>,
728 blackboard: Option<Arc<dyn Blackboard>>,
730 knowledge_base: Option<Arc<dyn KnowledgeBase>>,
731 on_event: Option<Arc<OnEvent>>,
732 on_text: Option<Arc<crate::llm::OnText>>,
734 lsp_manager: Option<Arc<crate::lsp::LspManager>>,
736 cached_definition: ToolDefinition,
737 observability_mode: super::observability::ObservabilityMode,
739 allow_shared_write: bool,
741 tenant_tracker: Option<Arc<crate::agent::tenant_tracker::TenantTokenTracker>>,
743}
744
745impl Tool for FormSquadTool {
746 fn definition(&self) -> ToolDefinition {
747 self.cached_definition.clone()
748 }
749
750 fn execute(
751 &self,
752 _ctx: &crate::ExecutionContext,
753 input: serde_json::Value,
754 ) -> Pin<Box<dyn Future<Output = Result<ToolOutput, Error>> + Send + '_>> {
755 Box::pin(async move {
756 let delegate_input: DelegateInput = serde_json::from_value(input)
757 .map_err(|e| Error::Agent(format!("Invalid form_squad input: {e}")))?;
758
759 let tasks = delegate_input.tasks;
760
761 if tasks.len() < 2 {
763 return Ok(ToolOutput::error(
764 "form_squad requires at least 2 tasks. Use delegate_task for single-agent tasks."
765 .to_string(),
766 ));
767 }
768
769 {
771 let mut seen = std::collections::HashSet::new();
772 for t in &tasks {
773 if !seen.insert(&t.agent) {
774 return Ok(ToolOutput::error(format!(
775 "Duplicate agent name in squad: '{}'",
776 t.agent
777 )));
778 }
779 }
780 }
781
782 for t in &tasks {
784 if !self.agent_pool.iter().any(|a| a.name == t.agent) {
785 return Ok(ToolOutput::error(format!(
786 "Unknown agent '{}'. Available agents: {}",
787 t.agent,
788 self.agent_pool
789 .iter()
790 .map(|a| a.name.as_str())
791 .collect::<Vec<_>>()
792 .join(", ")
793 )));
794 }
795 }
796
797 let task_count = tasks.len();
798 let agent_names: Vec<String> = tasks.iter().map(|t| t.agent.clone()).collect();
799
800 let private_bb: Arc<dyn Blackboard> = Arc::new(InMemoryBlackboard::new());
802
803 let _squad_span = info_span!(
804 "heartbit.orchestrator.squad",
805 agent_count = task_count,
806 agents = ?agent_names,
807 );
808
809 if let Some(ref cb) = self.on_event {
810 cb(AgentEvent::SubAgentsDispatched {
811 agent: "squad-leader".into(),
812 agents: agent_names.clone(),
813 });
814 }
815
816 let mut join_set = tokio::task::JoinSet::new();
818
819 for (idx, task) in tasks.into_iter().enumerate() {
820 let agent_def = match self.agent_pool.iter().find(|a| a.name == task.agent) {
822 Some(def) => def.clone(),
823 None => {
824 return Ok(ToolOutput::error(format!(
825 "Internal error: agent '{}' not found after validation",
826 task.agent
827 )));
828 }
829 };
830
831 let provider = agent_def
832 .provider_override
833 .clone()
834 .unwrap_or_else(|| self.shared_provider.clone());
835 let max_turns = agent_def.max_turns.unwrap_or(self.default_max_turns);
836 let max_tokens = agent_def.max_tokens.unwrap_or(self.default_max_tokens);
837 let shared_memory = self.shared_memory.clone();
838 let ns_prefix = self.memory_namespace_prefix.clone();
839 let bb = private_bb.clone();
840 let knowledge_base = self.knowledge_base.clone();
841 let on_event = self.on_event.clone();
842 let on_text = self.on_text.clone();
843 let lsp_manager = self.lsp_manager.clone();
844 let permission_rules = self.permission_rules.clone();
845 let observability_mode = self.observability_mode;
846 let allow_shared_write = self.allow_shared_write;
847 let tenant_tracker = self.tenant_tracker.clone();
848
849 info!(agent = %agent_def.name, task = %task.task, "spawning squad member");
850
851 join_set.spawn(async move {
852 let mut builder = AgentRunner::builder(provider)
853 .name(&agent_def.name)
854 .system_prompt(&agent_def.system_prompt)
855 .tools(agent_def.tools)
856 .max_turns(max_turns)
857 .max_tokens(max_tokens);
858
859 if let Some(strategy) = agent_def.context_strategy {
860 builder = builder.context_strategy(strategy);
861 }
862 if let Some(threshold) = agent_def.summarize_threshold {
863 builder = builder.summarize_threshold(threshold);
864 }
865 if let Some(timeout) = agent_def.tool_timeout {
866 builder = builder.tool_timeout(timeout);
867 }
868 if let Some(max) = agent_def.max_tool_output_bytes {
869 builder = builder.max_tool_output_bytes(max);
870 }
871 if let Some(schema) = agent_def.response_schema {
872 builder = builder.structured_schema(schema);
873 }
874 if !agent_def.guardrails.is_empty() {
875 builder = builder.guardrails(agent_def.guardrails);
876 }
877 if let Some(timeout) = agent_def.run_timeout {
878 builder = builder.run_timeout(timeout);
879 }
880 if let Some(effort) = agent_def.reasoning_effort {
881 builder = builder.reasoning_effort(effort);
882 }
883 if let Some(true) = agent_def.enable_reflection {
884 builder = builder.enable_reflection(true);
885 }
886 if let Some(threshold) = agent_def.tool_output_compression_threshold {
887 builder = builder.tool_output_compression_threshold(threshold);
888 }
889 if let Some(max) = agent_def.max_tools_per_turn {
890 builder = builder.max_tools_per_turn(max);
891 }
892 if let Some(profile) = agent_def.tool_profile {
893 builder = builder.tool_profile(profile);
894 }
895 if let Some(max) = agent_def.max_identical_tool_calls {
896 builder = builder.max_identical_tool_calls(max);
897 }
898 if let Some(max) = agent_def.max_fuzzy_identical_tool_calls {
899 builder = builder.max_fuzzy_identical_tool_calls(max);
900 }
901 if let Some(cap) = agent_def.max_tool_calls_per_turn {
902 builder = builder.max_tool_calls_per_turn(cap);
903 }
904 if let Some(ref config) = agent_def.session_prune_config {
905 builder = builder.session_prune_config(config.clone());
906 }
907 if let Some(true) = agent_def.enable_recursive_summarization {
908 builder = builder.enable_recursive_summarization(true);
909 }
910 if let Some(threshold) = agent_def.reflection_threshold {
911 builder = builder.reflection_threshold(threshold);
912 }
913 if let Some(true) = agent_def.consolidate_on_exit {
914 builder = builder.consolidate_on_exit(true);
915 }
916 if let Some(ref ws) = agent_def.workspace {
917 builder = builder.workspace(ws.clone());
918 }
919 if let Some(max) = agent_def.max_total_tokens {
920 builder = builder.max_total_tokens(max);
921 }
922 if let Some(trail) = agent_def.audit_trail {
923 builder = builder.audit_trail(trail);
924 }
925 if let Some(uid) = &agent_def.audit_user_id
926 && let Some(tid) = &agent_def.audit_tenant_id
927 {
928 builder = builder.audit_user_context(uid.clone(), tid.clone());
929 }
930 if !agent_def.audit_delegation_chain.is_empty() {
931 builder = builder
932 .audit_delegation_chain(agent_def.audit_delegation_chain.clone());
933 }
934
935 if !permission_rules.is_empty() {
937 builder = builder.permission_rules(permission_rules);
938 }
939
940 builder = builder.observability_mode(observability_mode);
942
943 if let Some(ref tracker) = tenant_tracker {
946 builder = builder.tenant_tracker(tracker.clone());
947 }
948
949 if let Some(ref lsp) = lsp_manager {
951 builder = builder.lsp_manager(lsp.clone());
952 }
953
954 if let Some(ref on_event) = on_event {
956 builder = builder.on_event(on_event.clone());
957 }
958 if let Some(ref on_text) = on_text {
960 builder = builder.on_text(on_text.clone());
961 }
962
963 if let Some(ref memory) = shared_memory {
965 let agent_ns = match &ns_prefix {
966 Some(prefix) => format!("{prefix}:{}", agent_def.name),
967 None => agent_def.name.clone(),
968 };
969 let ns = Arc::new(crate::memory::namespaced::NamespacedMemory::new(
970 memory.clone(),
971 &agent_ns,
972 ));
973 builder = builder.memory(ns);
974 let mem_scope = crate::auth::TenantScope::from_audit_fields(
975 agent_def.audit_tenant_id.as_deref(),
976 agent_def.audit_user_id.as_deref(),
977 );
978 builder = builder.tools(crate::memory::shared_tools::shared_memory_tools(
979 memory.clone(),
980 &agent_ns,
981 mem_scope,
982 allow_shared_write,
983 ));
984 }
985
986 builder = builder.tools(blackboard_tools(bb.clone(), &agent_def.name));
990
991 if let Some(ref kb) = knowledge_base {
993 builder = builder.knowledge(kb.clone());
994 }
995
996 let runner = match builder.build() {
997 Ok(r) => r,
998 Err(e) => {
999 return (
1000 idx,
1001 SubAgentResult {
1002 agent: agent_def.name,
1003 result: format!("Error building agent: {e}"),
1004 tokens_used: TokenUsage::default(),
1005 success: false,
1006 },
1007 );
1008 }
1009 };
1010
1011 let result = match runner.execute(&task.task).await {
1012 Ok(output) => {
1013 let key = format!("agent:{}", agent_def.name);
1015 if let Err(e) = bb
1016 .write(&key, serde_json::Value::String(output.result.clone()))
1017 .await
1018 {
1019 tracing::warn!(
1020 agent = %agent_def.name,
1021 error = %e,
1022 "failed to write result to private blackboard"
1023 );
1024 }
1025 SubAgentResult {
1026 agent: agent_def.name,
1027 result: output.result,
1028 tokens_used: output.tokens_used,
1029 success: true,
1030 }
1031 }
1032 Err(e) => SubAgentResult {
1033 agent: agent_def.name,
1034 result: format!("Error: {e}"),
1035 tokens_used: e.partial_usage(),
1036 success: false,
1037 },
1038 };
1039
1040 (idx, result)
1041 });
1042 }
1043
1044 let mut results: Vec<Option<(usize, SubAgentResult)>> = vec![None; task_count];
1045 while let Some(result) = join_set.join_next().await {
1046 match result {
1047 Ok((idx, sub_result)) => {
1048 results[idx] = Some((idx, sub_result));
1049 }
1050 Err(e) => {
1051 tracing::error!(error = %e, "squad member task panicked");
1052 }
1053 }
1054 }
1055
1056 let mut results: Vec<(usize, SubAgentResult)> = results
1058 .into_iter()
1059 .enumerate()
1060 .map(|(idx, r)| {
1061 r.unwrap_or_else(|| {
1062 (
1063 idx,
1064 SubAgentResult {
1065 agent: agent_names[idx].clone(),
1066 result: "Error: squad member task panicked".into(),
1067 tokens_used: TokenUsage::default(),
1068 success: false,
1069 },
1070 )
1071 })
1072 })
1073 .collect();
1074 results.sort_by_key(|(idx, _)| *idx);
1075
1076 let squad_label = format!("squad[{}]", agent_names.join(","));
1077 let bb_key = format!("squad:{}", agent_names.join("+"));
1078
1079 let mut total_tokens = TokenUsage::default();
1081 {
1082 let mut acc = self.accumulated_tokens.lock().expect("token lock poisoned");
1083 for (_, r) in &results {
1084 *acc += r.tokens_used;
1085 total_tokens += r.tokens_used;
1086 }
1087 }
1088
1089 if let Some(ref cb) = self.on_event {
1091 for (_, r) in &results {
1092 cb(AgentEvent::SubAgentCompleted {
1093 agent: r.agent.clone(),
1094 success: r.success,
1095 usage: r.tokens_used,
1096 });
1097 }
1098 }
1099
1100 let all_success = results.iter().all(|(_, r)| r.success);
1101
1102 let formatted = results
1103 .iter()
1104 .map(|(_, r)| format!("=== Agent: {} ===\n{}", r.agent, r.result))
1105 .collect::<Vec<_>>()
1106 .join("\n\n");
1107
1108 if let Some(ref cb) = self.on_event {
1110 cb(AgentEvent::SubAgentCompleted {
1111 agent: squad_label,
1112 success: all_success,
1113 usage: total_tokens,
1114 });
1115 }
1116
1117 if let Some(ref bb) = self.blackboard
1119 && let Err(e) = bb
1120 .write(&bb_key, serde_json::Value::String(formatted.clone()))
1121 .await
1122 {
1123 tracing::warn!(
1124 key = %bb_key,
1125 error = %e,
1126 "failed to write squad result to outer blackboard"
1127 );
1128 }
1129
1130 Ok(ToolOutput::success(formatted))
1131 })
1132 }
1133}
1134
1135struct SpawnAgentTool {
1141 shared_provider: Arc<BoxedProvider>,
1142 spawn_config: crate::types::SpawnConfig,
1143 tool_pool: std::collections::HashMap<String, Arc<dyn Tool>>,
1145 spawn_count: Arc<std::sync::atomic::AtomicU32>,
1147 spawned_names: Arc<Mutex<std::collections::HashSet<String>>>,
1149 accumulated_tokens: Arc<Mutex<TokenUsage>>,
1151 permission_rules: super::permission::PermissionRuleset,
1152 shared_memory: Option<Arc<dyn Memory>>,
1153 memory_namespace_prefix: Option<String>,
1154 on_event: Option<Arc<OnEvent>>,
1155 on_text: Option<Arc<crate::llm::OnText>>,
1156 lsp_manager: Option<Arc<crate::lsp::LspManager>>,
1157 observability_mode: super::observability::ObservabilityMode,
1158 workspace: Option<std::path::PathBuf>,
1159 guardrails: Vec<Arc<dyn Guardrail>>,
1160 audit_trail: Option<Arc<dyn super::audit::AuditTrail>>,
1161 audit_user_id: Option<String>,
1162 audit_tenant_id: Option<String>,
1163 audit_delegation_chain: Vec<String>,
1164 cached_definition: ToolDefinition,
1165 tenant_tracker: Option<Arc<crate::agent::tenant_tracker::TenantTokenTracker>>,
1167}
1168
1169#[derive(Deserialize)]
1170struct SpawnAgentInput {
1171 name: String,
1172 system_prompt: String,
1173 #[serde(default)]
1174 tools: Vec<String>,
1175 task: String,
1176}
1177
1178const SPAWN_MAX_PROMPT_BYTES: usize = 32 * 1024;
1180
1181impl SpawnAgentTool {
1182 fn build_definition(config: &crate::types::SpawnConfig) -> ToolDefinition {
1183 let allowlist = if config.tool_allowlist.is_empty() {
1184 "(none — reasoning-only agents)".to_string()
1185 } else {
1186 config.tool_allowlist.join(", ")
1187 };
1188 ToolDefinition {
1189 name: "spawn_agent".into(),
1190 description: format!(
1191 "Create a new specialist agent at runtime when no pre-configured agent fits the task. \
1192 The spawned agent runs with the given system prompt and tool subset, then returns its result.\n\n\
1193 Available tools for spawned agents: [{allowlist}]. Budget: {} agents max per run.",
1194 config.max_spawned_agents
1195 ),
1196 input_schema: json!({
1197 "type": "object",
1198 "required": ["name", "system_prompt", "task"],
1199 "properties": {
1200 "name": {
1201 "type": "string",
1202 "description": "Lowercase identifier for the agent (a-z, 0-9, underscores). Must start with a letter. E.g. 'tax_specialist', 'csv_analyzer'."
1203 },
1204 "system_prompt": {
1205 "type": "string",
1206 "description": "The agent's role and behavior instructions. Be specific about expertise and constraints."
1207 },
1208 "tools": {
1209 "type": "array",
1210 "items": { "type": "string" },
1211 "description": format!("Subset of available tools: [{allowlist}]. Empty array creates a reasoning-only agent.")
1212 },
1213 "task": {
1214 "type": "string",
1215 "description": "The specific task for this agent to accomplish."
1216 }
1217 },
1218 "additionalProperties": false
1219 }),
1220 }
1221 }
1222
1223 async fn spawn(&self, input: SpawnAgentInput) -> Result<ToolOutput, Error> {
1224 let current = self.spawn_count.load(std::sync::atomic::Ordering::Relaxed);
1226 if current >= self.spawn_config.max_spawned_agents {
1227 return Ok(ToolOutput::error(format!(
1228 "Spawn limit reached: {current}/{} agents already spawned this run.",
1229 self.spawn_config.max_spawned_agents
1230 )));
1231 }
1232
1233 let name_re =
1235 regex::Regex::new(r"^[a-z][a-z0-9_]{0,63}$").expect("spawn agent name regex is valid");
1236 if !name_re.is_match(&input.name) {
1237 return Ok(ToolOutput::error(format!(
1238 "Invalid agent name '{}'. Must match ^[a-z][a-z0-9_]{{0,63}}$ \
1239 (lowercase, starts with letter, alphanumeric + underscores, max 64 chars).",
1240 input.name
1241 )));
1242 }
1243
1244 {
1246 let mut names = self.spawned_names.lock().expect("spawned names lock");
1247 if !names.insert(input.name.clone()) {
1248 return Ok(ToolOutput::error(format!(
1249 "Agent name '{}' already used in this run. Choose a different name.",
1250 input.name
1251 )));
1252 }
1253 }
1254
1255 for tool_name in &input.tools {
1257 if !self.tool_pool.contains_key(tool_name) {
1258 let available: Vec<&str> = self.tool_pool.keys().map(|k| k.as_str()).collect();
1259 return Ok(ToolOutput::error(format!(
1260 "Tool '{}' not in allowlist. Available: [{}]",
1261 tool_name,
1262 available.join(", ")
1263 )));
1264 }
1265 }
1266
1267 {
1269 let acc = self.accumulated_tokens.lock().expect("token lock");
1270 let used = acc.total();
1271 if used >= self.spawn_config.max_total_tokens {
1272 return Ok(ToolOutput::error(format!(
1273 "Spawn token budget exhausted: {used}/{} tokens used across spawned agents.",
1274 self.spawn_config.max_total_tokens
1275 )));
1276 }
1277 }
1278
1279 if input.system_prompt.len() > SPAWN_MAX_PROMPT_BYTES {
1281 return Ok(ToolOutput::error(format!(
1282 "System prompt too long: {} bytes (max {SPAWN_MAX_PROMPT_BYTES}).",
1283 input.system_prompt.len()
1284 )));
1285 }
1286
1287 let spawned_name = format!("spawn:{}", input.name);
1288
1289 if let Some(ref cb) = self.on_event {
1291 cb(AgentEvent::AgentSpawned {
1292 agent: "orchestrator".into(),
1293 spawned_name: spawned_name.clone(),
1294 tools: input.tools.clone(),
1295 task: input.task.clone(),
1296 });
1297 }
1298
1299 let selected_tools: Vec<Arc<dyn Tool>> = input
1301 .tools
1302 .iter()
1303 .filter_map(|name| self.tool_pool.get(name).cloned())
1304 .collect();
1305
1306 let mut builder = AgentRunner::builder(self.shared_provider.clone())
1308 .name(&spawned_name)
1309 .system_prompt(&input.system_prompt)
1310 .tools(selected_tools)
1311 .max_turns(self.spawn_config.max_turns)
1312 .max_tokens(self.spawn_config.max_tokens)
1313 .observability_mode(self.observability_mode);
1314
1315 if !self.permission_rules.is_empty() {
1317 builder = builder.permission_rules(self.permission_rules.clone());
1318 }
1319 if !self.guardrails.is_empty() {
1320 builder = builder.guardrails(self.guardrails.clone());
1321 }
1322 if let Some(ref ws) = self.workspace {
1323 builder = builder.workspace(ws.clone());
1324 }
1325 if let Some(ref lsp) = self.lsp_manager {
1326 builder = builder.lsp_manager(lsp.clone());
1327 }
1328 if let Some(ref cb) = self.on_event {
1329 builder = builder.on_event(cb.clone());
1330 }
1331 if let Some(ref cb) = self.on_text {
1332 builder = builder.on_text(cb.clone());
1333 }
1334 if let Some(ref trail) = self.audit_trail {
1335 builder = builder.audit_trail(trail.clone());
1336 }
1337 if let (Some(uid), Some(tid)) = (&self.audit_user_id, &self.audit_tenant_id) {
1338 builder = builder.audit_user_context(uid.clone(), tid.clone());
1339 }
1340 if !self.audit_delegation_chain.is_empty() {
1341 let mut chain = self.audit_delegation_chain.clone();
1342 chain.push(spawned_name.clone());
1343 builder = builder.audit_delegation_chain(chain);
1344 }
1345
1346 if let Some(ref tracker) = self.tenant_tracker {
1349 builder = builder.tenant_tracker(tracker.clone());
1350 }
1351
1352 if let Some(ref memory) = self.shared_memory {
1354 let agent_ns = match &self.memory_namespace_prefix {
1355 Some(prefix) => format!("{prefix}:{spawned_name}"),
1356 None => spawned_name.clone(),
1357 };
1358 let ns = Arc::new(crate::memory::namespaced::NamespacedMemory::new(
1359 memory.clone(),
1360 &agent_ns,
1361 ));
1362 builder = builder.memory(ns);
1363 let mem_scope = crate::auth::TenantScope::from_audit_fields(
1364 self.audit_tenant_id.as_deref(),
1365 self.audit_user_id.as_deref(),
1366 );
1367 builder = builder.tools(crate::memory::shared_tools::shared_memory_tools(
1368 memory.clone(),
1369 &agent_ns,
1370 mem_scope,
1371 false, ));
1373 }
1374
1375 let runner = builder.build()?;
1376
1377 info!(
1378 agent = %spawned_name,
1379 tools = ?input.tools,
1380 "spawning dynamic agent"
1381 );
1382
1383 match runner.execute(&input.task).await {
1385 Ok(output) => {
1386 self.spawn_count
1388 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1389 {
1390 let mut acc = self.accumulated_tokens.lock().expect("token lock");
1391 *acc += output.tokens_used;
1392 }
1393 if let Some(ref cb) = self.on_event {
1394 cb(AgentEvent::SubAgentCompleted {
1395 agent: spawned_name.clone(),
1396 success: true,
1397 usage: output.tokens_used,
1398 });
1399 }
1400 Ok(ToolOutput::success(format!(
1401 "=== Spawned Agent: {} ===\n{}",
1402 spawned_name, output.result
1403 )))
1404 }
1405 Err(e) => {
1406 self.spawn_count
1407 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1408 let partial = e.partial_usage();
1409 {
1410 let mut acc = self.accumulated_tokens.lock().expect("token lock");
1411 *acc += partial;
1412 }
1413 if let Some(ref cb) = self.on_event {
1414 cb(AgentEvent::SubAgentCompleted {
1415 agent: spawned_name.clone(),
1416 success: false,
1417 usage: partial,
1418 });
1419 }
1420 Ok(ToolOutput::error(format!(
1421 "Spawned agent '{spawned_name}' failed: {e}"
1422 )))
1423 }
1424 }
1425 }
1426}
1427
1428impl Tool for SpawnAgentTool {
1429 fn definition(&self) -> ToolDefinition {
1430 self.cached_definition.clone()
1431 }
1432
1433 fn execute(
1434 &self,
1435 _ctx: &crate::ExecutionContext,
1436 input: serde_json::Value,
1437 ) -> Pin<Box<dyn Future<Output = Result<ToolOutput, Error>> + Send + '_>> {
1438 Box::pin(async move {
1439 let spawn_input: SpawnAgentInput = serde_json::from_value(input)
1440 .map_err(|e| Error::Agent(format!("Invalid spawn_agent input: {e}")))?;
1441 self.spawn(spawn_input).await
1442 })
1443 }
1444}
1445
1446pub fn build_system_prompt(
1457 agents: &[(&str, &str, &[String])],
1458 squads_enabled: bool,
1459 dispatch_mode: DispatchMode,
1460) -> String {
1461 let agent_list: String = agents
1462 .iter()
1463 .map(|(name, desc, tools)| {
1464 if tools.is_empty() {
1465 format!("- **{name}**: {desc}\n Tools: (none)")
1466 } else {
1467 format!("- **{name}**: {desc}\n Tools: {}", tools.join(", "))
1468 }
1469 })
1470 .collect::<Vec<_>>()
1471 .join("\n");
1472
1473 let delegation_instructions = match (squads_enabled, dispatch_mode) {
1474 (_, DispatchMode::Sequential) => {
1475 "## Delegation Tool\n\
1476 Delegate to ONE agent at a time using **delegate_task**. Wait for the result \
1477 before deciding the next agent. Do NOT batch multiple agents in a single call."
1478 }
1479 (true, DispatchMode::Parallel) => {
1480 "## Delegation Tools\n\
1481 You have two delegation tools:\n\n\
1482 1. **delegate_task** — Run independent subtasks in parallel. Each agent works in \
1483 isolation and cannot see other agents' output. Use when subtasks are independent.\n\n\
1484 2. **form_squad** — Run subtasks in parallel with a shared blackboard. \
1485 Unlike delegate_task, agents can read each other's results via the blackboard. \
1486 Agents run concurrently — use when they benefit from shared state, not when \
1487 strict ordering is needed.\n\n\
1488 After receiving results, synthesize them into a coherent response."
1489 }
1490 (false, DispatchMode::Parallel) => {
1491 "## Delegation Tool\n\
1492 Use the **delegate_task** tool to assign work to sub-agents. You can assign \
1493 multiple tasks at once for parallel execution. Each agent works in isolation. \
1494 After receiving results, synthesize them into a coherent response."
1495 }
1496 };
1497
1498 let choose_tool_step = match (squads_enabled, dispatch_mode) {
1499 (_, DispatchMode::Sequential) => {
1500 "3. DELEGATE: Use delegate_task with ONE agent at a time. Wait for results before \
1501 delegating to the next agent."
1502 }
1503 (true, DispatchMode::Parallel) => {
1504 "3. CHOOSE TOOL: Select delegate_task for independent parallel work, or form_squad \
1505 when agents benefit from shared state via a blackboard."
1506 }
1507 (false, DispatchMode::Parallel) => {
1508 "3. DELEGATE: Use delegate_task to assign subtasks to the best-fit agents."
1509 }
1510 };
1511
1512 format!(
1513 "You are an orchestrator agent. Analyze incoming tasks and delegate work to \
1514 specialized sub-agents.\n\n\
1515 ## Decision Process\n\
1516 1. DECOMPOSE: Break the task into distinct subtasks. Identify which require different expertise.\n\
1517 2. MATCH: For each subtask, pick the best-fit agent based on their description and tools.\n\
1518 {choose_tool_step}\n\n\
1519 ## Important\n\
1520 - ALWAYS delegate to a sub-agent using your delegation tools. You do NOT have any \
1521 direct capabilities — sub-agents have the tools. Never respond to the user directly \
1522 without delegating first.\n\n\
1523 ## Effort Scaling\n\
1524 - If only ONE agent is relevant, delegate a single task. Do NOT force-split across agents.\n\
1525 - If the task is simple enough for one agent, use one agent.\n\
1526 - Only use multiple agents when the task genuinely has multiple distinct parts \
1527 needing different expertise.\n\n\
1528 ## Task Quality\n\
1529 - Each delegated task must be self-contained: include all context the agent needs.\n\
1530 - Be specific: \"Read /path/to/file and extract X\" not \"look at the project\".\n\
1531 - Avoid overlapping tasks — no two agents should do the same work.\n\n\
1532 ## Available Sub-Agents\n\
1533 Choose agents based on their description and available tools:\n\
1534 {agent_list}\n\n\
1535 {delegation_instructions}"
1536 )
1537}
1538
1539pub fn build_delegate_tool_schema(
1547 agents: &[(&str, &str, &[String])],
1548 dispatch_mode: DispatchMode,
1549) -> ToolDefinition {
1550 let agent_descriptions: Vec<serde_json::Value> = agents
1551 .iter()
1552 .map(|(name, desc, tools)| json!({"name": name, "description": desc, "tools": tools}))
1553 .collect();
1554
1555 let (description, tasks_schema) = match dispatch_mode {
1556 DispatchMode::Sequential => (
1557 format!(
1558 "Delegate a task to ONE sub-agent at a time. Wait for the result before \
1559 delegating to the next agent. Each task runs in isolation. \
1560 Write clear, self-contained task descriptions with all necessary context. \
1561 Available agents: {}",
1562 serde_json::to_string(&agent_descriptions)
1563 .expect("agent list serialization is infallible")
1564 ),
1565 json!({
1566 "type": "array",
1567 "items": {
1568 "type": "object",
1569 "properties": {
1570 "agent": {
1571 "type": "string",
1572 "description": "Name of the sub-agent"
1573 },
1574 "task": {
1575 "type": "string",
1576 "description": "Task instruction for the sub-agent"
1577 }
1578 },
1579 "required": ["agent", "task"]
1580 },
1581 "minItems": 1,
1582 "maxItems": 1
1583 }),
1584 ),
1585 DispatchMode::Parallel => (
1586 format!(
1587 "Delegate independent tasks to sub-agents for parallel execution. \
1588 Each task runs in isolation — agents cannot see each other's work. \
1589 Write clear, self-contained task descriptions with all necessary context. \
1590 Available agents: {}",
1591 serde_json::to_string(&agent_descriptions)
1592 .expect("agent list serialization is infallible")
1593 ),
1594 json!({
1595 "type": "array",
1596 "items": {
1597 "type": "object",
1598 "properties": {
1599 "agent": {
1600 "type": "string",
1601 "description": "Name of the sub-agent"
1602 },
1603 "task": {
1604 "type": "string",
1605 "description": "Task instruction for the sub-agent"
1606 }
1607 },
1608 "required": ["agent", "task"]
1609 },
1610 "minItems": 1
1611 }),
1612 ),
1613 };
1614
1615 ToolDefinition {
1616 name: "delegate_task".into(),
1617 description,
1618 input_schema: json!({
1619 "type": "object",
1620 "properties": {
1621 "tasks": tasks_schema
1622 },
1623 "required": ["tasks"]
1624 }),
1625 }
1626}
1627
1628pub(crate) fn build_form_squad_tool_schema(agents: &[(&str, &str, &[String])]) -> ToolDefinition {
1635 let agent_descriptions: Vec<serde_json::Value> = agents
1636 .iter()
1637 .map(|(name, desc, tools)| json!({"name": name, "description": desc, "tools": tools}))
1638 .collect();
1639
1640 ToolDefinition {
1641 name: "form_squad".into(),
1642 description: format!(
1643 "Dispatch per-agent tasks in parallel with a shared blackboard for intra-squad coordination. \
1644 Unlike delegate_task, squad agents can read each other's results via the blackboard. \
1645 Use this when agents benefit from shared state (e.g., building on each other's work, \
1646 coordinating on a shared artifact). Agents run concurrently. \
1647 Requires at least 2 tasks (one per agent). \
1648 Available agents: {}",
1649 serde_json::to_string(&agent_descriptions)
1650 .expect("agent list serialization is infallible")
1651 ),
1652 input_schema: json!({
1653 "type": "object",
1654 "properties": {
1655 "tasks": {
1656 "type": "array",
1657 "items": {
1658 "type": "object",
1659 "properties": {
1660 "agent": {
1661 "type": "string",
1662 "description": "Name of the sub-agent"
1663 },
1664 "task": {
1665 "type": "string",
1666 "description": "Task instruction for the sub-agent"
1667 }
1668 },
1669 "required": ["agent", "task"]
1670 },
1671 "minItems": 2,
1672 "description": "Per-agent tasks for the squad (minimum 2)"
1673 }
1674 },
1675 "required": ["tasks"]
1676 }),
1677 }
1678}
1679
1680#[derive(Default)]
1684pub struct SubAgentConfig {
1685 pub name: String,
1687 pub description: String,
1689 pub system_prompt: String,
1691 pub tools: Vec<Arc<dyn Tool>>,
1693 pub context_strategy: Option<ContextStrategy>,
1695 pub summarize_threshold: Option<u32>,
1697 pub tool_timeout: Option<Duration>,
1699 pub max_tool_output_bytes: Option<usize>,
1701 pub max_turns: Option<usize>,
1703 pub max_tokens: Option<u32>,
1705 pub response_schema: Option<serde_json::Value>,
1708 pub run_timeout: Option<Duration>,
1711 pub guardrails: Vec<Arc<dyn Guardrail>>,
1713 pub provider: Option<Arc<BoxedProvider>>,
1717 pub reasoning_effort: Option<crate::llm::types::ReasoningEffort>,
1719 pub enable_reflection: Option<bool>,
1721 pub tool_output_compression_threshold: Option<usize>,
1723 pub max_tools_per_turn: Option<usize>,
1725 pub tool_profile: Option<super::tool_filter::ToolProfile>,
1727 pub max_identical_tool_calls: Option<u32>,
1729 pub max_fuzzy_identical_tool_calls: Option<u32>,
1731 pub max_tool_calls_per_turn: Option<u32>,
1733 pub session_prune_config: Option<crate::agent::pruner::SessionPruneConfig>,
1735 pub enable_recursive_summarization: Option<bool>,
1737 pub reflection_threshold: Option<u32>,
1739 pub consolidate_on_exit: Option<bool>,
1741 pub workspace: Option<std::path::PathBuf>,
1743 pub max_total_tokens: Option<u64>,
1745 pub audit_trail: Option<Arc<dyn super::audit::AuditTrail>>,
1747 pub audit_user_id: Option<String>,
1749 pub audit_tenant_id: Option<String>,
1751 #[allow(dead_code)]
1753 pub audit_delegation_chain: Vec<String>,
1754}
1755
1756pub struct OrchestratorBuilder<P: LlmProvider> {
1768 provider: Arc<P>,
1769 sub_agents: Vec<SubAgentDef>,
1770 max_turns: usize,
1771 max_tokens: u32,
1772 context_strategy: Option<ContextStrategy>,
1773 summarize_threshold: Option<u32>,
1774 tool_timeout: Option<Duration>,
1775 max_tool_output_bytes: Option<usize>,
1776 shared_memory: Option<Arc<dyn Memory>>,
1777 memory_namespace_prefix: Option<String>,
1778 blackboard: Option<Arc<dyn Blackboard>>,
1779 knowledge_base: Option<Arc<dyn KnowledgeBase>>,
1780 on_text: Option<Arc<crate::llm::OnText>>,
1781 on_approval: Option<Arc<crate::llm::OnApproval>>,
1782 on_event: Option<Arc<OnEvent>>,
1783 guardrails: Vec<Arc<dyn Guardrail>>,
1784 on_question: Option<Arc<OnQuestion>>,
1785 run_timeout: Option<Duration>,
1786 enable_squads: Option<bool>,
1787 reasoning_effort: Option<crate::llm::types::ReasoningEffort>,
1788 enable_reflection: bool,
1789 tool_output_compression_threshold: Option<usize>,
1790 max_tools_per_turn: Option<usize>,
1791 max_identical_tool_calls: Option<u32>,
1792 max_fuzzy_identical_tool_calls: Option<u32>,
1793 max_tool_calls_per_turn: Option<u32>,
1794 permission_rules: super::permission::PermissionRuleset,
1795 instruction_text: Option<String>,
1796 learned_permissions: Option<Arc<std::sync::Mutex<super::permission::LearnedPermissions>>>,
1797 lsp_manager: Option<Arc<crate::lsp::LspManager>>,
1798 observability_mode: Option<super::observability::ObservabilityMode>,
1799 dispatch_mode: DispatchMode,
1800 workspace: Option<std::path::PathBuf>,
1802 audit_trail: Option<Arc<dyn super::audit::AuditTrail>>,
1804 audit_user_id: Option<String>,
1806 audit_tenant_id: Option<String>,
1808 audit_delegation_chain: Vec<String>,
1810 allow_shared_write: bool,
1814 multi_agent_prompt: bool,
1817 spawn_config: Option<crate::types::SpawnConfig>,
1820 spawn_builtin_tools: Vec<Arc<dyn Tool>>,
1822 tenant_tracker: Option<Arc<crate::agent::tenant_tracker::TenantTokenTracker>>,
1827}
1828
1829impl<P: LlmProvider + 'static> OrchestratorBuilder<P> {
1830 pub fn sub_agent(
1832 mut self,
1833 name: impl Into<String>,
1834 description: impl Into<String>,
1835 system_prompt: impl Into<String>,
1836 ) -> Self {
1837 let mut def = SubAgentDef::new(name, description, system_prompt);
1838 def.workspace = self.workspace.clone();
1839 def.audit_trail = self.audit_trail.clone();
1840 def.audit_user_id = self.audit_user_id.clone();
1841 def.audit_tenant_id = self.audit_tenant_id.clone();
1842 def.audit_delegation_chain = self.audit_delegation_chain.clone();
1843 self.sub_agents.push(def);
1844 self
1845 }
1846
1847 pub fn sub_agent_with_tools(
1849 mut self,
1850 name: impl Into<String>,
1851 description: impl Into<String>,
1852 system_prompt: impl Into<String>,
1853 tools: Vec<Arc<dyn Tool>>,
1854 ) -> Self {
1855 let mut def = SubAgentDef::new(name, description, system_prompt);
1856 def.tools = tools;
1857 def.workspace = self.workspace.clone();
1858 def.audit_trail = self.audit_trail.clone();
1859 def.audit_user_id = self.audit_user_id.clone();
1860 def.audit_tenant_id = self.audit_tenant_id.clone();
1861 def.audit_delegation_chain = self.audit_delegation_chain.clone();
1862 self.sub_agents.push(def);
1863 self
1864 }
1865
1866 pub fn sub_agent_full(mut self, config: SubAgentConfig) -> Self {
1868 let mut def = SubAgentDef::from(config);
1869 if def.workspace.is_none() {
1870 def.workspace = self.workspace.clone();
1871 }
1872 if def.audit_trail.is_none() {
1873 def.audit_trail = self.audit_trail.clone();
1874 }
1875 if def.audit_user_id.is_none() {
1876 def.audit_user_id = self.audit_user_id.clone();
1877 }
1878 if def.audit_tenant_id.is_none() {
1879 def.audit_tenant_id = self.audit_tenant_id.clone();
1880 }
1881 if def.audit_delegation_chain.is_empty() {
1882 def.audit_delegation_chain = self.audit_delegation_chain.clone();
1883 }
1884 self.sub_agents.push(def);
1885 self
1886 }
1887
1888 pub fn max_turns(mut self, max_turns: usize) -> Self {
1890 self.max_turns = max_turns;
1891 self
1892 }
1893
1894 pub fn max_tokens(mut self, max_tokens: u32) -> Self {
1896 self.max_tokens = max_tokens;
1897 self
1898 }
1899
1900 pub fn context_strategy(mut self, strategy: ContextStrategy) -> Self {
1902 self.context_strategy = Some(strategy);
1903 self
1904 }
1905
1906 pub fn summarize_threshold(mut self, threshold: u32) -> Self {
1908 self.summarize_threshold = Some(threshold);
1909 self
1910 }
1911
1912 pub fn tool_timeout(mut self, timeout: Duration) -> Self {
1914 self.tool_timeout = Some(timeout);
1915 self
1916 }
1917
1918 pub fn max_tool_output_bytes(mut self, max: usize) -> Self {
1920 self.max_tool_output_bytes = Some(max);
1921 self
1922 }
1923
1924 pub fn shared_memory(mut self, memory: Arc<dyn Memory>) -> Self {
1928 self.shared_memory = Some(memory);
1929 self
1930 }
1931
1932 pub fn memory_namespace_prefix(mut self, prefix: impl Into<String>) -> Self {
1936 self.memory_namespace_prefix = Some(prefix.into());
1937 self
1938 }
1939
1940 pub fn allow_shared_write(mut self, allow: bool) -> Self {
1946 self.allow_shared_write = allow;
1947 self
1948 }
1949
1950 pub fn multi_agent_prompt(mut self, enabled: bool) -> Self {
1953 self.multi_agent_prompt = enabled;
1954 self
1955 }
1956
1957 pub fn spawn_config(
1963 mut self,
1964 config: crate::types::SpawnConfig,
1965 builtin_tools: Vec<Arc<dyn Tool>>,
1966 ) -> Self {
1967 self.spawn_config = Some(config);
1968 self.spawn_builtin_tools = builtin_tools;
1969 self
1970 }
1971
1972 pub fn tenant_tracker(
1981 mut self,
1982 tracker: Arc<crate::agent::tenant_tracker::TenantTokenTracker>,
1983 ) -> Self {
1984 self.tenant_tracker = Some(tracker);
1985 self
1986 }
1987
1988 pub fn blackboard(mut self, blackboard: Arc<dyn Blackboard>) -> Self {
1994 self.blackboard = Some(blackboard);
1995 self
1996 }
1997
1998 pub fn knowledge(mut self, kb: Arc<dyn KnowledgeBase>) -> Self {
2003 self.knowledge_base = Some(kb);
2004 self
2005 }
2006
2007 pub fn on_text(mut self, callback: Arc<crate::llm::OnText>) -> Self {
2011 self.on_text = Some(callback);
2012 self
2013 }
2014
2015 pub fn on_approval(mut self, callback: Arc<crate::llm::OnApproval>) -> Self {
2019 self.on_approval = Some(callback);
2020 self
2021 }
2022
2023 pub fn learned_permissions(
2025 mut self,
2026 learned: Arc<std::sync::Mutex<super::permission::LearnedPermissions>>,
2027 ) -> Self {
2028 self.learned_permissions = Some(learned);
2029 self
2030 }
2031
2032 pub fn lsp_manager(mut self, manager: Arc<crate::lsp::LspManager>) -> Self {
2034 self.lsp_manager = Some(manager);
2035 self
2036 }
2037
2038 pub fn on_event(mut self, callback: Arc<OnEvent>) -> Self {
2044 self.on_event = Some(callback);
2045 self
2046 }
2047
2048 pub fn guardrail(mut self, guardrail: Arc<dyn Guardrail>) -> Self {
2050 self.guardrails.push(guardrail);
2051 self
2052 }
2053
2054 pub fn guardrails(mut self, guardrails: Vec<Arc<dyn Guardrail>>) -> Self {
2056 self.guardrails.extend(guardrails);
2057 self
2058 }
2059
2060 pub fn on_question(mut self, callback: Arc<OnQuestion>) -> Self {
2062 self.on_question = Some(callback);
2063 self
2064 }
2065
2066 pub fn run_timeout(mut self, timeout: Duration) -> Self {
2069 self.run_timeout = Some(timeout);
2070 self
2071 }
2072
2073 pub fn enable_squads(mut self, enable: bool) -> Self {
2081 self.enable_squads = Some(enable);
2082 self
2083 }
2084
2085 pub fn reasoning_effort(mut self, effort: crate::llm::types::ReasoningEffort) -> Self {
2087 self.reasoning_effort = Some(effort);
2088 self
2089 }
2090
2091 pub fn enable_reflection(mut self, enabled: bool) -> Self {
2093 self.enable_reflection = enabled;
2094 self
2095 }
2096
2097 pub fn tool_output_compression_threshold(mut self, threshold: usize) -> Self {
2099 self.tool_output_compression_threshold = Some(threshold);
2100 self
2101 }
2102
2103 pub fn max_tools_per_turn(mut self, max: usize) -> Self {
2105 self.max_tools_per_turn = Some(max);
2106 self
2107 }
2108
2109 pub fn max_identical_tool_calls(mut self, max: u32) -> Self {
2111 self.max_identical_tool_calls = Some(max);
2112 self
2113 }
2114
2115 pub fn max_fuzzy_identical_tool_calls(mut self, max: u32) -> Self {
2117 self.max_fuzzy_identical_tool_calls = Some(max);
2118 self
2119 }
2120
2121 pub fn max_tool_calls_per_turn(mut self, cap: u32) -> Self {
2123 self.max_tool_calls_per_turn = Some(cap);
2124 self
2125 }
2126
2127 pub fn permission_rules(mut self, rules: super::permission::PermissionRuleset) -> Self {
2129 self.permission_rules = rules;
2130 self
2131 }
2132
2133 pub fn instruction_text(mut self, text: impl Into<String>) -> Self {
2135 let text = text.into();
2136 if !text.is_empty() {
2137 self.instruction_text = Some(text);
2138 }
2139 self
2140 }
2141
2142 pub fn observability_mode(mut self, mode: super::observability::ObservabilityMode) -> Self {
2144 self.observability_mode = Some(mode);
2145 self
2146 }
2147
2148 pub fn dispatch_mode(mut self, mode: DispatchMode) -> Self {
2154 self.dispatch_mode = mode;
2155 self
2156 }
2157
2158 pub fn workspace(mut self, path: impl Into<std::path::PathBuf>) -> Self {
2161 self.workspace = Some(path.into());
2162 self
2163 }
2164
2165 pub fn audit_trail(mut self, trail: Arc<dyn super::audit::AuditTrail>) -> Self {
2169 self.audit_trail = Some(trail);
2170 self
2171 }
2172
2173 pub fn audit_user_context(
2176 mut self,
2177 user_id: impl Into<String>,
2178 tenant_id: impl Into<String>,
2179 ) -> Self {
2180 self.audit_user_id = Some(user_id.into());
2181 self.audit_tenant_id = Some(tenant_id.into());
2182 self
2183 }
2184
2185 pub fn audit_delegation_chain(mut self, chain: Vec<String>) -> Self {
2187 self.audit_delegation_chain = chain;
2188 self
2189 }
2190
2191 pub fn build(mut self) -> Result<Orchestrator<P>, Error> {
2193 if self.multi_agent_prompt {
2195 for agent in &mut self.sub_agents {
2196 agent
2197 .system_prompt
2198 .push_str(&crate::agent::prompts::render_collab_prompt(
2199 &agent.name,
2200 &agent.description,
2201 ));
2202 }
2203 }
2204
2205 {
2207 let mut seen = std::collections::HashSet::new();
2208 for agent in &self.sub_agents {
2209 if agent.name.is_empty() {
2210 return Err(Error::Config("sub-agent name must not be empty".into()));
2211 }
2212 if !seen.insert(&agent.name) {
2213 return Err(Error::Config(format!(
2214 "duplicate sub-agent name: '{}'",
2215 agent.name
2216 )));
2217 }
2218 if agent.max_turns == Some(0) {
2219 return Err(Error::Config(format!(
2220 "sub-agent '{}': max_turns must be > 0",
2221 agent.name
2222 )));
2223 }
2224 if agent.max_tokens == Some(0) {
2225 return Err(Error::Config(format!(
2226 "sub-agent '{}': max_tokens must be > 0",
2227 agent.name
2228 )));
2229 }
2230 }
2231 }
2232
2233 if self.sub_agents.is_empty() {
2234 tracing::warn!(
2235 "orchestrator built with no sub-agents — delegate_task tool will list no agents"
2236 );
2237 }
2238
2239 let squads_enabled = if self.dispatch_mode == DispatchMode::Sequential {
2242 false
2243 } else {
2244 self.enable_squads.unwrap_or(self.sub_agents.len() >= 2)
2245 };
2246 if squads_enabled && self.sub_agents.len() < 2 {
2247 tracing::warn!(
2248 "enable_squads is true but fewer than 2 agents are registered — \
2249 form_squad requires at least 2 agents to be useful"
2250 );
2251 }
2252
2253 let tool_names: Vec<Vec<String>> = self
2254 .sub_agents
2255 .iter()
2256 .map(|a| a.tools.iter().map(|t| t.definition().name).collect())
2257 .collect();
2258 let triples: Vec<(&str, &str, &[String])> = self
2259 .sub_agents
2260 .iter()
2261 .zip(tool_names.iter())
2262 .map(|(a, names)| (a.name.as_str(), a.description.as_str(), names.as_slice()))
2263 .collect();
2264 let mut system = build_system_prompt(&triples, squads_enabled, self.dispatch_mode);
2265 if let (Some(uid), Some(tid)) = (&self.audit_user_id, &self.audit_tenant_id) {
2267 system.push_str(&format!(
2268 "\n---\nYou are operating on behalf of **{uid}** in organization **{tid}**.\nKeep this user's information private. Do not share their data with other users."
2269 ));
2270 }
2271 let cached_definition = build_delegate_tool_schema(&triples, self.dispatch_mode);
2272 let form_squad_definition = if squads_enabled {
2273 Some(build_form_squad_tool_schema(&triples))
2274 } else {
2275 None
2276 };
2277 drop(triples);
2279 drop(tool_names);
2280
2281 let sub_agent_tokens = Arc::new(Mutex::new(TokenUsage::default()));
2282
2283 let shared_provider = Arc::new(BoxedProvider::from_arc(self.provider.clone()));
2284
2285 let agent_pool = if squads_enabled {
2287 Some(self.sub_agents.clone())
2288 } else {
2289 None
2290 };
2291
2292 let resolved_mode = self
2293 .observability_mode
2294 .unwrap_or(super::observability::ObservabilityMode::Production);
2295
2296 let spawn_tool_data = if let Some(spawn_cfg) = self.spawn_config.take() {
2298 let mut tool_pool = std::collections::HashMap::new();
2299 for tool in &self.spawn_builtin_tools {
2300 let name = tool.definition().name;
2301 if spawn_cfg.tool_allowlist.contains(&name) {
2302 tool_pool.insert(name, tool.clone());
2303 }
2304 }
2305 for allowed in &spawn_cfg.tool_allowlist {
2307 if !tool_pool.contains_key(allowed) {
2308 return Err(Error::Config(format!(
2309 "orchestrator.spawn.tool_allowlist: unknown tool '{}'. \
2310 Available builtin tools: [{}]",
2311 allowed,
2312 self.spawn_builtin_tools
2313 .iter()
2314 .map(|t| t.definition().name)
2315 .collect::<Vec<_>>()
2316 .join(", ")
2317 )));
2318 }
2319 }
2320 tool_pool.remove("delegate_task");
2322 tool_pool.remove("form_squad");
2323 tool_pool.remove("spawn_agent");
2324
2325 let cached_definition = SpawnAgentTool::build_definition(&spawn_cfg);
2326
2327 system.push_str(
2328 "\n\n## Dynamic Agent Spawning\n\
2329 You also have the **spawn_agent** tool to create specialist agents at runtime \
2330 when no pre-configured agent fits the task. Use this as a secondary option — \
2331 prefer delegating to existing agents when they match the need.",
2332 );
2333
2334 Some((spawn_cfg, tool_pool, cached_definition))
2335 } else {
2336 None
2337 };
2338
2339 let delegate_tool: Arc<dyn Tool> = Arc::new(DelegateTaskTool {
2340 shared_provider: shared_provider.clone(),
2341 sub_agents: self.sub_agents,
2342 max_turns: self.max_turns,
2343 max_tokens: self.max_tokens,
2344 permission_rules: self.permission_rules.clone(),
2345 accumulated_tokens: sub_agent_tokens.clone(),
2346 shared_memory: self.shared_memory.clone(),
2347 memory_namespace_prefix: self.memory_namespace_prefix.clone(),
2348 blackboard: self.blackboard.clone(),
2349 knowledge_base: self.knowledge_base.clone(),
2350 cached_definition,
2351 on_event: self.on_event.clone(),
2352 on_text: self.on_text.clone(),
2353 lsp_manager: self.lsp_manager.clone(),
2354 observability_mode: resolved_mode,
2355 allow_shared_write: self.allow_shared_write,
2356 tenant_tracker: self.tenant_tracker.clone(),
2357 guardrails: self.guardrails.clone(),
2358 });
2359
2360 let mut runner_builder = AgentRunner::builder(self.provider)
2361 .name("orchestrator")
2362 .system_prompt(system)
2363 .tool(delegate_tool)
2364 .max_turns(self.max_turns)
2365 .max_tokens(self.max_tokens);
2366
2367 if let Some(agent_pool) = agent_pool {
2369 let squad_def = form_squad_definition.expect("squad definition computed when enabled");
2371 let form_squad_tool: Arc<dyn Tool> = Arc::new(FormSquadTool {
2372 shared_provider: shared_provider.clone(),
2373 agent_pool,
2374 default_max_turns: self.max_turns,
2375 default_max_tokens: self.max_tokens,
2376 permission_rules: self.permission_rules.clone(),
2377 accumulated_tokens: sub_agent_tokens.clone(),
2378 shared_memory: self.shared_memory.clone(),
2379 memory_namespace_prefix: self.memory_namespace_prefix.clone(),
2380 blackboard: self.blackboard.clone(),
2381 knowledge_base: self.knowledge_base.clone(),
2382 on_event: self.on_event.clone(),
2383 on_text: self.on_text.clone(),
2384 lsp_manager: self.lsp_manager.clone(),
2385 cached_definition: squad_def,
2386 observability_mode: resolved_mode,
2387 allow_shared_write: self.allow_shared_write,
2388 tenant_tracker: self.tenant_tracker.clone(),
2389 });
2390 runner_builder = runner_builder.tool(form_squad_tool);
2391 }
2392
2393 if let Some((spawn_cfg, tool_pool, spawn_def)) = spawn_tool_data {
2395 let spawn_tool: Arc<dyn Tool> = Arc::new(SpawnAgentTool {
2396 shared_provider: shared_provider.clone(),
2397 spawn_config: spawn_cfg,
2398 tool_pool,
2399 spawn_count: Arc::new(std::sync::atomic::AtomicU32::new(0)),
2400 spawned_names: Arc::new(Mutex::new(std::collections::HashSet::new())),
2401 accumulated_tokens: sub_agent_tokens.clone(),
2402 permission_rules: self.permission_rules.clone(),
2403 shared_memory: self.shared_memory.clone(),
2404 memory_namespace_prefix: self.memory_namespace_prefix.clone(),
2405 on_event: self.on_event.clone(),
2406 on_text: self.on_text.clone(),
2407 lsp_manager: self.lsp_manager.clone(),
2408 observability_mode: resolved_mode,
2409 workspace: self.workspace.clone(),
2410 guardrails: self.guardrails.clone(),
2411 audit_trail: self.audit_trail.clone(),
2412 audit_user_id: self.audit_user_id.clone(),
2413 audit_tenant_id: self.audit_tenant_id.clone(),
2414 audit_delegation_chain: self.audit_delegation_chain.clone(),
2415 cached_definition: spawn_def,
2416 tenant_tracker: self.tenant_tracker.clone(),
2417 });
2418 runner_builder = runner_builder.tool(spawn_tool);
2419 }
2420
2421 if let Some(ref memory) = self.shared_memory {
2424 let orch_ns = self
2425 .memory_namespace_prefix
2426 .as_deref()
2427 .unwrap_or("orchestrator");
2428 let mem_scope = crate::auth::TenantScope::from_audit_fields(
2429 self.audit_tenant_id.as_deref(),
2430 self.audit_user_id.as_deref(),
2431 );
2432 let mem_tools = crate::memory::shared_tools::shared_memory_tools(
2433 memory.clone(),
2434 orch_ns,
2435 mem_scope,
2436 self.allow_shared_write,
2437 );
2438 runner_builder = runner_builder.tools(mem_tools);
2439 }
2440
2441 if let Some(strategy) = self.context_strategy {
2442 runner_builder = runner_builder.context_strategy(strategy);
2443 }
2444 if let Some(threshold) = self.summarize_threshold {
2445 runner_builder = runner_builder.summarize_threshold(threshold);
2446 }
2447 if let Some(timeout) = self.tool_timeout {
2448 runner_builder = runner_builder.tool_timeout(timeout);
2449 }
2450 if let Some(max) = self.max_tool_output_bytes {
2451 runner_builder = runner_builder.max_tool_output_bytes(max);
2452 }
2453 if let Some(on_text) = self.on_text {
2454 runner_builder = runner_builder.on_text(on_text);
2455 }
2456 if let Some(on_approval) = self.on_approval {
2457 runner_builder = runner_builder.on_approval(on_approval);
2458 }
2459 if let Some(learned) = self.learned_permissions {
2460 runner_builder = runner_builder.learned_permissions(learned);
2461 }
2462 if let Some(lsp) = self.lsp_manager {
2463 runner_builder = runner_builder.lsp_manager(lsp);
2464 }
2465 if let Some(on_event) = self.on_event {
2466 runner_builder = runner_builder.on_event(on_event);
2467 }
2468 if !self.guardrails.is_empty() {
2469 runner_builder = runner_builder.guardrails(self.guardrails);
2470 }
2471 if let Some(on_question) = self.on_question {
2472 runner_builder = runner_builder.on_question(on_question);
2473 }
2474 if let Some(timeout) = self.run_timeout {
2475 runner_builder = runner_builder.run_timeout(timeout);
2476 }
2477 if let Some(effort) = self.reasoning_effort {
2478 runner_builder = runner_builder.reasoning_effort(effort);
2479 }
2480 if self.enable_reflection {
2481 runner_builder = runner_builder.enable_reflection(true);
2482 }
2483 if let Some(threshold) = self.tool_output_compression_threshold {
2484 runner_builder = runner_builder.tool_output_compression_threshold(threshold);
2485 }
2486 if let Some(max) = self.max_tools_per_turn {
2487 runner_builder = runner_builder.max_tools_per_turn(max);
2488 }
2489 if let Some(max) = self.max_identical_tool_calls {
2490 runner_builder = runner_builder.max_identical_tool_calls(max);
2491 }
2492 if let Some(max) = self.max_fuzzy_identical_tool_calls {
2493 runner_builder = runner_builder.max_fuzzy_identical_tool_calls(max);
2494 }
2495 if let Some(cap) = self.max_tool_calls_per_turn {
2496 runner_builder = runner_builder.max_tool_calls_per_turn(cap);
2497 }
2498 if !self.permission_rules.is_empty() {
2499 runner_builder = runner_builder.permission_rules(self.permission_rules);
2500 }
2501 if let Some(text) = self.instruction_text {
2502 runner_builder = runner_builder.instruction_text(text);
2503 }
2504 if let Some(mode) = self.observability_mode {
2505 runner_builder = runner_builder.observability_mode(mode);
2506 }
2507 if let Some(trail) = self.audit_trail {
2508 runner_builder = runner_builder.audit_trail(trail);
2509 }
2510 if let Some(uid) = self.audit_user_id
2511 && let Some(tid) = self.audit_tenant_id
2512 {
2513 runner_builder = runner_builder.audit_user_context(uid, tid);
2514 }
2515 if !self.audit_delegation_chain.is_empty() {
2516 runner_builder =
2517 runner_builder.audit_delegation_chain(self.audit_delegation_chain.clone());
2518 }
2519 if let Some(tracker) = self.tenant_tracker {
2520 runner_builder = runner_builder.tenant_tracker(tracker);
2521 }
2522
2523 let runner = runner_builder.build()?;
2524
2525 Ok(Orchestrator {
2526 runner,
2527 sub_agent_tokens,
2528 })
2529 }
2530}
2531
2532#[cfg(test)]
2533mod tests {
2534 use super::*;
2535 use crate::llm::types::{
2536 CompletionRequest, CompletionResponse, ContentBlock, StopReason, TokenUsage,
2537 };
2538 use crate::tool::ToolOutput;
2539 use std::sync::Mutex;
2540
2541 struct MockProvider {
2542 responses: Mutex<Vec<CompletionResponse>>,
2543 }
2544
2545 impl MockProvider {
2546 fn new(responses: Vec<CompletionResponse>) -> Self {
2547 Self {
2548 responses: Mutex::new(responses),
2549 }
2550 }
2551 }
2552
2553 impl LlmProvider for MockProvider {
2554 async fn complete(&self, _request: CompletionRequest) -> Result<CompletionResponse, Error> {
2555 let mut responses = self.responses.lock().expect("mock lock poisoned");
2556 if responses.is_empty() {
2557 return Err(Error::Agent("no more mock responses".into()));
2558 }
2559 Ok(responses.remove(0))
2560 }
2561
2562 fn model_name(&self) -> Option<&str> {
2563 Some("mock-model-v1")
2564 }
2565 }
2566
2567 struct MockTool {
2569 def: crate::llm::types::ToolDefinition,
2570 response: String,
2571 }
2572
2573 impl MockTool {
2574 fn new(name: &str, response: &str) -> Self {
2575 Self {
2576 def: crate::llm::types::ToolDefinition {
2577 name: name.into(),
2578 description: format!("Mock {name}"),
2579 input_schema: json!({"type": "object"}),
2580 },
2581 response: response.into(),
2582 }
2583 }
2584 }
2585
2586 impl crate::tool::Tool for MockTool {
2587 fn definition(&self) -> crate::llm::types::ToolDefinition {
2588 self.def.clone()
2589 }
2590
2591 fn execute(
2592 &self,
2593 _ctx: &crate::ExecutionContext,
2594 _input: serde_json::Value,
2595 ) -> std::pin::Pin<
2596 Box<dyn std::future::Future<Output = Result<ToolOutput, Error>> + Send + '_>,
2597 > {
2598 let response = self.response.clone();
2599 Box::pin(async move { Ok(ToolOutput::success(response)) })
2600 }
2601 }
2602
2603 #[test]
2604 fn system_prompt_includes_agents() {
2605 let tools_a = vec!["web_search".to_string(), "read_file".to_string()];
2606 let tools_b: Vec<String> = vec![];
2607 let agents: Vec<(&str, &str, &[String])> = vec![
2608 ("researcher", "Research specialist", tools_a.as_slice()),
2609 ("coder", "Coding expert", tools_b.as_slice()),
2610 ];
2611
2612 let prompt = build_system_prompt(&agents, false, DispatchMode::Parallel);
2613 assert!(prompt.contains("researcher"));
2614 assert!(prompt.contains("Research specialist"));
2615 assert!(prompt.contains("coder"));
2616 assert!(prompt.contains("Tools: web_search, read_file"));
2617 assert!(prompt.contains("Tools: (none)"));
2618 assert!(
2620 prompt.contains("Decision Process"),
2621 "prompt should contain Decision Process section: {prompt}"
2622 );
2623 assert!(
2624 prompt.contains("Effort Scaling"),
2625 "prompt should contain Effort Scaling section: {prompt}"
2626 );
2627 assert!(
2628 prompt.contains("Task Quality"),
2629 "prompt should contain Task Quality section: {prompt}"
2630 );
2631 assert!(
2632 prompt.contains("DECOMPOSE"),
2633 "prompt should contain decomposition guidance: {prompt}"
2634 );
2635 }
2636
2637 #[test]
2638 fn system_prompt_shows_tool_names() {
2639 let tools = vec![
2640 "web_search".to_string(),
2641 "read_file".to_string(),
2642 "knowledge_search".to_string(),
2643 ];
2644 let no_tools: Vec<String> = vec![];
2645 let agents: Vec<(&str, &str, &[String])> = vec![
2646 ("researcher", "Research specialist", tools.as_slice()),
2647 ("analyst", "Analytical thinker", no_tools.as_slice()),
2648 ];
2649
2650 let prompt = build_system_prompt(&agents, false, DispatchMode::Parallel);
2651 assert!(
2652 prompt.contains("Tools: web_search, read_file, knowledge_search"),
2653 "prompt should list tool names: {prompt}"
2654 );
2655 assert!(
2656 prompt.contains("Tools: (none)"),
2657 "prompt should show (none) for agents without tools: {prompt}"
2658 );
2659 }
2660
2661 #[test]
2662 fn system_prompt_sequential_says_one_at_a_time() {
2663 let tools: Vec<String> = vec![];
2664 let agents: Vec<(&str, &str, &[String])> =
2665 vec![("builder", "Builds stuff", tools.as_slice())];
2666 let prompt = build_system_prompt(&agents, false, DispatchMode::Sequential);
2667 assert!(
2668 prompt.contains("ONE agent at a time"),
2669 "sequential prompt should say one at a time: {prompt}"
2670 );
2671 assert!(
2672 !prompt.contains("parallel execution"),
2673 "sequential prompt should not mention parallel: {prompt}"
2674 );
2675 }
2676
2677 #[test]
2678 fn system_prompt_parallel_says_parallel() {
2679 let tools: Vec<String> = vec![];
2680 let agents: Vec<(&str, &str, &[String])> =
2681 vec![("builder", "Builds stuff", tools.as_slice())];
2682 let prompt = build_system_prompt(&agents, false, DispatchMode::Parallel);
2683 assert!(
2684 prompt.contains("parallel execution"),
2685 "parallel prompt should mention parallel: {prompt}"
2686 );
2687 }
2688
2689 #[test]
2690 fn delegate_schema_sequential_max_items_1() {
2691 let tools: Vec<String> = vec![];
2692 let agents: Vec<(&str, &str, &[String])> =
2693 vec![("builder", "Builds stuff", tools.as_slice())];
2694 let def = build_delegate_tool_schema(&agents, DispatchMode::Sequential);
2695 let tasks = &def.input_schema["properties"]["tasks"];
2696 assert_eq!(
2697 tasks["maxItems"], 1,
2698 "sequential schema should have maxItems=1: {tasks}"
2699 );
2700 assert!(
2701 def.description.contains("ONE sub-agent"),
2702 "sequential description should say ONE: {}",
2703 def.description
2704 );
2705 }
2706
2707 #[test]
2708 fn delegate_schema_parallel_no_max_items() {
2709 let tools: Vec<String> = vec![];
2710 let agents: Vec<(&str, &str, &[String])> =
2711 vec![("builder", "Builds stuff", tools.as_slice())];
2712 let def = build_delegate_tool_schema(&agents, DispatchMode::Parallel);
2713 let tasks = &def.input_schema["properties"]["tasks"];
2714 assert!(
2715 tasks.get("maxItems").is_none(),
2716 "parallel schema should not have maxItems: {tasks}"
2717 );
2718 }
2719
2720 #[tokio::test]
2721 async fn sequential_dispatch_disables_squads() {
2722 let provider = Arc::new(MockProvider::new(vec![CompletionResponse {
2726 content: vec![ContentBlock::Text {
2727 text: "done".into(),
2728 }],
2729 stop_reason: StopReason::EndTurn,
2730 usage: TokenUsage::default(),
2731 model: None,
2732 }]));
2733 let mut orch = Orchestrator::builder(provider)
2734 .sub_agent("a", "Agent A", "prompt a")
2735 .sub_agent("b", "Agent B", "prompt b")
2736 .dispatch_mode(DispatchMode::Sequential)
2737 .build()
2738 .unwrap();
2739 let output = orch.run("test").await.unwrap();
2740 assert_eq!(output.result, "done");
2741 }
2744
2745 #[test]
2746 fn sequential_dispatch_disables_squads_in_prompt() {
2747 let tools: Vec<String> = vec![];
2748 let agents: Vec<(&str, &str, &[String])> = vec![
2749 ("a", "Agent A", tools.as_slice()),
2750 ("b", "Agent B", tools.as_slice()),
2751 ];
2752 let prompt = build_system_prompt(&agents, false, DispatchMode::Sequential);
2754 assert!(
2755 !prompt.contains("form_squad"),
2756 "sequential prompt should not mention form_squad: {prompt}"
2757 );
2758 }
2759
2760 #[test]
2761 fn delegate_tool_schema_includes_agents() {
2762 let tools = vec!["web_search".to_string()];
2763 let agents: Vec<(&str, &str, &[String])> =
2764 vec![("researcher", "Research", tools.as_slice())];
2765 let def = build_delegate_tool_schema(&agents, DispatchMode::Parallel);
2766 assert_eq!(def.name, "delegate_task");
2767 assert!(def.description.contains("researcher"));
2768 assert!(
2769 def.description.contains("web_search"),
2770 "delegate tool description should contain tool names: {}",
2771 def.description
2772 );
2773 assert!(
2774 def.description.contains("isolation"),
2775 "delegate tool description should mention isolation: {}",
2776 def.description
2777 );
2778 assert!(
2779 def.description.contains("self-contained"),
2780 "delegate tool description should mention self-contained tasks: {}",
2781 def.description
2782 );
2783 }
2784
2785 #[test]
2786 fn delegate_tool_definition_includes_agents() {
2787 let agents: Vec<(&str, &str, &[String])> = vec![("researcher", "Research", &[])];
2788 let cached_definition = build_delegate_tool_schema(&agents, DispatchMode::Parallel);
2789
2790 let tool = DelegateTaskTool {
2791 shared_provider: Arc::new(BoxedProvider::new(MockProvider::new(vec![]))),
2792 sub_agents: vec![SubAgentDef {
2793 name: "researcher".into(),
2794 description: "Research".into(),
2795 system_prompt: "prompt".into(),
2796 tools: vec![],
2797 context_strategy: None,
2798 summarize_threshold: None,
2799 tool_timeout: None,
2800 max_tool_output_bytes: None,
2801 max_turns: None,
2802 max_tokens: None,
2803 response_schema: None,
2804 run_timeout: None,
2805 guardrails: vec![],
2806 provider_override: None,
2807 reasoning_effort: None,
2808 enable_reflection: None,
2809 tool_output_compression_threshold: None,
2810 max_tools_per_turn: None,
2811 tool_profile: None,
2812 max_identical_tool_calls: None,
2813 max_fuzzy_identical_tool_calls: None,
2814 max_tool_calls_per_turn: None,
2815 session_prune_config: None,
2816 enable_recursive_summarization: None,
2817 reflection_threshold: None,
2818 consolidate_on_exit: None,
2819 workspace: None,
2820 max_total_tokens: None,
2821 audit_trail: None,
2822 audit_user_id: None,
2823 audit_tenant_id: None,
2824 audit_delegation_chain: Vec::new(),
2825 }],
2826 shared_memory: None,
2827 memory_namespace_prefix: None,
2828 blackboard: None,
2829 knowledge_base: None,
2830 max_turns: 10,
2831 max_tokens: 4096,
2832 permission_rules: crate::agent::permission::PermissionRuleset::default(),
2833 accumulated_tokens: Arc::new(Mutex::new(TokenUsage::default())),
2834 cached_definition,
2835 on_event: None,
2836 on_text: None,
2837 lsp_manager: None,
2838 observability_mode: crate::ObservabilityMode::Production,
2839 allow_shared_write: true,
2840 tenant_tracker: None,
2841 guardrails: vec![],
2842 };
2843
2844 let def = tool.definition();
2845 assert_eq!(def.name, "delegate_task");
2846 assert!(def.description.contains("researcher"));
2847 assert!(
2848 def.description.contains("tools"),
2849 "delegate tool description should contain 'tools' key: {}",
2850 def.description
2851 );
2852 }
2853
2854 #[test]
2855 fn build_errors_on_duplicate_sub_agent_names() {
2856 let provider = Arc::new(MockProvider::new(vec![]));
2857 let result = Orchestrator::builder(provider)
2858 .sub_agent("researcher", "Research 1", "prompt1")
2859 .sub_agent("researcher", "Research 2", "prompt2")
2860 .build();
2861 assert!(result.is_err());
2862 let err = result.err().unwrap();
2863 assert!(
2864 err.to_string()
2865 .contains("duplicate sub-agent name: 'researcher'"),
2866 "error: {err}"
2867 );
2868 }
2869
2870 #[tokio::test]
2871 async fn orchestrator_direct_response_no_delegation() {
2872 let provider = Arc::new(MockProvider::new(vec![CompletionResponse {
2873 content: vec![ContentBlock::Text {
2874 text: "Simple answer.".into(),
2875 }],
2876 stop_reason: StopReason::EndTurn,
2877 usage: TokenUsage {
2878 input_tokens: 10,
2879 output_tokens: 5,
2880 ..Default::default()
2881 },
2882 model: None,
2883 }]));
2884
2885 let mut orch = Orchestrator::builder(provider)
2886 .sub_agent("researcher", "Research", "prompt")
2887 .build()
2888 .unwrap();
2889
2890 let output = orch.run("simple question").await.unwrap();
2891 assert_eq!(output.result, "Simple answer.");
2892 assert_eq!(output.tool_calls_made, 0);
2893 }
2894
2895 #[tokio::test]
2896 async fn orchestrator_delegates_and_synthesizes() {
2897 let provider = Arc::new(MockProvider::new(vec![
2901 CompletionResponse {
2903 content: vec![ContentBlock::ToolUse {
2904 id: "call-1".into(),
2905 name: "delegate_task".into(),
2906 input: json!({
2907 "tasks": [
2908 {"agent": "researcher", "task": "Research Rust"},
2909 {"agent": "analyst", "task": "Analyze findings"}
2910 ]
2911 }),
2912 }],
2913 stop_reason: StopReason::ToolUse,
2914 usage: TokenUsage {
2915 input_tokens: 50,
2916 output_tokens: 20,
2917 ..Default::default()
2918 },
2919 model: None,
2920 },
2921 CompletionResponse {
2923 content: vec![ContentBlock::Text {
2924 text: "Rust is fast and safe.".into(),
2925 }],
2926 stop_reason: StopReason::EndTurn,
2927 usage: TokenUsage {
2928 input_tokens: 10,
2929 output_tokens: 8,
2930 ..Default::default()
2931 },
2932 model: None,
2933 },
2934 CompletionResponse {
2936 content: vec![ContentBlock::Text {
2937 text: "Strengths: memory safety, performance.".into(),
2938 }],
2939 stop_reason: StopReason::EndTurn,
2940 usage: TokenUsage {
2941 input_tokens: 12,
2942 output_tokens: 10,
2943 ..Default::default()
2944 },
2945 model: None,
2946 },
2947 CompletionResponse {
2949 content: vec![ContentBlock::Text {
2950 text: "Based on research: Rust is excellent.".into(),
2951 }],
2952 stop_reason: StopReason::EndTurn,
2953 usage: TokenUsage {
2954 input_tokens: 80,
2955 output_tokens: 30,
2956 ..Default::default()
2957 },
2958 model: None,
2959 },
2960 ]));
2961
2962 let mut orch = Orchestrator::builder(provider)
2963 .sub_agent("researcher", "Research specialist", "You research.")
2964 .sub_agent("analyst", "Analysis expert", "You analyze.")
2965 .build()
2966 .unwrap();
2967
2968 let output = orch.run("Analyze Rust").await.unwrap();
2969 assert_eq!(output.result, "Based on research: Rust is excellent.");
2970 assert_eq!(output.tool_calls_made, 1); assert_eq!(output.tokens_used.input_tokens, 50 + 80 + 10 + 12);
2973 assert_eq!(output.tokens_used.output_tokens, 20 + 30 + 8 + 10);
2974 }
2975
2976 #[tokio::test]
2977 async fn orchestrator_handles_unknown_agent_gracefully() {
2978 let provider = Arc::new(MockProvider::new(vec![
2980 CompletionResponse {
2981 content: vec![ContentBlock::ToolUse {
2982 id: "call-1".into(),
2983 name: "delegate_task".into(),
2984 input: json!({
2985 "tasks": [{"agent": "nonexistent", "task": "do stuff"}]
2986 }),
2987 }],
2988 stop_reason: StopReason::ToolUse,
2989 usage: TokenUsage::default(),
2990 model: None,
2991 },
2992 CompletionResponse {
2994 content: vec![ContentBlock::Text {
2995 text: "No such agent available.".into(),
2996 }],
2997 stop_reason: StopReason::EndTurn,
2998 usage: TokenUsage::default(),
2999 model: None,
3000 },
3001 ]));
3002
3003 let mut orch = Orchestrator::builder(provider)
3004 .sub_agent("researcher", "Research", "prompt")
3005 .build()
3006 .unwrap();
3007
3008 let output = orch.run("delegate to unknown").await.unwrap();
3009 assert_eq!(output.result, "No such agent available.");
3010 }
3011
3012 #[tokio::test]
3013 async fn orchestrator_handles_invalid_tool_name() {
3014 let provider = Arc::new(MockProvider::new(vec![
3015 CompletionResponse {
3016 content: vec![ContentBlock::ToolUse {
3017 id: "call-1".into(),
3018 name: "wrong_tool".into(),
3019 input: json!({}),
3020 }],
3021 stop_reason: StopReason::ToolUse,
3022 usage: TokenUsage::default(),
3023 model: None,
3024 },
3025 CompletionResponse {
3026 content: vec![ContentBlock::Text {
3027 text: "Sorry, let me respond directly.".into(),
3028 }],
3029 stop_reason: StopReason::EndTurn,
3030 usage: TokenUsage::default(),
3031 model: None,
3032 },
3033 ]));
3034
3035 let mut orch = Orchestrator::builder(provider)
3036 .sub_agent("researcher", "Research", "prompt")
3037 .build()
3038 .unwrap();
3039
3040 let output = orch.run("do something").await.unwrap();
3041 assert_eq!(output.result, "Sorry, let me respond directly.");
3042 }
3043
3044 #[tokio::test]
3045 async fn orchestrator_handles_empty_delegate_tasks() {
3046 let provider = Arc::new(MockProvider::new(vec![
3047 CompletionResponse {
3049 content: vec![ContentBlock::ToolUse {
3050 id: "call-1".into(),
3051 name: "delegate_task".into(),
3052 input: json!({"tasks": []}),
3053 }],
3054 stop_reason: StopReason::ToolUse,
3055 usage: TokenUsage::default(),
3056 model: None,
3057 },
3058 CompletionResponse {
3060 content: vec![ContentBlock::Text {
3061 text: "Let me try again properly.".into(),
3062 }],
3063 stop_reason: StopReason::EndTurn,
3064 usage: TokenUsage::default(),
3065 model: None,
3066 },
3067 ]));
3068
3069 let mut orch = Orchestrator::builder(provider)
3070 .sub_agent("researcher", "Research", "prompt")
3071 .build()
3072 .unwrap();
3073
3074 let output = orch.run("do something").await.unwrap();
3075 assert_eq!(output.result, "Let me try again properly.");
3076 }
3077
3078 #[tokio::test]
3079 async fn orchestrator_handles_missing_tasks_field() {
3080 let provider = Arc::new(MockProvider::new(vec![
3081 CompletionResponse {
3083 content: vec![ContentBlock::ToolUse {
3084 id: "call-1".into(),
3085 name: "delegate_task".into(),
3086 input: json!({}),
3087 }],
3088 stop_reason: StopReason::ToolUse,
3089 usage: TokenUsage::default(),
3090 model: None,
3091 },
3092 CompletionResponse {
3094 content: vec![ContentBlock::Text {
3095 text: "I need to format correctly.".into(),
3096 }],
3097 stop_reason: StopReason::EndTurn,
3098 usage: TokenUsage::default(),
3099 model: None,
3100 },
3101 ]));
3102
3103 let mut orch = Orchestrator::builder(provider)
3104 .sub_agent("researcher", "Research", "prompt")
3105 .build()
3106 .unwrap();
3107
3108 let output = orch.run("do something").await.unwrap();
3109 assert_eq!(output.result, "I need to format correctly.");
3110 }
3111
3112 #[tokio::test]
3113 async fn blackboard_populated_after_delegation() {
3114 use crate::agent::blackboard::InMemoryBlackboard;
3115
3116 let bb = Arc::new(InMemoryBlackboard::new());
3117
3118 let provider = Arc::new(MockProvider::new(vec![
3119 CompletionResponse {
3121 content: vec![ContentBlock::ToolUse {
3122 id: "call-1".into(),
3123 name: "delegate_task".into(),
3124 input: json!({
3125 "tasks": [{"agent": "researcher", "task": "Find info"}]
3126 }),
3127 }],
3128 stop_reason: StopReason::ToolUse,
3129 usage: TokenUsage::default(),
3130 model: None,
3131 },
3132 CompletionResponse {
3134 content: vec![ContentBlock::Text {
3135 text: "Research result here.".into(),
3136 }],
3137 stop_reason: StopReason::EndTurn,
3138 usage: TokenUsage::default(),
3139 model: None,
3140 },
3141 CompletionResponse {
3143 content: vec![ContentBlock::Text {
3144 text: "Done.".into(),
3145 }],
3146 stop_reason: StopReason::EndTurn,
3147 usage: TokenUsage::default(),
3148 model: None,
3149 },
3150 ]));
3151
3152 let mut orch = Orchestrator::builder(provider)
3153 .sub_agent("researcher", "Research specialist", "You research.")
3154 .blackboard(bb.clone())
3155 .build()
3156 .unwrap();
3157
3158 orch.run("research something").await.unwrap();
3159
3160 let val: Option<serde_json::Value> = bb.read("agent:researcher").await.unwrap();
3162 assert!(val.is_some(), "blackboard should have agent:researcher key");
3163 assert_eq!(
3164 val.unwrap(),
3165 serde_json::Value::String("Research result here.".into())
3166 );
3167 }
3168
3169 #[tokio::test]
3170 async fn sub_agents_receive_blackboard_tools() {
3171 use crate::agent::blackboard::InMemoryBlackboard;
3172 use crate::llm::types::CompletionRequest;
3173
3174 struct ToolTrackingProvider {
3176 responses: Mutex<Vec<CompletionResponse>>,
3177 tool_names_seen: Mutex<Vec<Vec<String>>>,
3178 }
3179
3180 impl LlmProvider for ToolTrackingProvider {
3181 async fn complete(
3182 &self,
3183 request: CompletionRequest,
3184 ) -> Result<CompletionResponse, Error> {
3185 let names: Vec<String> = request.tools.iter().map(|t| t.name.clone()).collect();
3186 self.tool_names_seen.lock().expect("lock").push(names);
3187
3188 let mut responses = self.responses.lock().expect("lock");
3189 if responses.is_empty() {
3190 return Err(Error::Agent("no more mock responses".into()));
3191 }
3192 Ok(responses.remove(0))
3193 }
3194 }
3195
3196 let bb = Arc::new(InMemoryBlackboard::new());
3197
3198 let provider = Arc::new(ToolTrackingProvider {
3199 responses: Mutex::new(vec![
3200 CompletionResponse {
3202 content: vec![ContentBlock::ToolUse {
3203 id: "call-1".into(),
3204 name: "delegate_task".into(),
3205 input: json!({
3206 "tasks": [{"agent": "worker", "task": "do work"}]
3207 }),
3208 }],
3209 stop_reason: StopReason::ToolUse,
3210 usage: TokenUsage::default(),
3211 model: None,
3212 },
3213 CompletionResponse {
3215 content: vec![ContentBlock::Text {
3216 text: "Work done.".into(),
3217 }],
3218 stop_reason: StopReason::EndTurn,
3219 usage: TokenUsage::default(),
3220 model: None,
3221 },
3222 CompletionResponse {
3224 content: vec![ContentBlock::Text {
3225 text: "All done.".into(),
3226 }],
3227 stop_reason: StopReason::EndTurn,
3228 usage: TokenUsage::default(),
3229 model: None,
3230 },
3231 ]),
3232 tool_names_seen: Mutex::new(vec![]),
3233 });
3234
3235 let mut orch = Orchestrator::builder(provider.clone())
3236 .sub_agent("worker", "Worker agent", "You work.")
3237 .blackboard(bb)
3238 .build()
3239 .unwrap();
3240
3241 orch.run("do work").await.unwrap();
3242
3243 let all_tool_names = provider.tool_names_seen.lock().expect("lock");
3245 assert!(
3246 all_tool_names.len() >= 2,
3247 "expected at least 2 LLM calls, got {}",
3248 all_tool_names.len()
3249 );
3250 let sub_agent_tools = &all_tool_names[1];
3251 assert!(
3252 sub_agent_tools.contains(&"blackboard_read".to_string()),
3253 "sub-agent should have blackboard_read tool, got: {sub_agent_tools:?}"
3254 );
3255 assert!(
3256 sub_agent_tools.contains(&"blackboard_write".to_string()),
3257 "sub-agent should have blackboard_write tool, got: {sub_agent_tools:?}"
3258 );
3259 assert!(
3260 sub_agent_tools.contains(&"blackboard_list".to_string()),
3261 "sub-agent should have blackboard_list tool, got: {sub_agent_tools:?}"
3262 );
3263 }
3264
3265 #[test]
3266 fn blackboard_builder_method_works() {
3267 use crate::agent::blackboard::InMemoryBlackboard;
3268
3269 let bb = Arc::new(InMemoryBlackboard::new());
3270 let provider = Arc::new(MockProvider::new(vec![]));
3271
3272 let result = Orchestrator::builder(provider)
3274 .sub_agent("agent1", "Agent one", "You are agent 1.")
3275 .blackboard(bb)
3276 .build();
3277
3278 assert!(result.is_ok());
3279 }
3280
3281 #[test]
3282 fn knowledge_builder_method_works() {
3283 use crate::knowledge::in_memory::InMemoryKnowledgeBase;
3284
3285 let kb: Arc<dyn KnowledgeBase> = Arc::new(InMemoryKnowledgeBase::new());
3286 let provider = Arc::new(MockProvider::new(vec![]));
3287
3288 let result = Orchestrator::builder(provider)
3289 .sub_agent("agent1", "Agent one", "You are agent 1.")
3290 .knowledge(kb)
3291 .build();
3292
3293 assert!(result.is_ok());
3294 }
3295
3296 #[tokio::test]
3297 async fn sub_agents_receive_knowledge_tools() {
3298 use crate::knowledge::in_memory::InMemoryKnowledgeBase;
3299 use crate::llm::types::CompletionRequest;
3300
3301 struct ToolTrackingProvider {
3302 responses: Mutex<Vec<CompletionResponse>>,
3303 tool_names_seen: Mutex<Vec<Vec<String>>>,
3304 }
3305
3306 impl LlmProvider for ToolTrackingProvider {
3307 async fn complete(
3308 &self,
3309 request: CompletionRequest,
3310 ) -> Result<CompletionResponse, Error> {
3311 let names: Vec<String> = request.tools.iter().map(|t| t.name.clone()).collect();
3312 self.tool_names_seen.lock().expect("lock").push(names);
3313
3314 let mut responses = self.responses.lock().expect("lock");
3315 if responses.is_empty() {
3316 return Err(Error::Agent("no more mock responses".into()));
3317 }
3318 Ok(responses.remove(0))
3319 }
3320 }
3321
3322 let kb: Arc<dyn KnowledgeBase> = Arc::new(InMemoryKnowledgeBase::new());
3323
3324 let provider = Arc::new(ToolTrackingProvider {
3325 responses: Mutex::new(vec![
3326 CompletionResponse {
3328 content: vec![ContentBlock::ToolUse {
3329 id: "call-1".into(),
3330 name: "delegate_task".into(),
3331 input: json!({
3332 "tasks": [{"agent": "worker", "task": "do work"}]
3333 }),
3334 }],
3335 stop_reason: StopReason::ToolUse,
3336 usage: TokenUsage::default(),
3337 model: None,
3338 },
3339 CompletionResponse {
3341 content: vec![ContentBlock::Text {
3342 text: "Work done.".into(),
3343 }],
3344 stop_reason: StopReason::EndTurn,
3345 usage: TokenUsage::default(),
3346 model: None,
3347 },
3348 CompletionResponse {
3350 content: vec![ContentBlock::Text {
3351 text: "All done.".into(),
3352 }],
3353 stop_reason: StopReason::EndTurn,
3354 usage: TokenUsage::default(),
3355 model: None,
3356 },
3357 ]),
3358 tool_names_seen: Mutex::new(vec![]),
3359 });
3360
3361 let mut orch = Orchestrator::builder(provider.clone())
3362 .sub_agent("worker", "Worker agent", "You work.")
3363 .knowledge(kb)
3364 .build()
3365 .unwrap();
3366
3367 orch.run("do work").await.unwrap();
3368
3369 let all_tool_names = provider.tool_names_seen.lock().expect("lock");
3370 assert!(
3371 all_tool_names.len() >= 2,
3372 "expected at least 2 LLM calls, got {}",
3373 all_tool_names.len()
3374 );
3375 let sub_agent_tools = &all_tool_names[1];
3376 assert!(
3377 sub_agent_tools.contains(&"knowledge_search".to_string()),
3378 "sub-agent should have knowledge_search tool, got: {sub_agent_tools:?}"
3379 );
3380 }
3381
3382 #[tokio::test]
3383 async fn orchestrator_accumulates_cache_tokens_through_delegation() {
3384 let provider = Arc::new(MockProvider::new(vec![
3385 CompletionResponse {
3387 content: vec![ContentBlock::ToolUse {
3388 id: "call-1".into(),
3389 name: "delegate_task".into(),
3390 input: json!({
3391 "tasks": [{"agent": "researcher", "task": "Research Rust"}]
3392 }),
3393 }],
3394 stop_reason: StopReason::ToolUse,
3395 usage: TokenUsage {
3396 input_tokens: 50,
3397 output_tokens: 20,
3398 cache_creation_input_tokens: 100,
3399 cache_read_input_tokens: 0,
3400 reasoning_tokens: 0,
3401 },
3402 model: None,
3403 },
3404 CompletionResponse {
3406 content: vec![ContentBlock::Text {
3407 text: "Rust is fast.".into(),
3408 }],
3409 stop_reason: StopReason::EndTurn,
3410 usage: TokenUsage {
3411 input_tokens: 10,
3412 output_tokens: 8,
3413 cache_creation_input_tokens: 0,
3414 cache_read_input_tokens: 30,
3415 reasoning_tokens: 0,
3416 },
3417 model: None,
3418 },
3419 CompletionResponse {
3421 content: vec![ContentBlock::Text {
3422 text: "Rust is excellent.".into(),
3423 }],
3424 stop_reason: StopReason::EndTurn,
3425 usage: TokenUsage {
3426 input_tokens: 80,
3427 output_tokens: 30,
3428 cache_creation_input_tokens: 0,
3429 cache_read_input_tokens: 90,
3430 reasoning_tokens: 0,
3431 },
3432 model: None,
3433 },
3434 ]));
3435
3436 let mut orch = Orchestrator::builder(provider)
3437 .sub_agent("researcher", "Research specialist", "You research.")
3438 .build()
3439 .unwrap();
3440
3441 let output = orch.run("Analyze Rust").await.unwrap();
3442 assert_eq!(output.tokens_used.input_tokens, 50 + 80 + 10);
3445 assert_eq!(output.tokens_used.output_tokens, 20 + 30 + 8);
3446 assert_eq!(output.tokens_used.cache_creation_input_tokens, 100);
3447 assert_eq!(output.tokens_used.cache_read_input_tokens, 90 + 30);
3448 }
3449
3450 #[tokio::test]
3451 async fn orchestrator_error_includes_sub_agent_tokens() {
3452 let provider = Arc::new(MockProvider::new(vec![
3456 CompletionResponse {
3458 content: vec![ContentBlock::ToolUse {
3459 id: "call-1".into(),
3460 name: "delegate_task".into(),
3461 input: json!({
3462 "tasks": [{"agent": "researcher", "task": "Research Rust"}]
3463 }),
3464 }],
3465 stop_reason: StopReason::ToolUse,
3466 usage: TokenUsage {
3467 input_tokens: 50,
3468 output_tokens: 20,
3469 ..Default::default()
3470 },
3471 model: None,
3472 },
3473 CompletionResponse {
3475 content: vec![ContentBlock::Text {
3476 text: "Rust is fast.".into(),
3477 }],
3478 stop_reason: StopReason::EndTurn,
3479 usage: TokenUsage {
3480 input_tokens: 15,
3481 output_tokens: 10,
3482 ..Default::default()
3483 },
3484 model: None,
3485 },
3486 CompletionResponse {
3488 content: vec![ContentBlock::ToolUse {
3489 id: "call-2".into(),
3490 name: "delegate_task".into(),
3491 input: json!({
3492 "tasks": [{"agent": "researcher", "task": "More research"}]
3493 }),
3494 }],
3495 stop_reason: StopReason::ToolUse,
3496 usage: TokenUsage {
3497 input_tokens: 80,
3498 output_tokens: 25,
3499 ..Default::default()
3500 },
3501 model: None,
3502 },
3503 CompletionResponse {
3505 content: vec![ContentBlock::Text {
3506 text: "More info.".into(),
3507 }],
3508 stop_reason: StopReason::EndTurn,
3509 usage: TokenUsage {
3510 input_tokens: 12,
3511 output_tokens: 8,
3512 ..Default::default()
3513 },
3514 model: None,
3515 },
3516 ]));
3517
3518 let mut orch = Orchestrator::builder(provider)
3519 .sub_agent("researcher", "Research", "prompt")
3520 .max_turns(2)
3521 .build()
3522 .unwrap();
3523
3524 let err = orch.run("research deeply").await.unwrap_err();
3525
3526 match &err {
3528 Error::WithPartialUsage { source, .. } => {
3529 assert!(
3530 matches!(**source, Error::MaxTurnsExceeded(2)),
3531 "inner error should be MaxTurnsExceeded(2), got: {source}"
3532 );
3533 }
3534 other => panic!("expected WithPartialUsage, got: {other}"),
3535 }
3536
3537 let usage = err.partial_usage();
3538 assert_eq!(
3540 usage.input_tokens,
3541 50 + 80 + 15 + 12,
3542 "input tokens: orchestrator(50+80) + sub-agent(15+12)"
3543 );
3544 assert_eq!(
3545 usage.output_tokens,
3546 20 + 25 + 10 + 8,
3547 "output tokens: orchestrator(20+25) + sub-agent(10+8)"
3548 );
3549 }
3550
3551 #[tokio::test]
3552 async fn on_event_emits_sub_agent_events() {
3553 use crate::agent::events::AgentEvent;
3554
3555 let events: Arc<std::sync::Mutex<Vec<AgentEvent>>> =
3556 Arc::new(std::sync::Mutex::new(vec![]));
3557 let events_clone = events.clone();
3558
3559 let provider = Arc::new(MockProvider::new(vec![
3560 CompletionResponse {
3562 content: vec![ContentBlock::ToolUse {
3563 id: "call-1".into(),
3564 name: "delegate_task".into(),
3565 input: json!({
3566 "tasks": [{"agent": "researcher", "task": "Research Rust"}]
3567 }),
3568 }],
3569 stop_reason: StopReason::ToolUse,
3570 usage: TokenUsage::default(),
3571 model: None,
3572 },
3573 CompletionResponse {
3575 content: vec![ContentBlock::Text {
3576 text: "Rust is fast.".into(),
3577 }],
3578 stop_reason: StopReason::EndTurn,
3579 usage: TokenUsage {
3580 input_tokens: 10,
3581 output_tokens: 5,
3582 ..Default::default()
3583 },
3584 model: None,
3585 },
3586 CompletionResponse {
3588 content: vec![ContentBlock::Text {
3589 text: "Summary: Rust is fast.".into(),
3590 }],
3591 stop_reason: StopReason::EndTurn,
3592 usage: TokenUsage::default(),
3593 model: None,
3594 },
3595 ]));
3596
3597 let mut orch = Orchestrator::builder(provider)
3598 .sub_agent("researcher", "Research", "prompt")
3599 .on_event(Arc::new(move |e| {
3600 events_clone.lock().unwrap().push(e);
3601 }))
3602 .build()
3603 .unwrap();
3604
3605 orch.run("research task").await.unwrap();
3606
3607 let events = events.lock().unwrap();
3608
3609 let dispatched: Vec<_> = events
3610 .iter()
3611 .filter(|e| matches!(e, AgentEvent::SubAgentsDispatched { .. }))
3612 .collect();
3613 assert_eq!(dispatched.len(), 1, "expected 1 SubAgentsDispatched");
3614 match &dispatched[0] {
3615 AgentEvent::SubAgentsDispatched { agent, agents } => {
3616 assert_eq!(agent, "orchestrator");
3617 assert_eq!(agents, &["researcher"]);
3618 }
3619 _ => unreachable!(),
3620 }
3621
3622 let completed: Vec<_> = events
3623 .iter()
3624 .filter(|e| matches!(e, AgentEvent::SubAgentCompleted { .. }))
3625 .collect();
3626 assert_eq!(completed.len(), 1, "expected 1 SubAgentCompleted");
3627 match &completed[0] {
3628 AgentEvent::SubAgentCompleted {
3629 agent,
3630 success,
3631 usage,
3632 } => {
3633 assert_eq!(agent, "researcher");
3634 assert!(success);
3635 assert_eq!(usage.input_tokens, 10);
3636 }
3637 _ => unreachable!(),
3638 }
3639 }
3640
3641 #[tokio::test]
3642 async fn sub_agent_receives_guardrails() {
3643 use crate::agent::guardrail::Guardrail;
3644 use crate::llm::types::CompletionRequest;
3645
3646 struct MarkerGuardrail;
3648 impl Guardrail for MarkerGuardrail {
3649 fn pre_llm(
3650 &self,
3651 request: &mut CompletionRequest,
3652 ) -> std::pin::Pin<
3653 Box<dyn std::future::Future<Output = Result<(), crate::error::Error>> + Send + '_>,
3654 > {
3655 request.system = format!("{} [GUARDRAIL_ACTIVE]", request.system);
3656 Box::pin(async { Ok(()) })
3657 }
3658 }
3659
3660 struct CapturingProvider {
3661 responses: Mutex<Vec<CompletionResponse>>,
3662 systems_seen: Mutex<Vec<String>>,
3663 }
3664
3665 impl LlmProvider for CapturingProvider {
3666 async fn complete(
3667 &self,
3668 request: CompletionRequest,
3669 ) -> Result<CompletionResponse, crate::error::Error> {
3670 self.systems_seen
3671 .lock()
3672 .unwrap()
3673 .push(request.system.clone());
3674 let mut responses = self.responses.lock().unwrap();
3675 if responses.is_empty() {
3676 return Err(crate::error::Error::Agent("no more responses".into()));
3677 }
3678 Ok(responses.remove(0))
3679 }
3680 }
3681
3682 let guardrail: Arc<dyn Guardrail> = Arc::new(MarkerGuardrail);
3683
3684 let provider = Arc::new(CapturingProvider {
3685 responses: Mutex::new(vec![
3686 CompletionResponse {
3688 content: vec![ContentBlock::ToolUse {
3689 id: "call-1".into(),
3690 name: "delegate_task".into(),
3691 input: json!({
3692 "tasks": [{"agent": "worker", "task": "do work"}]
3693 }),
3694 }],
3695 stop_reason: StopReason::ToolUse,
3696 usage: TokenUsage::default(),
3697 model: None,
3698 },
3699 CompletionResponse {
3701 content: vec![ContentBlock::Text {
3702 text: "Work done.".into(),
3703 }],
3704 stop_reason: StopReason::EndTurn,
3705 usage: TokenUsage::default(),
3706 model: None,
3707 },
3708 CompletionResponse {
3710 content: vec![ContentBlock::Text {
3711 text: "All done.".into(),
3712 }],
3713 stop_reason: StopReason::EndTurn,
3714 usage: TokenUsage::default(),
3715 model: None,
3716 },
3717 ]),
3718 systems_seen: Mutex::new(vec![]),
3719 });
3720
3721 let mut orch = Orchestrator::builder(provider.clone())
3722 .sub_agent_full(SubAgentConfig {
3723 name: "worker".into(),
3724 description: "Worker agent".into(),
3725 system_prompt: "You work.".into(),
3726 tools: vec![],
3727 context_strategy: None,
3728 summarize_threshold: None,
3729 tool_timeout: None,
3730 max_tool_output_bytes: None,
3731 max_turns: None,
3732 max_tokens: None,
3733 response_schema: None,
3734 run_timeout: None,
3735 guardrails: vec![guardrail],
3736 provider: None,
3737 reasoning_effort: None,
3738 enable_reflection: None,
3739 tool_output_compression_threshold: None,
3740 max_tools_per_turn: None,
3741 tool_profile: None,
3742 max_identical_tool_calls: None,
3743 max_fuzzy_identical_tool_calls: None,
3744 max_tool_calls_per_turn: None,
3745 session_prune_config: None,
3746 enable_recursive_summarization: None,
3747 reflection_threshold: None,
3748 consolidate_on_exit: None,
3749 workspace: None,
3750 max_total_tokens: None,
3751 audit_trail: None,
3752 audit_user_id: None,
3753 audit_tenant_id: None,
3754 audit_delegation_chain: Vec::new(),
3755 })
3756 .build()
3757 .unwrap();
3758
3759 orch.run("do work").await.unwrap();
3760
3761 let systems = provider.systems_seen.lock().unwrap();
3763 assert!(
3764 systems.len() >= 2,
3765 "expected at least 2 LLM calls, got {}",
3766 systems.len()
3767 );
3768 assert!(
3770 systems[1].contains("[GUARDRAIL_ACTIVE]"),
3771 "sub-agent system prompt should contain guardrail marker: {}",
3772 systems[1]
3773 );
3774 assert!(
3776 !systems[0].contains("[GUARDRAIL_ACTIVE]"),
3777 "orchestrator system prompt should NOT contain guardrail marker: {}",
3778 systems[0]
3779 );
3780 }
3781
3782 #[tokio::test]
3788 async fn orchestrator_guardrails_propagate_to_delegated_sub_agents() {
3789 use crate::agent::guardrail::Guardrail;
3790 use crate::llm::types::CompletionRequest;
3791
3792 struct MarkerGuardrail;
3793 impl Guardrail for MarkerGuardrail {
3794 fn pre_llm(
3795 &self,
3796 request: &mut CompletionRequest,
3797 ) -> std::pin::Pin<
3798 Box<dyn std::future::Future<Output = Result<(), crate::error::Error>> + Send + '_>,
3799 > {
3800 request.system = format!("{} [ORCH_GUARD_ACTIVE]", request.system);
3801 Box::pin(async { Ok(()) })
3802 }
3803 }
3804
3805 struct CapturingProvider {
3806 responses: Mutex<Vec<CompletionResponse>>,
3807 systems_seen: Mutex<Vec<String>>,
3808 }
3809
3810 impl LlmProvider for CapturingProvider {
3811 async fn complete(
3812 &self,
3813 request: CompletionRequest,
3814 ) -> Result<CompletionResponse, crate::error::Error> {
3815 self.systems_seen
3816 .lock()
3817 .unwrap()
3818 .push(request.system.clone());
3819 let mut responses = self.responses.lock().unwrap();
3820 if responses.is_empty() {
3821 return Err(crate::error::Error::Agent("no more responses".into()));
3822 }
3823 Ok(responses.remove(0))
3824 }
3825 }
3826
3827 let guardrail: Arc<dyn Guardrail> = Arc::new(MarkerGuardrail);
3828
3829 let provider = Arc::new(CapturingProvider {
3830 responses: Mutex::new(vec![
3831 CompletionResponse {
3833 content: vec![ContentBlock::ToolUse {
3834 id: "call-1".into(),
3835 name: "delegate_task".into(),
3836 input: json!({
3837 "tasks": [{"agent": "worker", "task": "do work"}]
3838 }),
3839 }],
3840 stop_reason: StopReason::ToolUse,
3841 usage: TokenUsage::default(),
3842 model: None,
3843 },
3844 CompletionResponse {
3846 content: vec![ContentBlock::Text {
3847 text: "Work done.".into(),
3848 }],
3849 stop_reason: StopReason::EndTurn,
3850 usage: TokenUsage::default(),
3851 model: None,
3852 },
3853 CompletionResponse {
3855 content: vec![ContentBlock::Text {
3856 text: "Synthesized.".into(),
3857 }],
3858 stop_reason: StopReason::EndTurn,
3859 usage: TokenUsage::default(),
3860 model: None,
3861 },
3862 ]),
3863 systems_seen: Mutex::new(vec![]),
3864 });
3865
3866 let mut orch = Orchestrator::builder(provider.clone())
3869 .guardrail(guardrail)
3870 .sub_agent_full(SubAgentConfig {
3871 name: "worker".into(),
3872 description: "Worker agent".into(),
3873 system_prompt: "You work.".into(),
3874 tools: vec![],
3875 context_strategy: None,
3876 summarize_threshold: None,
3877 tool_timeout: None,
3878 max_tool_output_bytes: None,
3879 max_turns: None,
3880 max_tokens: None,
3881 response_schema: None,
3882 run_timeout: None,
3883 guardrails: vec![],
3884 provider: None,
3885 reasoning_effort: None,
3886 enable_reflection: None,
3887 tool_output_compression_threshold: None,
3888 max_tools_per_turn: None,
3889 tool_profile: None,
3890 max_identical_tool_calls: None,
3891 max_fuzzy_identical_tool_calls: None,
3892 max_tool_calls_per_turn: None,
3893 session_prune_config: None,
3894 enable_recursive_summarization: None,
3895 reflection_threshold: None,
3896 consolidate_on_exit: None,
3897 workspace: None,
3898 max_total_tokens: None,
3899 audit_trail: None,
3900 audit_user_id: None,
3901 audit_tenant_id: None,
3902 audit_delegation_chain: Vec::new(),
3903 })
3904 .build()
3905 .unwrap();
3906
3907 orch.run("do work").await.unwrap();
3908
3909 let systems = provider.systems_seen.lock().unwrap();
3910 assert!(systems.len() >= 2, "expected at least 2 LLM calls");
3911 assert!(
3914 systems[1].contains("[ORCH_GUARD_ACTIVE]"),
3915 "sub-agent system prompt should contain orchestrator guardrail marker; got: {}",
3916 systems[1]
3917 );
3918 }
3919
3920 #[test]
3921 fn build_rejects_sub_agent_with_zero_max_turns() {
3922 let provider = Arc::new(MockProvider::new(vec![]));
3923 let result = Orchestrator::builder(provider)
3924 .sub_agent_full(SubAgentConfig {
3925 name: "agent1".into(),
3926 description: "Test agent".into(),
3927 system_prompt: "prompt".into(),
3928 tools: vec![],
3929 context_strategy: None,
3930 summarize_threshold: None,
3931 tool_timeout: None,
3932 max_tool_output_bytes: None,
3933 max_turns: Some(0),
3934 max_tokens: None,
3935 response_schema: None,
3936 run_timeout: None,
3937 guardrails: vec![],
3938 provider: None,
3939 reasoning_effort: None,
3940 enable_reflection: None,
3941 tool_output_compression_threshold: None,
3942 max_tools_per_turn: None,
3943 tool_profile: None,
3944 max_identical_tool_calls: None,
3945 max_fuzzy_identical_tool_calls: None,
3946 max_tool_calls_per_turn: None,
3947 session_prune_config: None,
3948 enable_recursive_summarization: None,
3949 reflection_threshold: None,
3950 consolidate_on_exit: None,
3951 workspace: None,
3952 max_total_tokens: None,
3953 audit_trail: None,
3954 audit_user_id: None,
3955 audit_tenant_id: None,
3956 audit_delegation_chain: Vec::new(),
3957 })
3958 .build();
3959
3960 match result {
3961 Err(e) => assert!(
3962 e.to_string().contains("max_turns must be > 0"),
3963 "expected max_turns error, got: {e}"
3964 ),
3965 Ok(_) => panic!("expected build to fail with zero max_turns"),
3966 }
3967 }
3968
3969 #[test]
3970 fn build_rejects_sub_agent_with_zero_max_tokens() {
3971 let provider = Arc::new(MockProvider::new(vec![]));
3972 let result = Orchestrator::builder(provider)
3973 .sub_agent_full(SubAgentConfig {
3974 name: "agent1".into(),
3975 description: "Test agent".into(),
3976 system_prompt: "prompt".into(),
3977 tools: vec![],
3978 context_strategy: None,
3979 summarize_threshold: None,
3980 tool_timeout: None,
3981 max_tool_output_bytes: None,
3982 max_turns: None,
3983 max_tokens: Some(0),
3984 response_schema: None,
3985 run_timeout: None,
3986 guardrails: vec![],
3987 provider: None,
3988 reasoning_effort: None,
3989 enable_reflection: None,
3990 tool_output_compression_threshold: None,
3991 max_tools_per_turn: None,
3992 tool_profile: None,
3993 max_identical_tool_calls: None,
3994 max_fuzzy_identical_tool_calls: None,
3995 max_tool_calls_per_turn: None,
3996 session_prune_config: None,
3997 enable_recursive_summarization: None,
3998 reflection_threshold: None,
3999 consolidate_on_exit: None,
4000 workspace: None,
4001 max_total_tokens: None,
4002 audit_trail: None,
4003 audit_user_id: None,
4004 audit_tenant_id: None,
4005 audit_delegation_chain: Vec::new(),
4006 })
4007 .build();
4008
4009 match result {
4010 Err(e) => assert!(
4011 e.to_string().contains("max_tokens must be > 0"),
4012 "expected max_tokens error, got: {e}"
4013 ),
4014 Ok(_) => panic!("expected build to fail with zero max_tokens"),
4015 }
4016 }
4017
4018 #[tokio::test]
4019 async fn sub_agent_uses_override_provider() {
4020 use crate::llm::types::CompletionRequest;
4021
4022 struct IdentifiedProvider {
4024 id: String,
4025 responses: Mutex<Vec<CompletionResponse>>,
4026 }
4027
4028 impl LlmProvider for IdentifiedProvider {
4029 async fn complete(
4030 &self,
4031 _request: CompletionRequest,
4032 ) -> Result<CompletionResponse, Error> {
4033 let mut responses = self.responses.lock().expect("lock");
4034 if responses.is_empty() {
4035 return Err(Error::Agent(format!("no more responses for {}", self.id)));
4036 }
4037 Ok(responses.remove(0))
4038 }
4039 }
4040
4041 let opus_provider = Arc::new(IdentifiedProvider {
4043 id: "opus".into(),
4044 responses: Mutex::new(vec![
4045 CompletionResponse {
4047 content: vec![ContentBlock::ToolUse {
4048 id: "call-1".into(),
4049 name: "delegate_task".into(),
4050 input: json!({
4051 "tasks": [{"agent": "cheap", "task": "do cheap work"}]
4052 }),
4053 }],
4054 stop_reason: StopReason::ToolUse,
4055 usage: TokenUsage::default(),
4056 model: None,
4057 },
4058 CompletionResponse {
4060 content: vec![ContentBlock::Text {
4061 text: "Done.".into(),
4062 }],
4063 stop_reason: StopReason::EndTurn,
4064 usage: TokenUsage::default(),
4065 model: None,
4066 },
4067 ]),
4068 });
4069
4070 let haiku_provider: Arc<BoxedProvider> = Arc::new(BoxedProvider::new(IdentifiedProvider {
4071 id: "haiku".into(),
4072 responses: Mutex::new(vec![
4073 CompletionResponse {
4075 content: vec![ContentBlock::Text {
4076 text: "Cheap work done.".into(),
4077 }],
4078 stop_reason: StopReason::EndTurn,
4079 usage: TokenUsage {
4080 input_tokens: 5,
4081 output_tokens: 3,
4082 ..Default::default()
4083 },
4084 model: None,
4085 },
4086 ]),
4087 }));
4088
4089 let mut orch = Orchestrator::builder(opus_provider)
4090 .sub_agent_full(SubAgentConfig {
4091 name: "cheap".into(),
4092 description: "Cheap agent".into(),
4093 system_prompt: "You do cheap work.".into(),
4094 tools: vec![],
4095 context_strategy: None,
4096 summarize_threshold: None,
4097 tool_timeout: None,
4098 max_tool_output_bytes: None,
4099 max_turns: None,
4100 max_tokens: None,
4101 response_schema: None,
4102 run_timeout: None,
4103 guardrails: vec![],
4104 provider: Some(haiku_provider),
4105 reasoning_effort: None,
4106 enable_reflection: None,
4107 tool_output_compression_threshold: None,
4108 max_tools_per_turn: None,
4109 tool_profile: None,
4110 max_identical_tool_calls: None,
4111 max_fuzzy_identical_tool_calls: None,
4112 max_tool_calls_per_turn: None,
4113 session_prune_config: None,
4114 enable_recursive_summarization: None,
4115 reflection_threshold: None,
4116 consolidate_on_exit: None,
4117 workspace: None,
4118 max_total_tokens: None,
4119 audit_trail: None,
4120 audit_user_id: None,
4121 audit_tenant_id: None,
4122 audit_delegation_chain: Vec::new(),
4123 })
4124 .build()
4125 .unwrap();
4126
4127 let output = orch.run("do work cheaply").await.unwrap();
4128 assert_eq!(output.result, "Done.");
4129 assert_eq!(output.tokens_used.input_tokens, 5);
4131 }
4132
4133 #[tokio::test]
4134 async fn sub_agent_inherits_default_provider() {
4135 let provider = Arc::new(MockProvider::new(vec![
4137 CompletionResponse {
4139 content: vec![ContentBlock::ToolUse {
4140 id: "call-1".into(),
4141 name: "delegate_task".into(),
4142 input: json!({
4143 "tasks": [{"agent": "worker", "task": "do work"}]
4144 }),
4145 }],
4146 stop_reason: StopReason::ToolUse,
4147 usage: TokenUsage::default(),
4148 model: None,
4149 },
4150 CompletionResponse {
4152 content: vec![ContentBlock::Text {
4153 text: "Work done.".into(),
4154 }],
4155 stop_reason: StopReason::EndTurn,
4156 usage: TokenUsage::default(),
4157 model: None,
4158 },
4159 CompletionResponse {
4161 content: vec![ContentBlock::Text {
4162 text: "All done.".into(),
4163 }],
4164 stop_reason: StopReason::EndTurn,
4165 usage: TokenUsage::default(),
4166 model: None,
4167 },
4168 ]));
4169
4170 let mut orch = Orchestrator::builder(provider)
4171 .sub_agent_full(SubAgentConfig {
4172 name: "worker".into(),
4173 description: "Worker".into(),
4174 system_prompt: "Work.".into(),
4175 tools: vec![],
4176 context_strategy: None,
4177 summarize_threshold: None,
4178 tool_timeout: None,
4179 max_tool_output_bytes: None,
4180 max_turns: None,
4181 max_tokens: None,
4182 response_schema: None,
4183 run_timeout: None,
4184 guardrails: vec![],
4185 provider: None,
4186 reasoning_effort: None,
4187 enable_reflection: None,
4188 tool_output_compression_threshold: None,
4189 max_tools_per_turn: None,
4190 tool_profile: None,
4191 max_identical_tool_calls: None,
4192 max_fuzzy_identical_tool_calls: None,
4193 max_tool_calls_per_turn: None,
4194 session_prune_config: None,
4195 enable_recursive_summarization: None,
4196 reflection_threshold: None,
4197 consolidate_on_exit: None,
4198 workspace: None,
4199 max_total_tokens: None,
4200 audit_trail: None,
4201 audit_user_id: None,
4202 audit_tenant_id: None,
4203 audit_delegation_chain: Vec::new(),
4204 })
4205 .build()
4206 .unwrap();
4207
4208 let output = orch.run("do work").await.unwrap();
4209 assert_eq!(output.result, "All done.");
4210 }
4211
4212 #[test]
4215 fn form_squad_tool_definition_schema() {
4216 let tools = vec!["web_search".to_string()];
4217 let agents: Vec<(&str, &str, &[String])> = vec![
4218 ("researcher", "Research specialist", tools.as_slice()),
4219 ("analyst", "Analysis expert", &[]),
4220 ];
4221 let def = build_form_squad_tool_schema(&agents);
4222 assert_eq!(def.name, "form_squad");
4223 assert!(
4224 def.description.contains("researcher"),
4225 "description should list agents: {}",
4226 def.description
4227 );
4228 assert!(
4229 def.description.contains("analyst"),
4230 "description should list agents: {}",
4231 def.description
4232 );
4233 assert!(
4234 def.description.contains("blackboard"),
4235 "description should mention shared blackboard: {}",
4236 def.description
4237 );
4238 assert!(
4239 def.description.contains("Unlike delegate_task"),
4240 "description should contrast with delegate_task: {}",
4241 def.description
4242 );
4243 assert_eq!(
4245 def.input_schema["properties"]["tasks"]["type"], "array",
4246 "schema should have tasks array"
4247 );
4248 assert_eq!(
4249 def.input_schema["properties"]["tasks"]["items"]["properties"]["agent"]["type"],
4250 "string",
4251 "tasks items should have agent field"
4252 );
4253 assert_eq!(
4254 def.input_schema["properties"]["tasks"]["items"]["properties"]["task"]["type"],
4255 "string",
4256 "tasks items should have task field"
4257 );
4258 let required = def.input_schema["required"]
4259 .as_array()
4260 .expect("required should be array");
4261 assert!(
4262 required.contains(&json!("tasks")),
4263 "tasks should be required"
4264 );
4265 }
4266
4267 #[tokio::test]
4268 async fn form_squad_dispatches_directly() {
4269 let provider = Arc::new(MockProvider::new(vec![
4272 CompletionResponse {
4274 content: vec![ContentBlock::ToolUse {
4275 id: "call-1".into(),
4276 name: "form_squad".into(),
4277 input: json!({
4278 "tasks": [
4279 {"agent": "researcher", "task": "Research Rust"},
4280 {"agent": "analyst", "task": "Analyze findings"}
4281 ]
4282 }),
4283 }],
4284 stop_reason: StopReason::ToolUse,
4285 usage: TokenUsage {
4286 input_tokens: 50,
4287 output_tokens: 20,
4288 ..Default::default()
4289 },
4290 model: None,
4291 },
4292 CompletionResponse {
4294 content: vec![ContentBlock::Text {
4295 text: "Rust is fast and safe.".into(),
4296 }],
4297 stop_reason: StopReason::EndTurn,
4298 usage: TokenUsage {
4299 input_tokens: 10,
4300 output_tokens: 8,
4301 ..Default::default()
4302 },
4303 model: None,
4304 },
4305 CompletionResponse {
4307 content: vec![ContentBlock::Text {
4308 text: "Strengths: memory safety.".into(),
4309 }],
4310 stop_reason: StopReason::EndTurn,
4311 usage: TokenUsage {
4312 input_tokens: 12,
4313 output_tokens: 10,
4314 ..Default::default()
4315 },
4316 model: None,
4317 },
4318 CompletionResponse {
4320 content: vec![ContentBlock::Text {
4321 text: "Final: Rust is excellent.".into(),
4322 }],
4323 stop_reason: StopReason::EndTurn,
4324 usage: TokenUsage {
4325 input_tokens: 60,
4326 output_tokens: 25,
4327 ..Default::default()
4328 },
4329 model: None,
4330 },
4331 ]));
4332
4333 let mut orch = Orchestrator::builder(provider)
4334 .sub_agent("researcher", "Research specialist", "You research.")
4335 .sub_agent("analyst", "Analysis expert", "You analyze.")
4336 .sub_agent("coder", "Coding expert", "You code.")
4337 .build()
4338 .unwrap();
4339
4340 let output = orch.run("Analyze Rust deeply").await.unwrap();
4341 assert_eq!(output.result, "Final: Rust is excellent.");
4342 }
4343
4344 #[tokio::test]
4345 async fn form_squad_tokens_roll_up() {
4346 let provider = Arc::new(MockProvider::new(vec![
4347 CompletionResponse {
4349 content: vec![ContentBlock::ToolUse {
4350 id: "call-1".into(),
4351 name: "form_squad".into(),
4352 input: json!({
4353 "tasks": [
4354 {"agent": "agent_a", "task": "Task A"},
4355 {"agent": "agent_b", "task": "Task B"}
4356 ]
4357 }),
4358 }],
4359 stop_reason: StopReason::ToolUse,
4360 usage: TokenUsage {
4361 input_tokens: 50,
4362 output_tokens: 20,
4363 ..Default::default()
4364 },
4365 model: None,
4366 },
4367 CompletionResponse {
4369 content: vec![ContentBlock::Text {
4370 text: "Done A.".into(),
4371 }],
4372 stop_reason: StopReason::EndTurn,
4373 usage: TokenUsage {
4374 input_tokens: 10,
4375 output_tokens: 5,
4376 ..Default::default()
4377 },
4378 model: None,
4379 },
4380 CompletionResponse {
4382 content: vec![ContentBlock::Text {
4383 text: "Done B.".into(),
4384 }],
4385 stop_reason: StopReason::EndTurn,
4386 usage: TokenUsage {
4387 input_tokens: 12,
4388 output_tokens: 6,
4389 ..Default::default()
4390 },
4391 model: None,
4392 },
4393 CompletionResponse {
4395 content: vec![ContentBlock::Text {
4396 text: "All done.".into(),
4397 }],
4398 stop_reason: StopReason::EndTurn,
4399 usage: TokenUsage {
4400 input_tokens: 60,
4401 output_tokens: 25,
4402 ..Default::default()
4403 },
4404 model: None,
4405 },
4406 ]));
4407
4408 let mut orch = Orchestrator::builder(provider)
4409 .sub_agent("agent_a", "Agent A", "You are A.")
4410 .sub_agent("agent_b", "Agent B", "You are B.")
4411 .build()
4412 .unwrap();
4413
4414 let output = orch.run("Collaborate").await.unwrap();
4415 assert_eq!(
4419 output.tokens_used.input_tokens,
4420 50 + 60 + 10 + 12,
4421 "all token levels should roll up"
4422 );
4423 assert_eq!(
4424 output.tokens_used.output_tokens,
4425 20 + 25 + 5 + 6,
4426 "all token levels should roll up"
4427 );
4428 }
4429
4430 #[tokio::test]
4431 async fn form_squad_returns_error_for_unknown_agent() {
4432 let provider = Arc::new(MockProvider::new(vec![
4433 CompletionResponse {
4435 content: vec![ContentBlock::ToolUse {
4436 id: "call-1".into(),
4437 name: "form_squad".into(),
4438 input: json!({
4439 "tasks": [
4440 {"agent": "researcher", "task": "Do research"},
4441 {"agent": "nonexistent", "task": "Do stuff"}
4442 ]
4443 }),
4444 }],
4445 stop_reason: StopReason::ToolUse,
4446 usage: TokenUsage::default(),
4447 model: None,
4448 },
4449 CompletionResponse {
4451 content: vec![ContentBlock::Text {
4452 text: "No such agent available.".into(),
4453 }],
4454 stop_reason: StopReason::EndTurn,
4455 usage: TokenUsage::default(),
4456 model: None,
4457 },
4458 ]));
4459
4460 let mut orch = Orchestrator::builder(provider)
4461 .sub_agent("researcher", "Research", "prompt")
4462 .sub_agent("analyst", "Analysis", "prompt")
4463 .build()
4464 .unwrap();
4465
4466 let output = orch.run("delegate to unknown squad").await.unwrap();
4467 assert_eq!(output.result, "No such agent available.");
4468 }
4469
4470 #[tokio::test]
4471 async fn form_squad_requires_at_least_two_agents() {
4472 let provider = Arc::new(MockProvider::new(vec![
4473 CompletionResponse {
4475 content: vec![ContentBlock::ToolUse {
4476 id: "call-1".into(),
4477 name: "form_squad".into(),
4478 input: json!({
4479 "tasks": [
4480 {"agent": "researcher", "task": "Solo task"}
4481 ]
4482 }),
4483 }],
4484 stop_reason: StopReason::ToolUse,
4485 usage: TokenUsage::default(),
4486 model: None,
4487 },
4488 CompletionResponse {
4490 content: vec![ContentBlock::Text {
4491 text: "Using delegate_task instead.".into(),
4492 }],
4493 stop_reason: StopReason::EndTurn,
4494 usage: TokenUsage::default(),
4495 model: None,
4496 },
4497 ]));
4498
4499 let mut orch = Orchestrator::builder(provider)
4500 .sub_agent("researcher", "Research", "prompt")
4501 .sub_agent("analyst", "Analysis", "prompt")
4502 .build()
4503 .unwrap();
4504
4505 let output = orch.run("form solo squad").await.unwrap();
4506 assert_eq!(output.result, "Using delegate_task instead.");
4507 }
4508
4509 #[tokio::test]
4510 async fn form_squad_rejects_duplicate_agents() {
4511 let provider = Arc::new(MockProvider::new(vec![
4512 CompletionResponse {
4514 content: vec![ContentBlock::ToolUse {
4515 id: "call-1".into(),
4516 name: "form_squad".into(),
4517 input: json!({
4518 "tasks": [
4519 {"agent": "researcher", "task": "Task 1"},
4520 {"agent": "researcher", "task": "Task 2"}
4521 ]
4522 }),
4523 }],
4524 stop_reason: StopReason::ToolUse,
4525 usage: TokenUsage::default(),
4526 model: None,
4527 },
4528 CompletionResponse {
4530 content: vec![ContentBlock::Text {
4531 text: "Fixed duplicate issue.".into(),
4532 }],
4533 stop_reason: StopReason::EndTurn,
4534 usage: TokenUsage::default(),
4535 model: None,
4536 },
4537 ]));
4538
4539 let mut orch = Orchestrator::builder(provider)
4540 .sub_agent("researcher", "Research", "prompt")
4541 .sub_agent("analyst", "Analysis", "prompt")
4542 .build()
4543 .unwrap();
4544
4545 let output = orch.run("form squad with dupes").await.unwrap();
4546 assert_eq!(output.result, "Fixed duplicate issue.");
4547 }
4548
4549 #[tokio::test]
4550 async fn form_squad_private_blackboard() {
4551 use crate::agent::blackboard::InMemoryBlackboard;
4552
4553 let outer_bb = Arc::new(InMemoryBlackboard::new());
4554
4555 let provider = Arc::new(MockProvider::new(vec![
4556 CompletionResponse {
4558 content: vec![ContentBlock::ToolUse {
4559 id: "call-1".into(),
4560 name: "form_squad".into(),
4561 input: json!({
4562 "tasks": [
4563 {"agent": "writer_a", "task": "Write something"},
4564 {"agent": "writer_b", "task": "Write something else"}
4565 ]
4566 }),
4567 }],
4568 stop_reason: StopReason::ToolUse,
4569 usage: TokenUsage::default(),
4570 model: None,
4571 },
4572 CompletionResponse {
4574 content: vec![ContentBlock::Text {
4575 text: "Written to squad blackboard.".into(),
4576 }],
4577 stop_reason: StopReason::EndTurn,
4578 usage: TokenUsage::default(),
4579 model: None,
4580 },
4581 CompletionResponse {
4583 content: vec![ContentBlock::Text {
4584 text: "Also written.".into(),
4585 }],
4586 stop_reason: StopReason::EndTurn,
4587 usage: TokenUsage::default(),
4588 model: None,
4589 },
4590 CompletionResponse {
4592 content: vec![ContentBlock::Text {
4593 text: "Done.".into(),
4594 }],
4595 stop_reason: StopReason::EndTurn,
4596 usage: TokenUsage::default(),
4597 model: None,
4598 },
4599 ]));
4600
4601 let mut orch = Orchestrator::builder(provider)
4602 .sub_agent("writer_a", "Writer A", "You write.")
4603 .sub_agent("writer_b", "Writer B", "You write.")
4604 .blackboard(outer_bb.clone())
4605 .build()
4606 .unwrap();
4607
4608 orch.run("write to blackboard").await.unwrap();
4609
4610 let squad_key = "squad:writer_a+writer_b";
4612 let val = outer_bb.read(squad_key).await.unwrap();
4613 assert!(
4614 val.is_some(),
4615 "outer blackboard should have squad result under '{squad_key}'"
4616 );
4617
4618 let agent_key = "agent:writer_a";
4621 let val = outer_bb.read(agent_key).await.unwrap();
4622 assert!(
4623 val.is_none(),
4624 "outer blackboard should NOT have '{agent_key}' — that's on the private blackboard"
4625 );
4626 }
4627
4628 #[tokio::test]
4629 async fn form_squad_error_returns_tool_error_not_hard_error() {
4630 let failing_provider = Arc::new(BoxedProvider::new(MockProvider::new(vec![])));
4637
4638 let provider = Arc::new(MockProvider::new(vec![
4639 CompletionResponse {
4641 content: vec![ContentBlock::ToolUse {
4642 id: "call-1".into(),
4643 name: "form_squad".into(),
4644 input: json!({
4645 "tasks": [
4646 {"agent": "agent_a", "task": "Do A"},
4647 {"agent": "agent_b", "task": "Do B"}
4648 ]
4649 }),
4650 }],
4651 stop_reason: StopReason::ToolUse,
4652 usage: TokenUsage {
4653 input_tokens: 50,
4654 output_tokens: 20,
4655 ..Default::default()
4656 },
4657 model: None,
4658 },
4659 CompletionResponse {
4661 content: vec![ContentBlock::Text {
4662 text: "Done A.".into(),
4663 }],
4664 stop_reason: StopReason::EndTurn,
4665 usage: TokenUsage {
4666 input_tokens: 10,
4667 output_tokens: 5,
4668 ..Default::default()
4669 },
4670 model: None,
4671 },
4672 CompletionResponse {
4674 content: vec![ContentBlock::Text {
4675 text: "Squad failed, falling back.".into(),
4676 }],
4677 stop_reason: StopReason::EndTurn,
4678 usage: TokenUsage {
4679 input_tokens: 60,
4680 output_tokens: 25,
4681 ..Default::default()
4682 },
4683 model: None,
4684 },
4685 ]));
4686
4687 let mut orch = Orchestrator::builder(provider)
4688 .sub_agent("agent_a", "Agent A", "You are A.")
4689 .sub_agent_full(SubAgentConfig {
4690 name: "agent_b".into(),
4691 description: "Agent B".into(),
4692 system_prompt: "You are B.".into(),
4693 tools: vec![],
4694 context_strategy: None,
4695 summarize_threshold: None,
4696 tool_timeout: None,
4697 max_tool_output_bytes: None,
4698 max_turns: None,
4699 max_tokens: None,
4700 response_schema: None,
4701 run_timeout: None,
4702 guardrails: vec![],
4703 provider: Some(failing_provider),
4704 reasoning_effort: None,
4705 enable_reflection: None,
4706 tool_output_compression_threshold: None,
4707 max_tools_per_turn: None,
4708 tool_profile: None,
4709 max_identical_tool_calls: None,
4710 max_fuzzy_identical_tool_calls: None,
4711 max_tool_calls_per_turn: None,
4712 session_prune_config: None,
4713 enable_recursive_summarization: None,
4714 reflection_threshold: None,
4715 consolidate_on_exit: None,
4716 workspace: None,
4717 max_total_tokens: None,
4718 audit_trail: None,
4719 audit_user_id: None,
4720 audit_tenant_id: None,
4721 audit_delegation_chain: Vec::new(),
4722 })
4723 .build()
4724 .unwrap();
4725
4726 let output = orch.run("complex task").await.unwrap();
4727 assert_eq!(output.result, "Squad failed, falling back.");
4728 assert!(
4730 output.tokens_used.input_tokens > 50 + 60,
4731 "should include partial squad tokens: {}",
4732 output.tokens_used.input_tokens
4733 );
4734 }
4735
4736 #[tokio::test]
4737 async fn orchestrator_registers_both_tools() {
4738 use crate::llm::types::CompletionRequest;
4739
4740 struct ToolCapturingProvider {
4741 responses: Mutex<Vec<CompletionResponse>>,
4742 tool_names_seen: Mutex<Vec<Vec<String>>>,
4743 }
4744
4745 impl LlmProvider for ToolCapturingProvider {
4746 async fn complete(
4747 &self,
4748 request: CompletionRequest,
4749 ) -> Result<CompletionResponse, Error> {
4750 let names: Vec<String> = request.tools.iter().map(|t| t.name.clone()).collect();
4751 self.tool_names_seen.lock().expect("lock").push(names);
4752 let mut responses = self.responses.lock().expect("lock");
4753 if responses.is_empty() {
4754 return Err(Error::Agent("no more responses".into()));
4755 }
4756 Ok(responses.remove(0))
4757 }
4758 }
4759
4760 let provider = Arc::new(ToolCapturingProvider {
4761 responses: Mutex::new(vec![CompletionResponse {
4762 content: vec![ContentBlock::Text {
4763 text: "Direct answer.".into(),
4764 }],
4765 stop_reason: StopReason::EndTurn,
4766 usage: TokenUsage::default(),
4767 model: None,
4768 }]),
4769 tool_names_seen: Mutex::new(vec![]),
4770 });
4771
4772 let mut orch = Orchestrator::builder(provider.clone())
4774 .sub_agent("researcher", "Research", "prompt")
4775 .sub_agent("analyst", "Analysis", "prompt")
4776 .build()
4777 .unwrap();
4778
4779 orch.run("test").await.unwrap();
4780
4781 let tool_names = provider.tool_names_seen.lock().unwrap();
4782 assert!(
4783 tool_names[0].contains(&"delegate_task".to_string()),
4784 "should have delegate_task: {:?}",
4785 tool_names[0]
4786 );
4787 assert!(
4788 tool_names[0].contains(&"form_squad".to_string()),
4789 "should have form_squad: {:?}",
4790 tool_names[0]
4791 );
4792 }
4793
4794 #[test]
4795 fn orchestrator_single_agent_no_squads() {
4796 let provider = Arc::new(MockProvider::new(vec![]));
4797
4798 let result = Orchestrator::builder(provider)
4800 .sub_agent("researcher", "Research", "prompt")
4801 .build();
4802
4803 assert!(result.is_ok());
4804 }
4805
4806 #[tokio::test]
4807 async fn orchestrator_squads_disabled_explicitly() {
4808 use crate::llm::types::CompletionRequest;
4809
4810 struct ToolCapturingProvider {
4811 responses: Mutex<Vec<CompletionResponse>>,
4812 tool_names_seen: Mutex<Vec<Vec<String>>>,
4813 }
4814
4815 impl LlmProvider for ToolCapturingProvider {
4816 async fn complete(
4817 &self,
4818 request: CompletionRequest,
4819 ) -> Result<CompletionResponse, Error> {
4820 let names: Vec<String> = request.tools.iter().map(|t| t.name.clone()).collect();
4821 self.tool_names_seen.lock().expect("lock").push(names);
4822 let mut responses = self.responses.lock().expect("lock");
4823 if responses.is_empty() {
4824 return Err(Error::Agent("no more responses".into()));
4825 }
4826 Ok(responses.remove(0))
4827 }
4828 }
4829
4830 let provider = Arc::new(ToolCapturingProvider {
4831 responses: Mutex::new(vec![CompletionResponse {
4832 content: vec![ContentBlock::Text {
4833 text: "Direct answer.".into(),
4834 }],
4835 stop_reason: StopReason::EndTurn,
4836 usage: TokenUsage::default(),
4837 model: None,
4838 }]),
4839 tool_names_seen: Mutex::new(vec![]),
4840 });
4841
4842 let mut orch = Orchestrator::builder(provider.clone())
4844 .sub_agent("researcher", "Research", "prompt")
4845 .sub_agent("analyst", "Analysis", "prompt")
4846 .enable_squads(false)
4847 .build()
4848 .unwrap();
4849
4850 orch.run("test").await.unwrap();
4851
4852 let tool_names = provider.tool_names_seen.lock().unwrap();
4853 assert!(
4854 tool_names[0].contains(&"delegate_task".to_string()),
4855 "should have delegate_task: {:?}",
4856 tool_names[0]
4857 );
4858 assert!(
4859 !tool_names[0].contains(&"form_squad".to_string()),
4860 "should NOT have form_squad when disabled: {:?}",
4861 tool_names[0]
4862 );
4863 }
4864
4865 #[test]
4866 fn system_prompt_mentions_both_tools_when_squads_enabled() {
4867 let tools = vec!["web_search".to_string()];
4868 let agents: Vec<(&str, &str, &[String])> = vec![
4869 ("researcher", "Research specialist", tools.as_slice()),
4870 ("analyst", "Analysis expert", &[]),
4871 ];
4872
4873 let prompt = build_system_prompt(&agents, true, DispatchMode::Parallel);
4874 assert!(
4875 prompt.contains("delegate_task"),
4876 "prompt should mention delegate_task: {prompt}"
4877 );
4878 assert!(
4879 prompt.contains("form_squad"),
4880 "prompt should mention form_squad: {prompt}"
4881 );
4882 assert!(
4883 prompt.contains("two delegation tools"),
4884 "prompt should explain both tools: {prompt}"
4885 );
4886 assert!(
4888 prompt.contains("isolation"),
4889 "prompt should mention isolation for delegate_task: {prompt}"
4890 );
4891 assert!(
4892 prompt.contains("blackboard"),
4893 "prompt should mention shared blackboard for form_squad: {prompt}"
4894 );
4895 }
4896
4897 #[test]
4898 fn system_prompt_only_delegate_when_squads_disabled() {
4899 let agents: Vec<(&str, &str, &[String])> = vec![("researcher", "Research specialist", &[])];
4900
4901 let prompt = build_system_prompt(&agents, false, DispatchMode::Parallel);
4902 assert!(
4903 prompt.contains("delegate_task"),
4904 "prompt should mention delegate_task: {prompt}"
4905 );
4906 assert!(
4907 !prompt.contains("form_squad"),
4908 "prompt should NOT mention form_squad: {prompt}"
4909 );
4910 assert!(
4912 prompt.contains("Decision Process"),
4913 "prompt should contain Decision Process even without squads: {prompt}"
4914 );
4915 assert!(
4916 prompt.contains("Effort Scaling"),
4917 "prompt should contain Effort Scaling even without squads: {prompt}"
4918 );
4919 }
4920
4921 #[tokio::test]
4922 async fn delegate_forwards_on_event_to_sub_agents() {
4923 let provider = Arc::new(MockProvider::new(vec![
4925 CompletionResponse {
4927 content: vec![ContentBlock::ToolUse {
4928 id: "call-1".into(),
4929 name: "delegate_task".into(),
4930 input: json!({
4931 "tasks": [
4932 {"agent": "worker", "task": "do work"}
4933 ]
4934 }),
4935 }],
4936 stop_reason: StopReason::ToolUse,
4937 usage: TokenUsage::default(),
4938 model: None,
4939 },
4940 CompletionResponse {
4942 content: vec![ContentBlock::Text {
4943 text: "done".into(),
4944 }],
4945 stop_reason: StopReason::EndTurn,
4946 usage: TokenUsage::default(),
4947 model: None,
4948 },
4949 CompletionResponse {
4951 content: vec![ContentBlock::Text {
4952 text: "All done.".into(),
4953 }],
4954 stop_reason: StopReason::EndTurn,
4955 usage: TokenUsage::default(),
4956 model: None,
4957 },
4958 ]));
4959
4960 let events = Arc::new(Mutex::new(Vec::<AgentEvent>::new()));
4961 let events_clone = events.clone();
4962 let on_event: Arc<OnEvent> = Arc::new(move |event: AgentEvent| {
4963 events_clone.lock().expect("test lock").push(event);
4964 });
4965
4966 let mut orch = Orchestrator::builder(provider)
4967 .sub_agent("worker", "Worker agent", "You do work.")
4968 .on_event(on_event)
4969 .build()
4970 .unwrap();
4971
4972 let _output = orch.run("delegate some work").await.unwrap();
4973
4974 let events = events.lock().expect("test lock");
4975
4976 let orchestrator_events: Vec<_> = events
4978 .iter()
4979 .filter(|e| match e {
4980 AgentEvent::RunStarted { agent, .. }
4981 | AgentEvent::TurnStarted { agent, .. }
4982 | AgentEvent::LlmResponse { agent, .. }
4983 | AgentEvent::RunCompleted { agent, .. } => agent == "orchestrator",
4984 _ => false,
4985 })
4986 .collect();
4987 let worker_events: Vec<_> = events
4988 .iter()
4989 .filter(|e| match e {
4990 AgentEvent::RunStarted { agent, .. }
4991 | AgentEvent::TurnStarted { agent, .. }
4992 | AgentEvent::LlmResponse { agent, .. }
4993 | AgentEvent::RunCompleted { agent, .. } => agent == "worker",
4994 _ => false,
4995 })
4996 .collect();
4997
4998 assert!(
4999 !orchestrator_events.is_empty(),
5000 "should have orchestrator events"
5001 );
5002 assert!(
5003 !worker_events.is_empty(),
5004 "should have sub-agent worker events (forwarded via on_event)"
5005 );
5006
5007 let worker_run_started = events
5009 .iter()
5010 .any(|e| matches!(e, AgentEvent::RunStarted { agent, .. } if agent == "worker"));
5011 assert!(
5012 worker_run_started,
5013 "sub-agent should emit RunStarted via forwarded on_event"
5014 );
5015 }
5016
5017 #[tokio::test]
5036 async fn full_audit_trail_end_to_end() {
5037 let long_output = "x".repeat(70_000); let provider = Arc::new(MockProvider::new(vec![
5041 CompletionResponse {
5043 content: vec![
5044 ContentBlock::Text {
5045 text: "I'll delegate to the researcher and coder.".into(),
5046 },
5047 ContentBlock::ToolUse {
5048 id: "orch-call-1".into(),
5049 name: "delegate_task".into(),
5050 input: json!({
5051 "tasks": [
5052 {"agent": "researcher", "task": "Search for Rust concurrency patterns"},
5053 {"agent": "coder", "task": "Read the main.rs file"}
5054 ]
5055 }),
5056 },
5057 ],
5058 stop_reason: StopReason::ToolUse,
5059 usage: TokenUsage {
5060 input_tokens: 100,
5061 output_tokens: 40,
5062 ..Default::default()
5063 },
5064 model: None,
5065 },
5066 CompletionResponse {
5068 content: vec![
5069 ContentBlock::Text {
5070 text: "Let me search for Rust concurrency info.".into(),
5071 },
5072 ContentBlock::ToolUse {
5073 id: "res-call-1".into(),
5074 name: "web_search".into(),
5075 input: json!({"query": "rust async concurrency"}),
5076 },
5077 ],
5078 stop_reason: StopReason::ToolUse,
5079 usage: TokenUsage {
5080 input_tokens: 20,
5081 output_tokens: 10,
5082 ..Default::default()
5083 },
5084 model: None,
5085 },
5086 CompletionResponse {
5088 content: vec![ContentBlock::Text {
5089 text: "Rust uses async/await with tokio for concurrency.".into(),
5090 }],
5091 stop_reason: StopReason::EndTurn,
5092 usage: TokenUsage {
5093 input_tokens: 30,
5094 output_tokens: 15,
5095 ..Default::default()
5096 },
5097 model: None,
5098 },
5099 CompletionResponse {
5101 content: vec![
5102 ContentBlock::Text {
5103 text: "I'll read the main.rs file.".into(),
5104 },
5105 ContentBlock::ToolUse {
5106 id: "cod-call-1".into(),
5107 name: "read_file".into(),
5108 input: json!({"path": "/src/main.rs"}),
5109 },
5110 ],
5111 stop_reason: StopReason::ToolUse,
5112 usage: TokenUsage {
5113 input_tokens: 15,
5114 output_tokens: 8,
5115 ..Default::default()
5116 },
5117 model: None,
5118 },
5119 CompletionResponse {
5121 content: vec![ContentBlock::Text {
5122 text: "The main.rs contains the entry point.".into(),
5123 }],
5124 stop_reason: StopReason::EndTurn,
5125 usage: TokenUsage {
5126 input_tokens: 25,
5127 output_tokens: 12,
5128 ..Default::default()
5129 },
5130 model: None,
5131 },
5132 CompletionResponse {
5134 content: vec![ContentBlock::Text {
5135 text: "Combined analysis: Rust async is great for concurrency.".into(),
5136 }],
5137 stop_reason: StopReason::EndTurn,
5138 usage: TokenUsage {
5139 input_tokens: 200,
5140 output_tokens: 50,
5141 ..Default::default()
5142 },
5143 model: None,
5144 },
5145 ]));
5146
5147 let events = Arc::new(Mutex::new(Vec::<AgentEvent>::new()));
5148 let events_clone = events.clone();
5149 let on_event: Arc<OnEvent> = Arc::new(move |event: AgentEvent| {
5150 events_clone.lock().expect("test lock").push(event);
5151 });
5152
5153 let long_output_clone = long_output.clone();
5156 let shared_tools: Vec<Arc<dyn Tool>> = vec![
5157 Arc::new(MockTool::new("web_search", &long_output_clone)),
5158 Arc::new(MockTool::new(
5159 "read_file",
5160 "fn main() { println!(\"hello\"); }",
5161 )),
5162 ];
5163 let mut orch = Orchestrator::builder(provider)
5164 .sub_agent_with_tools(
5165 "researcher",
5166 "Research specialist",
5167 "You research topics.",
5168 shared_tools.clone(),
5169 )
5170 .sub_agent_with_tools(
5171 "coder",
5172 "Code expert",
5173 "You read and analyze code.",
5174 shared_tools.clone(),
5175 )
5176 .on_event(on_event)
5177 .build()
5178 .unwrap();
5179
5180 let output = orch
5181 .run("Analyze Rust concurrency and the main.rs file")
5182 .await
5183 .unwrap();
5184
5185 assert_eq!(
5187 output.result,
5188 "Combined analysis: Rust async is great for concurrency."
5189 );
5190 assert_eq!(output.tool_calls_made, 1); let events = events.lock().expect("test lock");
5194
5195 fn agent_of(e: &AgentEvent) -> &str {
5197 match e {
5198 AgentEvent::RunStarted { agent, .. }
5199 | AgentEvent::TurnStarted { agent, .. }
5200 | AgentEvent::LlmResponse { agent, .. }
5201 | AgentEvent::ToolCallStarted { agent, .. }
5202 | AgentEvent::ToolCallCompleted { agent, .. }
5203 | AgentEvent::RunCompleted { agent, .. }
5204 | AgentEvent::RunFailed { agent, .. }
5205 | AgentEvent::SubAgentsDispatched { agent, .. }
5206 | AgentEvent::SubAgentCompleted { agent, .. }
5207 | AgentEvent::ApprovalRequested { agent, .. }
5208 | AgentEvent::ApprovalDecision { agent, .. }
5209 | AgentEvent::ContextSummarized { agent, .. }
5210 | AgentEvent::GuardrailDenied { agent, .. }
5211 | AgentEvent::GuardrailWarned { agent, .. }
5212 | AgentEvent::RetryAttempt { agent, .. }
5213 | AgentEvent::DoomLoopDetected { agent, .. }
5214 | AgentEvent::FuzzyDoomLoopDetected { agent, .. }
5215 | AgentEvent::AutoCompactionTriggered { agent, .. }
5216 | AgentEvent::SessionPruned { agent, .. }
5217 | AgentEvent::ModelEscalated { agent, .. }
5218 | AgentEvent::BudgetExceeded { agent, .. }
5219 | AgentEvent::AgentSpawned { agent, .. }
5220 | AgentEvent::KillSwitchActivated { agent, .. }
5221 | AgentEvent::ToolNameRepaired { agent, .. } => agent,
5222 AgentEvent::SensorEventProcessed { sensor_name, .. } => sensor_name,
5223 AgentEvent::StoryUpdated { story_id, .. } => story_id,
5224 AgentEvent::TaskRouted { decision, .. } => decision,
5225 AgentEvent::WorkflowNodeStarted { node, .. }
5226 | AgentEvent::WorkflowNodeCompleted { node, .. }
5227 | AgentEvent::WorkflowNodeFailed { node, .. } => node,
5228 }
5229 }
5230
5231 let event_summary: Vec<String> = events
5233 .iter()
5234 .enumerate()
5235 .map(|(i, e)| format!("{i}: [{:>12}] {:?}", agent_of(e), std::mem::discriminant(e)))
5236 .collect();
5237
5238 let agents_seen: std::collections::HashSet<&str> = events.iter().map(agent_of).collect();
5240 assert!(
5241 agents_seen.contains("orchestrator"),
5242 "missing orchestrator events.\nEvent stream:\n{}",
5243 event_summary.join("\n")
5244 );
5245 assert!(
5246 agents_seen.contains("researcher"),
5247 "missing researcher events (should be forwarded).\nEvent stream:\n{}",
5248 event_summary.join("\n")
5249 );
5250 assert!(
5251 agents_seen.contains("coder"),
5252 "missing coder events (should be forwarded).\nEvent stream:\n{}",
5253 event_summary.join("\n")
5254 );
5255
5256 let orch_events: Vec<&AgentEvent> = events
5258 .iter()
5259 .filter(|e| agent_of(e) == "orchestrator")
5260 .collect();
5261
5262 assert!(
5264 matches!(orch_events[0], AgentEvent::RunStarted { task, .. } if task.contains("Analyze Rust")),
5265 "first orch event should be RunStarted, got: {:?}",
5266 orch_events[0]
5267 );
5268 assert!(
5270 matches!(orch_events.last().unwrap(), AgentEvent::RunCompleted { .. }),
5271 "last orch event should be RunCompleted, got: {:?}",
5272 orch_events.last().unwrap()
5273 );
5274
5275 let llm_responses: Vec<&AgentEvent> = events
5277 .iter()
5278 .filter(|e| matches!(e, AgentEvent::LlmResponse { .. }))
5279 .collect();
5280 assert!(
5281 llm_responses.len() >= 3,
5282 "expected >= 3 LlmResponse events (1 orch + at least 1 per sub-agent), got {}.\nEvents:\n{}",
5283 llm_responses.len(),
5284 event_summary.join("\n")
5285 );
5286
5287 for llm_event in &llm_responses {
5288 match llm_event {
5289 AgentEvent::LlmResponse {
5290 agent, text, model, ..
5291 } => {
5292 assert_eq!(
5294 model.as_deref(),
5295 Some("mock-model-v1"),
5296 "LlmResponse for '{agent}' should have model name"
5297 );
5298 assert!(
5300 !text.is_empty(),
5301 "LlmResponse for '{agent}' should have non-empty text"
5302 );
5303 }
5304 _ => unreachable!(),
5305 }
5306 }
5307
5308 let tool_started: Vec<&AgentEvent> = events
5310 .iter()
5311 .filter(|e| matches!(e, AgentEvent::ToolCallStarted { .. }))
5312 .collect();
5313 assert!(
5315 tool_started.len() >= 3,
5316 "expected >= 3 ToolCallStarted events, got {}.\nEvents:\n{}",
5317 tool_started.len(),
5318 event_summary.join("\n")
5319 );
5320
5321 let web_search_started = tool_started.iter().find(|e| {
5324 matches!(e, AgentEvent::ToolCallStarted { tool_name, .. } if tool_name == "web_search")
5325 });
5326 assert!(
5327 web_search_started.is_some(),
5328 "should have a web_search ToolCallStarted"
5329 );
5330 match web_search_started.unwrap() {
5331 AgentEvent::ToolCallStarted { input, .. } => {
5332 assert!(
5333 input.contains("rust async concurrency"),
5334 "web_search input should contain query, got: {input}"
5335 );
5336 }
5337 _ => unreachable!(),
5338 }
5339
5340 let read_file_started = tool_started.iter().find(|e| {
5342 matches!(e, AgentEvent::ToolCallStarted { tool_name, .. } if tool_name == "read_file")
5343 });
5344 assert!(
5345 read_file_started.is_some(),
5346 "should have a read_file ToolCallStarted"
5347 );
5348 match read_file_started.unwrap() {
5349 AgentEvent::ToolCallStarted { input, .. } => {
5350 assert!(
5351 input.contains("/src/main.rs"),
5352 "read_file input should contain path, got: {input}"
5353 );
5354 }
5355 _ => unreachable!(),
5356 }
5357
5358 let delegate_started = tool_started.iter().find(|e| {
5360 matches!(e, AgentEvent::ToolCallStarted { tool_name, .. } if tool_name == "delegate_task")
5361 });
5362 assert!(
5363 delegate_started.is_some(),
5364 "should have a delegate_task ToolCallStarted"
5365 );
5366 match delegate_started.unwrap() {
5367 AgentEvent::ToolCallStarted { agent, input, .. } => {
5368 assert_eq!(agent, "orchestrator");
5369 assert!(
5370 input.contains("researcher"),
5371 "delegate_task input should contain agent names, got: {input}"
5372 );
5373 }
5374 _ => unreachable!(),
5375 }
5376
5377 let tool_completed: Vec<&AgentEvent> = events
5379 .iter()
5380 .filter(|e| matches!(e, AgentEvent::ToolCallCompleted { .. }))
5381 .collect();
5382 assert!(
5383 tool_completed.len() >= 3,
5384 "expected >= 3 ToolCallCompleted events, got {}",
5385 tool_completed.len()
5386 );
5387
5388 let web_search_completed = tool_completed.iter().find(|e| {
5390 matches!(e, AgentEvent::ToolCallCompleted { tool_name, .. } if tool_name == "web_search")
5391 });
5392 assert!(
5393 web_search_completed.is_some(),
5394 "should have a web_search ToolCallCompleted"
5395 );
5396 match web_search_completed.unwrap() {
5397 AgentEvent::ToolCallCompleted {
5398 output, is_error, ..
5399 } => {
5400 assert!(!is_error);
5401 assert!(
5403 output.contains("[truncated:"),
5404 "web_search output (70000 bytes) should be truncated in event, got {} bytes: {}",
5405 output.len(),
5406 &output[..output.len().min(100)]
5407 );
5408 }
5409 _ => unreachable!(),
5410 }
5411
5412 let read_file_completed = tool_completed.iter().find(|e| {
5414 matches!(e, AgentEvent::ToolCallCompleted { tool_name, .. } if tool_name == "read_file")
5415 });
5416 assert!(
5417 read_file_completed.is_some(),
5418 "should have a read_file ToolCallCompleted"
5419 );
5420 match read_file_completed.unwrap() {
5421 AgentEvent::ToolCallCompleted {
5422 output, is_error, ..
5423 } => {
5424 assert!(!is_error);
5425 assert!(
5426 output.contains("fn main()"),
5427 "read_file output should contain file content, got: {output}"
5428 );
5429 assert!(
5430 !output.contains("[truncated:"),
5431 "read_file output should NOT be truncated"
5432 );
5433 }
5434 _ => unreachable!(),
5435 }
5436
5437 let dispatched: Vec<&AgentEvent> = events
5439 .iter()
5440 .filter(|e| matches!(e, AgentEvent::SubAgentsDispatched { .. }))
5441 .collect();
5442 assert_eq!(dispatched.len(), 1, "expected 1 SubAgentsDispatched event");
5443 match dispatched[0] {
5444 AgentEvent::SubAgentsDispatched { agents, .. } => {
5445 assert!(
5446 agents.contains(&"researcher".to_string()),
5447 "dispatched agents should include researcher"
5448 );
5449 assert!(
5450 agents.contains(&"coder".to_string()),
5451 "dispatched agents should include coder"
5452 );
5453 }
5454 _ => unreachable!(),
5455 }
5456
5457 let completed: Vec<&AgentEvent> = events
5458 .iter()
5459 .filter(|e| matches!(e, AgentEvent::SubAgentCompleted { .. }))
5460 .collect();
5461 assert_eq!(
5462 completed.len(),
5463 2,
5464 "expected 2 SubAgentCompleted events (one per sub-agent)"
5465 );
5466 for c in &completed {
5467 match c {
5468 AgentEvent::SubAgentCompleted { success, agent, .. } => {
5469 assert!(success, "sub-agent '{agent}' should succeed");
5470 }
5471 _ => unreachable!(),
5472 }
5473 }
5474
5475 for agent_name in &["orchestrator", "researcher", "coder"] {
5477 let agent_events: Vec<&AgentEvent> = events
5478 .iter()
5479 .filter(|e| agent_of(e) == *agent_name)
5480 .collect();
5481 if !agent_events.is_empty() {
5482 assert!(
5483 matches!(agent_events[0], AgentEvent::RunStarted { .. }),
5484 "first event for '{agent_name}' should be RunStarted, got: {:?}",
5485 agent_events[0]
5486 );
5487 }
5488 }
5489
5490 assert_eq!(
5496 output.tokens_used.input_tokens,
5497 100 + 200 + 20 + 30 + 15 + 25,
5498 "total input tokens should include orchestrator + sub-agents"
5499 );
5500 assert_eq!(
5501 output.tokens_used.output_tokens,
5502 40 + 50 + 10 + 15 + 8 + 12,
5503 "total output tokens should include orchestrator + sub-agents"
5504 );
5505
5506 let sub_agent_llm = llm_responses.iter().find(|e| {
5509 matches!(e, AgentEvent::LlmResponse { text, .. }
5510 if text.contains("async/await"))
5511 });
5512 assert!(
5513 sub_agent_llm.is_some(),
5514 "should have a sub-agent LlmResponse with text about async/await"
5515 );
5516
5517 assert!(
5527 events.len() >= 20,
5528 "expected at least 20 events for full audit trail, got {}.\nEvents:\n{}",
5529 events.len(),
5530 event_summary.join("\n")
5531 );
5532 }
5533
5534 #[tokio::test]
5535 async fn sub_agent_run_timeout_fires_when_configured() {
5536 let provider = Arc::new(MockProvider::new(vec![
5540 CompletionResponse {
5542 content: vec![ContentBlock::ToolUse {
5543 id: "call-1".into(),
5544 name: "delegate_task".into(),
5545 input: json!({
5546 "tasks": [{"agent": "slow-agent", "task": "do something"}]
5547 }),
5548 }],
5549 stop_reason: StopReason::ToolUse,
5550 usage: TokenUsage::default(),
5551 model: None,
5552 },
5553 ]));
5557
5558 struct SlowProvider;
5560 impl LlmProvider for SlowProvider {
5561 async fn complete(
5562 &self,
5563 _request: CompletionRequest,
5564 ) -> Result<CompletionResponse, Error> {
5565 tokio::time::sleep(Duration::from_secs(3600)).await;
5566 unreachable!()
5567 }
5568 }
5569 let slow_provider = Arc::new(BoxedProvider::new(SlowProvider));
5570
5571 let mut orch = Orchestrator::builder(provider)
5572 .sub_agent_full(SubAgentConfig {
5573 name: "slow-agent".into(),
5574 description: "A slow agent".into(),
5575 system_prompt: "sys".into(),
5576 tools: vec![],
5577 context_strategy: None,
5578 summarize_threshold: None,
5579 tool_timeout: None,
5580 max_tool_output_bytes: None,
5581 max_turns: None,
5582 max_tokens: None,
5583 response_schema: None,
5584 run_timeout: Some(Duration::from_millis(100)),
5585 guardrails: vec![],
5586 provider: Some(slow_provider),
5587 reasoning_effort: None,
5588 enable_reflection: None,
5589 tool_output_compression_threshold: None,
5590 max_tools_per_turn: None,
5591 tool_profile: None,
5592 max_identical_tool_calls: None,
5593 max_fuzzy_identical_tool_calls: None,
5594 max_tool_calls_per_turn: None,
5595 session_prune_config: None,
5596 enable_recursive_summarization: None,
5597 reflection_threshold: None,
5598 consolidate_on_exit: None,
5599 workspace: None,
5600 max_total_tokens: None,
5601 audit_trail: None,
5602 audit_user_id: None,
5603 audit_tenant_id: None,
5604 audit_delegation_chain: Vec::new(),
5605 })
5606 .build()
5607 .unwrap();
5608
5609 let result = orch.run("go").await;
5613 match result {
5616 Ok(output) => {
5617 assert!(
5619 output.result.contains("timeout") || output.result.contains("Timeout"),
5620 "expected timeout in result, got: {}",
5621 output.result
5622 );
5623 }
5624 Err(e) => {
5625 let msg = e.to_string();
5628 assert!(
5629 msg.contains("no more mock responses")
5630 || msg.contains("timeout")
5631 || msg.contains("Timeout"),
5632 "expected timeout-related error, got: {msg}"
5633 );
5634 }
5635 }
5636 }
5637
5638 #[tokio::test]
5647 async fn form_squad_complex_with_tools_events_and_failure() {
5648 use crate::agent::blackboard::InMemoryBlackboard;
5649
5650 let outer_bb = Arc::new(InMemoryBlackboard::new());
5651
5652 let failing_provider = Arc::new(BoxedProvider::new(MockProvider::new(vec![])));
5654
5655 let provider = Arc::new(MockProvider::new(vec![
5658 CompletionResponse {
5660 content: vec![ContentBlock::ToolUse {
5661 id: "orch-call-1".into(),
5662 name: "form_squad".into(),
5663 input: json!({
5664 "tasks": [
5665 {"agent": "planner", "task": "Create a plan for the analysis"},
5666 {"agent": "worker", "task": "Compute the metrics"},
5667 {"agent": "reviewer", "task": "Review all findings"}
5668 ]
5669 }),
5670 }],
5671 stop_reason: StopReason::ToolUse,
5672 usage: TokenUsage {
5673 input_tokens: 100,
5674 output_tokens: 40,
5675 cache_creation_input_tokens: 5,
5676 cache_read_input_tokens: 3,
5677 reasoning_tokens: 0,
5678 },
5679 model: None,
5680 },
5681 CompletionResponse {
5683 content: vec![ContentBlock::Text {
5684 text: "Plan: Step 1 gather data, Step 2 compute, Step 3 review.".into(),
5685 }],
5686 stop_reason: StopReason::EndTurn,
5687 usage: TokenUsage {
5688 input_tokens: 20,
5689 output_tokens: 15,
5690 reasoning_tokens: 8,
5691 ..Default::default()
5692 },
5693 model: None,
5694 },
5695 CompletionResponse {
5697 content: vec![
5698 ContentBlock::Text {
5699 text: "I'll compute the metrics now.".into(),
5700 },
5701 ContentBlock::ToolUse {
5702 id: "worker-call-1".into(),
5703 name: "compute".into(),
5704 input: json!({"expression": "42 * 17"}),
5705 },
5706 ],
5707 stop_reason: StopReason::ToolUse,
5708 usage: TokenUsage {
5709 input_tokens: 25,
5710 output_tokens: 12,
5711 ..Default::default()
5712 },
5713 model: None,
5714 },
5715 CompletionResponse {
5717 content: vec![ContentBlock::Text {
5718 text: "Computation result: 714. Analysis complete.".into(),
5719 }],
5720 stop_reason: StopReason::EndTurn,
5721 usage: TokenUsage {
5722 input_tokens: 35,
5723 output_tokens: 18,
5724 ..Default::default()
5725 },
5726 model: None,
5727 },
5728 CompletionResponse {
5731 content: vec![ContentBlock::Text {
5732 text: "Squad partial success: plan and computation done, review failed.".into(),
5733 }],
5734 stop_reason: StopReason::EndTurn,
5735 usage: TokenUsage {
5736 input_tokens: 200,
5737 output_tokens: 60,
5738 ..Default::default()
5739 },
5740 model: None,
5741 },
5742 ]));
5743
5744 let events = Arc::new(Mutex::new(Vec::<AgentEvent>::new()));
5746 let events_clone = events.clone();
5747 let on_event: Arc<OnEvent> = Arc::new(move |event: AgentEvent| {
5748 events_clone.lock().expect("test lock").push(event);
5749 });
5750
5751 let compute_tool: Arc<dyn Tool> = Arc::new(MockTool::new("compute", "714"));
5754
5755 let mut orch = Orchestrator::builder(provider)
5756 .sub_agent_with_tools(
5757 "planner",
5758 "Planning specialist",
5759 "You create plans.",
5760 vec![compute_tool.clone()],
5761 )
5762 .sub_agent_with_tools(
5763 "worker",
5764 "Computation worker",
5765 "You compute metrics.",
5766 vec![compute_tool.clone()],
5767 )
5768 .sub_agent_full(SubAgentConfig {
5769 name: "reviewer".into(),
5770 description: "Review specialist".into(),
5771 system_prompt: "You review findings.".into(),
5772 tools: vec![],
5773 context_strategy: None,
5774 summarize_threshold: None,
5775 tool_timeout: None,
5776 max_tool_output_bytes: None,
5777 max_turns: None,
5778 max_tokens: None,
5779 response_schema: None,
5780 run_timeout: None,
5781 guardrails: vec![],
5782 provider: Some(failing_provider),
5783 reasoning_effort: None,
5784 enable_reflection: None,
5785 tool_output_compression_threshold: None,
5786 max_tools_per_turn: None,
5787 tool_profile: None,
5788 max_identical_tool_calls: None,
5789 max_fuzzy_identical_tool_calls: None,
5790 max_tool_calls_per_turn: None,
5791 session_prune_config: None,
5792 enable_recursive_summarization: None,
5793 reflection_threshold: None,
5794 consolidate_on_exit: None,
5795 workspace: None,
5796 max_total_tokens: None,
5797 audit_trail: None,
5798 audit_user_id: None,
5799 audit_tenant_id: None,
5800 audit_delegation_chain: Vec::new(),
5801 })
5802 .blackboard(outer_bb.clone())
5803 .on_event(on_event)
5804 .build()
5805 .unwrap();
5806
5807 let output = orch.run("Analyze the system performance").await.unwrap();
5808
5809 assert_eq!(
5811 output.result,
5812 "Squad partial success: plan and computation done, review failed."
5813 );
5814
5815 let expected_input = 100 + 200 + 20 + 25 + 35;
5821 let expected_output = 40 + 60 + 15 + 12 + 18;
5822 assert_eq!(
5823 output.tokens_used.input_tokens, expected_input,
5824 "input tokens should sum orchestrator + planner + worker (reviewer failed)"
5825 );
5826 assert_eq!(
5827 output.tokens_used.output_tokens, expected_output,
5828 "output tokens should sum orchestrator + planner + worker"
5829 );
5830 assert_eq!(
5831 output.tokens_used.reasoning_tokens, 8,
5832 "reasoning tokens should come from planner"
5833 );
5834 assert_eq!(
5835 output.tokens_used.cache_creation_input_tokens, 5,
5836 "cache creation tokens from orchestrator"
5837 );
5838 assert_eq!(
5839 output.tokens_used.cache_read_input_tokens, 3,
5840 "cache read tokens from orchestrator"
5841 );
5842
5843 let squad_key = "squad:planner+worker+reviewer";
5846 let squad_val = outer_bb.read(squad_key).await.unwrap();
5847 assert!(
5848 squad_val.is_some(),
5849 "outer blackboard should have squad result under '{squad_key}'"
5850 );
5851 let squad_text = squad_val.unwrap().to_string();
5852 assert!(
5854 squad_text.contains("Plan: Step 1"),
5855 "squad result should include planner's output"
5856 );
5857 assert!(
5859 squad_text.contains("Computation result: 714"),
5860 "squad result should include worker's output"
5861 );
5862 assert!(
5864 squad_text.contains("Error"),
5865 "squad result should include reviewer's error"
5866 );
5867
5868 assert!(
5870 outer_bb.read("agent:planner").await.unwrap().is_none(),
5871 "outer blackboard should NOT have agent:planner"
5872 );
5873 assert!(
5874 outer_bb.read("agent:worker").await.unwrap().is_none(),
5875 "outer blackboard should NOT have agent:worker"
5876 );
5877
5878 let events = events.lock().expect("test lock");
5880
5881 fn agent_of(e: &AgentEvent) -> &str {
5882 match e {
5883 AgentEvent::RunStarted { agent, .. }
5884 | AgentEvent::TurnStarted { agent, .. }
5885 | AgentEvent::LlmResponse { agent, .. }
5886 | AgentEvent::ToolCallStarted { agent, .. }
5887 | AgentEvent::ToolCallCompleted { agent, .. }
5888 | AgentEvent::RunCompleted { agent, .. }
5889 | AgentEvent::RunFailed { agent, .. }
5890 | AgentEvent::SubAgentsDispatched { agent, .. }
5891 | AgentEvent::SubAgentCompleted { agent, .. }
5892 | AgentEvent::ApprovalRequested { agent, .. }
5893 | AgentEvent::ApprovalDecision { agent, .. }
5894 | AgentEvent::ContextSummarized { agent, .. }
5895 | AgentEvent::GuardrailDenied { agent, .. }
5896 | AgentEvent::GuardrailWarned { agent, .. }
5897 | AgentEvent::RetryAttempt { agent, .. }
5898 | AgentEvent::DoomLoopDetected { agent, .. }
5899 | AgentEvent::FuzzyDoomLoopDetected { agent, .. }
5900 | AgentEvent::AutoCompactionTriggered { agent, .. }
5901 | AgentEvent::SessionPruned { agent, .. }
5902 | AgentEvent::ModelEscalated { agent, .. }
5903 | AgentEvent::BudgetExceeded { agent, .. }
5904 | AgentEvent::AgentSpawned { agent, .. }
5905 | AgentEvent::KillSwitchActivated { agent, .. }
5906 | AgentEvent::ToolNameRepaired { agent, .. } => agent,
5907 AgentEvent::SensorEventProcessed { sensor_name, .. } => sensor_name,
5908 AgentEvent::StoryUpdated { story_id, .. } => story_id,
5909 AgentEvent::TaskRouted { decision, .. } => decision,
5910 AgentEvent::WorkflowNodeStarted { node, .. }
5911 | AgentEvent::WorkflowNodeCompleted { node, .. }
5912 | AgentEvent::WorkflowNodeFailed { node, .. } => node,
5913 }
5914 }
5915
5916 fn event_type(e: &AgentEvent) -> &'static str {
5917 match e {
5918 AgentEvent::RunStarted { .. } => "RunStarted",
5919 AgentEvent::TurnStarted { .. } => "TurnStarted",
5920 AgentEvent::LlmResponse { .. } => "LlmResponse",
5921 AgentEvent::ToolCallStarted { .. } => "ToolCallStarted",
5922 AgentEvent::ToolCallCompleted { .. } => "ToolCallCompleted",
5923 AgentEvent::RunCompleted { .. } => "RunCompleted",
5924 AgentEvent::RunFailed { .. } => "RunFailed",
5925 AgentEvent::SubAgentsDispatched { .. } => "SubAgentsDispatched",
5926 AgentEvent::SubAgentCompleted { .. } => "SubAgentCompleted",
5927 AgentEvent::ApprovalRequested { .. } => "ApprovalRequested",
5928 AgentEvent::ApprovalDecision { .. } => "ApprovalDecision",
5929 AgentEvent::ContextSummarized { .. } => "ContextSummarized",
5930 AgentEvent::GuardrailDenied { .. } => "GuardrailDenied",
5931 AgentEvent::GuardrailWarned { .. } => "GuardrailWarned",
5932 AgentEvent::RetryAttempt { .. } => "RetryAttempt",
5933 AgentEvent::DoomLoopDetected { .. } => "DoomLoopDetected",
5934 AgentEvent::FuzzyDoomLoopDetected { .. } => "FuzzyDoomLoopDetected",
5935 AgentEvent::AutoCompactionTriggered { .. } => "AutoCompactionTriggered",
5936 AgentEvent::SessionPruned { .. } => "SessionPruned",
5937 AgentEvent::SensorEventProcessed { .. } => "SensorEventProcessed",
5938 AgentEvent::StoryUpdated { .. } => "StoryUpdated",
5939 AgentEvent::TaskRouted { .. } => "TaskRouted",
5940 AgentEvent::ModelEscalated { .. } => "ModelEscalated",
5941 AgentEvent::BudgetExceeded { .. } => "BudgetExceeded",
5942 AgentEvent::AgentSpawned { .. } => "AgentSpawned",
5943 AgentEvent::KillSwitchActivated { .. } => "KillSwitchActivated",
5944 AgentEvent::WorkflowNodeStarted { .. } => "WorkflowNodeStarted",
5945 AgentEvent::WorkflowNodeCompleted { .. } => "WorkflowNodeCompleted",
5946 AgentEvent::WorkflowNodeFailed { .. } => "WorkflowNodeFailed",
5947 AgentEvent::ToolNameRepaired { .. } => "ToolNameRepaired",
5948 }
5949 }
5950
5951 let event_summary: Vec<String> = events
5952 .iter()
5953 .enumerate()
5954 .map(|(i, e)| format!("{i}: [{:>12}] {}", agent_of(e), event_type(e)))
5955 .collect();
5956 let event_log = event_summary.join("\n");
5957
5958 let agents_seen: std::collections::HashSet<&str> = events.iter().map(agent_of).collect();
5960 assert!(
5961 agents_seen.contains("orchestrator"),
5962 "missing orchestrator events.\n{event_log}"
5963 );
5964 let has_planner = agents_seen.contains("planner");
5966 let has_worker = agents_seen.contains("worker");
5967 assert!(
5968 has_planner || has_worker,
5969 "should have events from at least one successful squad member.\n{event_log}"
5970 );
5971
5972 let dispatched: Vec<&AgentEvent> = events
5974 .iter()
5975 .filter(|e| matches!(e, AgentEvent::SubAgentsDispatched { .. }))
5976 .collect();
5977 assert_eq!(
5978 dispatched.len(),
5979 1,
5980 "expected exactly 1 SubAgentsDispatched event.\n{event_log}"
5981 );
5982 match dispatched[0] {
5983 AgentEvent::SubAgentsDispatched { agents, agent } => {
5984 assert_eq!(
5985 agent, "squad-leader",
5986 "form_squad uses 'squad-leader' label"
5987 );
5988 assert_eq!(agents.len(), 3, "should dispatch 3 squad members");
5989 assert!(agents.contains(&"planner".to_string()));
5990 assert!(agents.contains(&"worker".to_string()));
5991 assert!(agents.contains(&"reviewer".to_string()));
5992 }
5993 _ => unreachable!(),
5994 }
5995
5996 let completed: Vec<&AgentEvent> = events
5998 .iter()
5999 .filter(|e| matches!(e, AgentEvent::SubAgentCompleted { .. }))
6000 .collect();
6001 assert_eq!(
6002 completed.len(),
6003 4,
6004 "expected 4 SubAgentCompleted events (3 per-agent + 1 aggregate).\n{event_log}"
6005 );
6006
6007 let per_agent: Vec<&AgentEvent> = completed
6009 .iter()
6010 .filter(|e| {
6011 matches!(e, AgentEvent::SubAgentCompleted { agent, .. }
6012 if !agent.starts_with("squad["))
6013 })
6014 .copied()
6015 .collect();
6016 assert_eq!(per_agent.len(), 3, "3 per-agent completion events");
6017
6018 let reviewer_completed = per_agent.iter().find(
6020 |e| matches!(e, AgentEvent::SubAgentCompleted { agent, .. } if agent == "reviewer"),
6021 );
6022 assert!(
6023 reviewer_completed.is_some(),
6024 "should have reviewer SubAgentCompleted"
6025 );
6026 match reviewer_completed.unwrap() {
6027 AgentEvent::SubAgentCompleted { success, .. } => {
6028 assert!(!success, "reviewer should have failed");
6029 }
6030 _ => unreachable!(),
6031 }
6032
6033 let planner_completed = per_agent.iter().find(
6035 |e| matches!(e, AgentEvent::SubAgentCompleted { agent, .. } if agent == "planner"),
6036 );
6037 assert!(
6038 planner_completed.is_some(),
6039 "should have planner SubAgentCompleted"
6040 );
6041 match planner_completed.unwrap() {
6042 AgentEvent::SubAgentCompleted { success, usage, .. } => {
6043 assert!(success, "planner should have succeeded");
6044 assert_eq!(usage.input_tokens, 20);
6045 assert_eq!(usage.output_tokens, 15);
6046 assert_eq!(usage.reasoning_tokens, 8);
6047 }
6048 _ => unreachable!(),
6049 }
6050
6051 let squad_completed = completed.iter().find(|e| {
6053 matches!(e, AgentEvent::SubAgentCompleted { agent, .. }
6054 if agent.starts_with("squad["))
6055 });
6056 assert!(
6057 squad_completed.is_some(),
6058 "should have aggregate squad completion event.\n{event_log}"
6059 );
6060 match squad_completed.unwrap() {
6061 AgentEvent::SubAgentCompleted {
6062 agent,
6063 success,
6064 usage,
6065 } => {
6066 assert!(
6067 agent.contains("planner")
6068 && agent.contains("worker")
6069 && agent.contains("reviewer"),
6070 "aggregate label should list all agents: {agent}"
6071 );
6072 assert!(
6073 !success,
6074 "aggregate should be false because reviewer failed"
6075 );
6076 assert_eq!(usage.input_tokens, 20 + 25 + 35, "aggregate input tokens");
6078 assert_eq!(usage.output_tokens, 15 + 12 + 18, "aggregate output tokens");
6079 }
6080 _ => unreachable!(),
6081 }
6082
6083 let tool_started: Vec<&AgentEvent> = events
6085 .iter()
6086 .filter(|e| {
6087 matches!(e, AgentEvent::ToolCallStarted { tool_name, .. } if tool_name == "compute")
6088 })
6089 .collect();
6090 assert!(
6093 !tool_started.is_empty(),
6094 "should have at least one compute ToolCallStarted.\n{event_log}"
6095 );
6096
6097 let tool_completed: Vec<&AgentEvent> = events
6098 .iter()
6099 .filter(|e| {
6100 matches!(e, AgentEvent::ToolCallCompleted { tool_name, .. } if tool_name == "compute")
6101 })
6102 .collect();
6103 assert!(
6104 !tool_completed.is_empty(),
6105 "should have at least one compute ToolCallCompleted.\n{event_log}"
6106 );
6107 match tool_completed[0] {
6108 AgentEvent::ToolCallCompleted {
6109 output, is_error, ..
6110 } => {
6111 assert!(!is_error, "compute tool should succeed");
6112 assert!(
6113 output.contains("714"),
6114 "compute output should be '714', got: {output}"
6115 );
6116 }
6117 _ => unreachable!(),
6118 }
6119
6120 for agent_name in &["orchestrator", "planner", "worker"] {
6122 let agent_events: Vec<&AgentEvent> = events
6123 .iter()
6124 .filter(|e| agent_of(e) == *agent_name)
6125 .collect();
6126 if !agent_events.is_empty() {
6127 assert!(
6128 matches!(agent_events[0], AgentEvent::RunStarted { .. }),
6129 "first event for '{agent_name}' should be RunStarted, got: {:?}\n{event_log}",
6130 agent_events[0]
6131 );
6132 }
6133 }
6134
6135 let reviewer_events: Vec<&AgentEvent> = events
6137 .iter()
6138 .filter(|e| agent_of(e) == "reviewer")
6139 .collect();
6140 if !reviewer_events.is_empty() {
6141 assert!(
6142 matches!(reviewer_events[0], AgentEvent::RunStarted { .. }),
6143 "reviewer first event should be RunStarted"
6144 );
6145 let has_failed = reviewer_events
6147 .iter()
6148 .any(|e| matches!(e, AgentEvent::RunFailed { .. }));
6149 assert!(
6150 has_failed,
6151 "reviewer should have a RunFailed event.\n{event_log}"
6152 );
6153 }
6154
6155 let dispatch_idx = events
6157 .iter()
6158 .position(|e| matches!(e, AgentEvent::SubAgentsDispatched { .. }));
6159 let first_completed_idx = events
6160 .iter()
6161 .position(|e| matches!(e, AgentEvent::SubAgentCompleted { .. }));
6162 if let (Some(d), Some(c)) = (dispatch_idx, first_completed_idx) {
6163 assert!(
6164 d < c,
6165 "SubAgentsDispatched (idx {d}) should precede SubAgentCompleted (idx {c})\n{event_log}"
6166 );
6167 }
6168
6169 let llm_responses: Vec<&AgentEvent> = events
6171 .iter()
6172 .filter(|e| matches!(e, AgentEvent::LlmResponse { .. }))
6173 .collect();
6174 assert!(
6175 !llm_responses.is_empty(),
6176 "should have LlmResponse events.\n{event_log}"
6177 );
6178 for lr in &llm_responses {
6179 match lr {
6180 AgentEvent::LlmResponse { model, .. } => {
6181 assert_eq!(
6182 model.as_deref(),
6183 Some("mock-model-v1"),
6184 "LlmResponse should carry provider model name"
6185 );
6186 }
6187 _ => unreachable!(),
6188 }
6189 }
6190
6191 assert!(
6201 events.len() >= 15,
6202 "expected at least 15 events for complex squad test, got {}.\n{event_log}",
6203 events.len(),
6204 );
6205 }
6206
6207 #[test]
6208 fn build_rejects_empty_sub_agent_name() {
6209 let provider = Arc::new(MockProvider::new(vec![]));
6210 let result = Orchestrator::builder(provider)
6211 .sub_agent("", "Empty name agent", "prompt")
6212 .build();
6213 match result {
6214 Err(Error::Config(msg)) => {
6215 assert!(
6216 msg.contains("must not be empty"),
6217 "expected empty name error, got: {msg}"
6218 );
6219 }
6220 Err(other) => panic!("expected Config error, got: {other:?}"),
6221 Ok(_) => panic!("expected error for empty sub-agent name"),
6222 }
6223 }
6224
6225 #[tokio::test]
6226 async fn instruction_text_wired_to_orchestrator_system_prompt() {
6227 struct CapturingProvider {
6229 captured_systems: Mutex<Vec<String>>,
6230 }
6231 impl LlmProvider for CapturingProvider {
6232 async fn complete(
6233 &self,
6234 request: CompletionRequest,
6235 ) -> Result<CompletionResponse, Error> {
6236 self.captured_systems
6237 .lock()
6238 .expect("lock")
6239 .push(request.system.clone());
6240 Ok(CompletionResponse {
6242 content: vec![ContentBlock::Text {
6243 text: "Task complete.".into(),
6244 }],
6245 stop_reason: StopReason::EndTurn,
6246 usage: TokenUsage::default(),
6247 model: None,
6248 })
6249 }
6250 }
6251
6252 let provider = Arc::new(CapturingProvider {
6253 captured_systems: Mutex::new(Vec::new()),
6254 });
6255 let mut orchestrator = Orchestrator::builder(provider.clone())
6256 .sub_agent("agent-a", "Does things", "You are agent A.")
6257 .instruction_text("Always verify your work.")
6258 .build()
6259 .unwrap();
6260
6261 let _output = orchestrator.run("test task").await.unwrap();
6262 let systems = provider.captured_systems.lock().expect("lock").clone();
6263 assert!(!systems.is_empty(), "should have at least one LLM call");
6265 let orchestrator_system = &systems[0];
6266 assert!(
6267 orchestrator_system.contains("# Project Instructions"),
6268 "orchestrator system prompt should contain instruction header"
6269 );
6270 assert!(
6271 orchestrator_system.contains("Always verify your work."),
6272 "orchestrator system prompt should contain instruction text"
6273 );
6274 }
6275
6276 #[tokio::test]
6277 async fn permission_rules_propagate_to_sub_agents() {
6278 let provider = Arc::new(MockProvider::new(vec![
6282 CompletionResponse {
6284 content: vec![ContentBlock::ToolUse {
6285 id: "orch-1".into(),
6286 name: "delegate_task".into(),
6287 input: json!({
6288 "tasks": [{"agent": "worker", "task": "run a bash command"}]
6289 }),
6290 }],
6291 stop_reason: StopReason::ToolUse,
6292 usage: TokenUsage::default(),
6293 model: None,
6294 },
6295 CompletionResponse {
6297 content: vec![ContentBlock::ToolUse {
6298 id: "worker-1".into(),
6299 name: "bash".into(),
6300 input: json!({"command": "echo hello"}),
6301 }],
6302 stop_reason: StopReason::ToolUse,
6303 usage: TokenUsage::default(),
6304 model: None,
6305 },
6306 CompletionResponse {
6308 content: vec![ContentBlock::Text {
6309 text: "Bash was denied.".into(),
6310 }],
6311 stop_reason: StopReason::EndTurn,
6312 usage: TokenUsage::default(),
6313 model: None,
6314 },
6315 CompletionResponse {
6317 content: vec![ContentBlock::Text {
6318 text: "Worker reported bash was denied.".into(),
6319 }],
6320 stop_reason: StopReason::EndTurn,
6321 usage: TokenUsage::default(),
6322 model: None,
6323 },
6324 ]));
6325
6326 let events = Arc::new(Mutex::new(Vec::<AgentEvent>::new()));
6327 let events_clone = events.clone();
6328 let on_event: Arc<OnEvent> = Arc::new(move |event: AgentEvent| {
6329 events_clone.lock().expect("test lock").push(event);
6330 });
6331
6332 let deny_bash = crate::agent::permission::PermissionRuleset::new(vec![
6333 crate::agent::permission::PermissionRule {
6334 tool: "bash".into(),
6335 pattern: "*".into(),
6336 action: crate::agent::permission::PermissionAction::Deny,
6337 },
6338 ]);
6339
6340 let bash_tool: Arc<dyn Tool> = Arc::new(MockTool::new("bash", "executed"));
6341
6342 let mut orch = Orchestrator::builder(provider)
6343 .sub_agent_with_tools("worker", "Bash worker", "You run bash.", vec![bash_tool])
6344 .permission_rules(deny_bash)
6345 .on_event(on_event)
6346 .build()
6347 .unwrap();
6348
6349 let output = orch.run("run bash via worker").await.unwrap();
6350 assert_eq!(output.result, "Worker reported bash was denied.");
6351
6352 let events = events.lock().expect("test lock");
6355 let worker_tool_events: Vec<_> = events
6356 .iter()
6357 .filter(|e| {
6358 matches!(
6359 e,
6360 AgentEvent::ToolCallStarted { agent, tool_name, .. }
6361 | AgentEvent::ToolCallCompleted { agent, tool_name, .. }
6362 if agent == "worker" && tool_name == "bash"
6363 )
6364 })
6365 .collect();
6366 assert!(
6367 worker_tool_events.is_empty(),
6368 "bash tool calls in worker should be denied (no events emitted), got: {worker_tool_events:?}"
6369 );
6370 }
6371
6372 #[tokio::test]
6373 async fn permission_rules_propagate_to_squad_members() {
6374 let provider = Arc::new(MockProvider::new(vec![
6376 CompletionResponse {
6378 content: vec![ContentBlock::ToolUse {
6379 id: "orch-1".into(),
6380 name: "form_squad".into(),
6381 input: json!({
6382 "tasks": [
6383 {"agent": "alpha", "task": "run bash"},
6384 {"agent": "beta", "task": "say hello"}
6385 ]
6386 }),
6387 }],
6388 stop_reason: StopReason::ToolUse,
6389 usage: TokenUsage::default(),
6390 model: None,
6391 },
6392 CompletionResponse {
6394 content: vec![ContentBlock::ToolUse {
6395 id: "alpha-1".into(),
6396 name: "bash".into(),
6397 input: json!({"command": "ls"}),
6398 }],
6399 stop_reason: StopReason::ToolUse,
6400 usage: TokenUsage::default(),
6401 model: None,
6402 },
6403 CompletionResponse {
6405 content: vec![ContentBlock::Text {
6406 text: "Bash denied.".into(),
6407 }],
6408 stop_reason: StopReason::EndTurn,
6409 usage: TokenUsage::default(),
6410 model: None,
6411 },
6412 CompletionResponse {
6414 content: vec![ContentBlock::Text {
6415 text: "Hello!".into(),
6416 }],
6417 stop_reason: StopReason::EndTurn,
6418 usage: TokenUsage::default(),
6419 model: None,
6420 },
6421 CompletionResponse {
6423 content: vec![ContentBlock::Text {
6424 text: "Squad done.".into(),
6425 }],
6426 stop_reason: StopReason::EndTurn,
6427 usage: TokenUsage::default(),
6428 model: None,
6429 },
6430 ]));
6431
6432 let deny_bash = crate::agent::permission::PermissionRuleset::new(vec![
6433 crate::agent::permission::PermissionRule {
6434 tool: "bash".into(),
6435 pattern: "*".into(),
6436 action: crate::agent::permission::PermissionAction::Deny,
6437 },
6438 ]);
6439
6440 let bash_tool: Arc<dyn Tool> = Arc::new(MockTool::new("bash", "executed"));
6441
6442 let events = Arc::new(Mutex::new(Vec::<AgentEvent>::new()));
6443 let events_clone = events.clone();
6444 let on_event: Arc<OnEvent> = Arc::new(move |event: AgentEvent| {
6445 events_clone.lock().expect("test lock").push(event);
6446 });
6447
6448 let mut orch = Orchestrator::builder(provider)
6449 .sub_agent_with_tools(
6450 "alpha",
6451 "Alpha agent",
6452 "You run bash.",
6453 vec![bash_tool.clone()],
6454 )
6455 .sub_agent("beta", "Beta agent", "You say hello.")
6456 .permission_rules(deny_bash)
6457 .on_event(on_event)
6458 .build()
6459 .unwrap();
6460
6461 let output = orch.run("form a squad").await.unwrap();
6462 assert_eq!(output.result, "Squad done.");
6463
6464 let events = events.lock().expect("test lock");
6466 let bash_events: Vec<_> = events
6467 .iter()
6468 .filter(|e| {
6469 matches!(
6470 e,
6471 AgentEvent::ToolCallStarted { tool_name, .. }
6472 | AgentEvent::ToolCallCompleted { tool_name, .. }
6473 if tool_name == "bash"
6474 )
6475 })
6476 .collect();
6477 assert!(
6478 bash_events.is_empty(),
6479 "bash tool calls in squad should be denied (no events), got: {bash_events:?}"
6480 );
6481 }
6482
6483 #[test]
6484 fn workspace_propagates_from_builder_to_sub_agents() {
6485 let provider = Arc::new(MockProvider::new(vec![]));
6486 let builder = Orchestrator::builder(provider)
6487 .workspace("/shared/workspace")
6488 .sub_agent_full(SubAgentConfig {
6489 name: "agent1".into(),
6490 description: "test".into(),
6491 workspace: None, ..Default::default()
6493 });
6494
6495 let agent = &builder.sub_agents[0];
6496 assert_eq!(
6497 agent.workspace.as_deref(),
6498 Some(std::path::Path::new("/shared/workspace")),
6499 "sub-agent should inherit workspace from builder"
6500 );
6501 }
6502
6503 #[test]
6504 fn sub_agent_workspace_overrides_builder() {
6505 let provider = Arc::new(MockProvider::new(vec![]));
6506 let builder = Orchestrator::builder(provider)
6507 .workspace("/shared/workspace")
6508 .sub_agent_full(SubAgentConfig {
6509 name: "agent1".into(),
6510 description: "test".into(),
6511 workspace: Some("/custom/workspace".into()),
6512 ..Default::default()
6513 });
6514
6515 let agent = &builder.sub_agents[0];
6516 assert_eq!(
6517 agent.workspace.as_deref(),
6518 Some(std::path::Path::new("/custom/workspace")),
6519 "sub-agent should use its own workspace over builder's"
6520 );
6521 }
6522
6523 #[test]
6524 fn no_workspace_when_builder_has_none() {
6525 let provider = Arc::new(MockProvider::new(vec![]));
6526 let builder = Orchestrator::builder(provider).sub_agent_full(SubAgentConfig {
6527 name: "agent1".into(),
6528 description: "test".into(),
6529 workspace: None,
6530 ..Default::default()
6531 });
6532
6533 let agent = &builder.sub_agents[0];
6534 assert!(
6535 agent.workspace.is_none(),
6536 "sub-agent should have no workspace when builder has none"
6537 );
6538 }
6539
6540 #[test]
6541 fn multi_agent_prompt_enabled_by_default() {
6542 let provider = Arc::new(MockProvider::new(vec![]));
6543 let builder = Orchestrator::builder(provider);
6544 assert!(builder.multi_agent_prompt);
6545 }
6546
6547 #[test]
6548 fn multi_agent_prompt_can_be_disabled() {
6549 let provider = Arc::new(MockProvider::new(vec![]));
6550 let builder = Orchestrator::builder(provider).multi_agent_prompt(false);
6551 assert!(!builder.multi_agent_prompt);
6552 }
6553
6554 #[test]
6555 fn build_injects_collab_prompt_when_enabled() {
6556 let provider = Arc::new(MockProvider::new(vec![]));
6557 let mut builder = Orchestrator::builder(provider).sub_agent(
6558 "writer",
6559 "Writes content",
6560 "You are a writer.",
6561 );
6562 assert!(
6564 !builder.sub_agents[0]
6565 .system_prompt
6566 .contains("MULTI-AGENT COLLABORATION PROTOCOL")
6567 );
6568 if builder.multi_agent_prompt {
6570 for agent in &mut builder.sub_agents {
6571 agent
6572 .system_prompt
6573 .push_str(&crate::agent::prompts::render_collab_prompt(
6574 &agent.name,
6575 &agent.description,
6576 ));
6577 }
6578 }
6579 assert!(
6580 builder.sub_agents[0]
6581 .system_prompt
6582 .contains("MULTI-AGENT COLLABORATION PROTOCOL")
6583 );
6584 assert!(builder.sub_agents[0].system_prompt.contains("`writer`"));
6585 assert!(
6586 builder.sub_agents[0]
6587 .system_prompt
6588 .contains("Writes content")
6589 );
6590 }
6591
6592 #[test]
6593 fn build_omits_collab_prompt_when_disabled() {
6594 let provider = Arc::new(MockProvider::new(vec![]));
6595 let mut builder = Orchestrator::builder(provider)
6596 .multi_agent_prompt(false)
6597 .sub_agent("writer", "Writes content", "You are a writer.");
6598 if builder.multi_agent_prompt {
6600 for agent in &mut builder.sub_agents {
6601 agent
6602 .system_prompt
6603 .push_str(&crate::agent::prompts::render_collab_prompt(
6604 &agent.name,
6605 &agent.description,
6606 ));
6607 }
6608 }
6609 assert!(
6610 !builder.sub_agents[0]
6611 .system_prompt
6612 .contains("MULTI-AGENT COLLABORATION PROTOCOL")
6613 );
6614 assert_eq!(builder.sub_agents[0].system_prompt, "You are a writer.");
6615 }
6616
6617 fn make_spawn_config() -> crate::types::SpawnConfig {
6620 crate::types::SpawnConfig {
6621 max_spawned_agents: 3,
6622 tool_allowlist: vec![],
6623 max_turns: 5,
6624 max_tokens: 1024,
6625 max_total_tokens: 10_000,
6626 }
6627 }
6628
6629 fn build_spawn_tool(
6630 provider: Arc<MockProvider>,
6631 config: crate::types::SpawnConfig,
6632 tools: Vec<Arc<dyn Tool>>,
6633 ) -> SpawnAgentTool {
6634 let mut tool_pool = std::collections::HashMap::new();
6635 for tool in &tools {
6636 let name = tool.definition().name;
6637 if config.tool_allowlist.contains(&name) {
6638 tool_pool.insert(name, tool.clone());
6639 }
6640 }
6641 let cached_definition = SpawnAgentTool::build_definition(&config);
6642 SpawnAgentTool {
6643 shared_provider: Arc::new(BoxedProvider::from_arc(provider)),
6644 spawn_config: config,
6645 tool_pool,
6646 spawn_count: Arc::new(std::sync::atomic::AtomicU32::new(0)),
6647 spawned_names: Arc::new(Mutex::new(std::collections::HashSet::new())),
6648 accumulated_tokens: Arc::new(Mutex::new(TokenUsage::default())),
6649 permission_rules: crate::agent::permission::PermissionRuleset::default(),
6650 shared_memory: None,
6651 memory_namespace_prefix: None,
6652 on_event: None,
6653 on_text: None,
6654 lsp_manager: None,
6655 observability_mode: crate::agent::observability::ObservabilityMode::Production,
6656 workspace: None,
6657 guardrails: vec![],
6658 audit_trail: None,
6659 audit_user_id: None,
6660 audit_tenant_id: None,
6661 audit_delegation_chain: vec![],
6662 cached_definition,
6663 tenant_tracker: None,
6664 }
6665 }
6666
6667 #[tokio::test]
6668 async fn spawn_agent_basic_execution() {
6669 let provider = Arc::new(MockProvider::new(vec![CompletionResponse {
6670 content: vec![ContentBlock::Text {
6671 text: "Tax analysis complete.".into(),
6672 }],
6673 usage: TokenUsage {
6674 input_tokens: 50,
6675 output_tokens: 20,
6676 ..Default::default()
6677 },
6678 stop_reason: StopReason::EndTurn,
6679 model: None,
6680 }]));
6681
6682 let tool = build_spawn_tool(provider, make_spawn_config(), vec![]);
6683
6684 let result = tool
6685 .spawn(SpawnAgentInput {
6686 name: "tax_specialist".into(),
6687 system_prompt: "You are a tax law expert.".into(),
6688 tools: vec![],
6689 task: "Analyze tax implications.".into(),
6690 })
6691 .await
6692 .unwrap();
6693
6694 assert!(!result.is_error);
6695 assert!(result.content.contains("Tax analysis complete."));
6696 assert!(result.content.contains("spawn:tax_specialist"));
6697 assert_eq!(
6698 tool.spawn_count.load(std::sync::atomic::Ordering::Relaxed),
6699 1
6700 );
6701 }
6702
6703 #[tokio::test]
6704 async fn spawn_agent_rejects_invalid_name() {
6705 let provider = Arc::new(MockProvider::new(vec![]));
6706 let tool = build_spawn_tool(provider, make_spawn_config(), vec![]);
6707
6708 let invalid_names = vec![
6709 "Tax-Specialist",
6710 "123abc",
6711 "",
6712 "has spaces",
6713 "../path",
6714 "a/b",
6715 "UPPER",
6716 ];
6717
6718 for name in invalid_names {
6719 let result = tool
6720 .spawn(SpawnAgentInput {
6721 name: name.into(),
6722 system_prompt: "test".into(),
6723 tools: vec![],
6724 task: "test".into(),
6725 })
6726 .await
6727 .unwrap();
6728 assert!(
6729 result.is_error,
6730 "expected error for name '{name}', got success: {}",
6731 result.content
6732 );
6733 assert!(
6734 result.content.contains("Invalid agent name"),
6735 "expected 'Invalid agent name' in error for name '{name}', got: {}",
6736 result.content
6737 );
6738 }
6739 }
6740
6741 #[tokio::test]
6742 async fn spawn_agent_rejects_duplicate_name() {
6743 let provider = Arc::new(MockProvider::new(vec![
6744 CompletionResponse {
6745 content: vec![ContentBlock::Text {
6746 text: "done".into(),
6747 }],
6748 usage: TokenUsage::default(),
6749 stop_reason: StopReason::EndTurn,
6750 model: None,
6751 },
6752 ]));
6754 let tool = build_spawn_tool(provider, make_spawn_config(), vec![]);
6755
6756 let r1 = tool
6758 .spawn(SpawnAgentInput {
6759 name: "helper".into(),
6760 system_prompt: "test".into(),
6761 tools: vec![],
6762 task: "test".into(),
6763 })
6764 .await
6765 .unwrap();
6766 assert!(!r1.is_error);
6767
6768 let r2 = tool
6770 .spawn(SpawnAgentInput {
6771 name: "helper".into(),
6772 system_prompt: "test".into(),
6773 tools: vec![],
6774 task: "test".into(),
6775 })
6776 .await
6777 .unwrap();
6778 assert!(r2.is_error);
6779 assert!(r2.content.contains("already used"));
6780 }
6781
6782 #[tokio::test]
6783 async fn spawn_agent_enforces_tool_allowlist() {
6784 let provider = Arc::new(MockProvider::new(vec![]));
6785 let mut config = make_spawn_config();
6786 config.tool_allowlist = vec!["mock_read".into()];
6787
6788 let mock = MockTool::new("mock_read", "file content");
6789 let tool = build_spawn_tool(provider, config, vec![Arc::new(mock)]);
6790
6791 let result = tool
6793 .spawn(SpawnAgentInput {
6794 name: "reader".into(),
6795 system_prompt: "test".into(),
6796 tools: vec!["bash".into()],
6797 task: "test".into(),
6798 })
6799 .await
6800 .unwrap();
6801 assert!(result.is_error);
6802 assert!(result.content.contains("not in allowlist"));
6803 }
6804
6805 #[tokio::test]
6806 async fn spawn_agent_count_cap() {
6807 let provider = Arc::new(MockProvider::new(vec![CompletionResponse {
6808 content: vec![ContentBlock::Text {
6809 text: "done".into(),
6810 }],
6811 usage: TokenUsage::default(),
6812 stop_reason: StopReason::EndTurn,
6813 model: None,
6814 }]));
6815 let mut config = make_spawn_config();
6816 config.max_spawned_agents = 1;
6817 let tool = build_spawn_tool(provider, config, vec![]);
6818
6819 let r1 = tool
6821 .spawn(SpawnAgentInput {
6822 name: "first".into(),
6823 system_prompt: "test".into(),
6824 tools: vec![],
6825 task: "test".into(),
6826 })
6827 .await
6828 .unwrap();
6829 assert!(!r1.is_error);
6830
6831 let r2 = tool
6833 .spawn(SpawnAgentInput {
6834 name: "second".into(),
6835 system_prompt: "test".into(),
6836 tools: vec![],
6837 task: "test".into(),
6838 })
6839 .await
6840 .unwrap();
6841 assert!(r2.is_error);
6842 assert!(r2.content.contains("Spawn limit reached"));
6843 }
6844
6845 #[tokio::test]
6846 async fn spawn_agent_token_budget_enforcement() {
6847 let provider = Arc::new(MockProvider::new(vec![]));
6848 let mut config = make_spawn_config();
6849 config.max_total_tokens = 100;
6850 let tool = build_spawn_tool(provider, config, vec![]);
6851
6852 {
6854 let mut acc = tool.accumulated_tokens.lock().unwrap();
6855 *acc = TokenUsage {
6856 input_tokens: 60,
6857 output_tokens: 50,
6858 ..Default::default()
6859 };
6860 }
6861
6862 let result = tool
6863 .spawn(SpawnAgentInput {
6864 name: "spender".into(),
6865 system_prompt: "test".into(),
6866 tools: vec![],
6867 task: "test".into(),
6868 })
6869 .await
6870 .unwrap();
6871 assert!(result.is_error);
6872 assert!(result.content.contains("budget exhausted"));
6873 }
6874
6875 #[tokio::test]
6876 async fn spawn_agent_no_delegation_tools() {
6877 let provider = Arc::new(MockProvider::new(vec![]));
6880 let mut config = make_spawn_config();
6881 config.tool_allowlist = vec!["mock_read".into()];
6882
6883 let mock = MockTool::new("mock_read", "content");
6884 let tools: Vec<Arc<dyn Tool>> = vec![Arc::new(mock)];
6885
6886 let mut tool_pool = std::collections::HashMap::new();
6887 for t in &tools {
6888 let name = t.definition().name;
6889 if config.tool_allowlist.contains(&name) {
6890 tool_pool.insert(name, t.clone());
6891 }
6892 }
6893 tool_pool.insert(
6895 "delegate_task".into(),
6896 Arc::new(MockTool::new("delegate_task", "bad")),
6897 );
6898 tool_pool.remove("delegate_task");
6899 tool_pool.remove("form_squad");
6900 tool_pool.remove("spawn_agent");
6901
6902 assert!(tool_pool.contains_key("mock_read"));
6904 assert!(!tool_pool.contains_key("delegate_task"));
6905 assert!(!tool_pool.contains_key("form_squad"));
6906 assert!(!tool_pool.contains_key("spawn_agent"));
6907 drop(provider);
6908 }
6909
6910 #[tokio::test]
6911 async fn spawn_agent_emits_events() {
6912 let events: Arc<Mutex<Vec<AgentEvent>>> = Arc::new(Mutex::new(vec![]));
6913 let events_clone = events.clone();
6914
6915 let provider = Arc::new(MockProvider::new(vec![CompletionResponse {
6916 content: vec![ContentBlock::Text {
6917 text: "result".into(),
6918 }],
6919 usage: TokenUsage::default(),
6920 stop_reason: StopReason::EndTurn,
6921 model: None,
6922 }]));
6923
6924 let mut tool = build_spawn_tool(provider, make_spawn_config(), vec![]);
6925 tool.on_event = Some(Arc::new(move |e: AgentEvent| {
6926 events_clone.lock().unwrap().push(e);
6927 }));
6928
6929 let _ = tool
6930 .spawn(SpawnAgentInput {
6931 name: "emitter".into(),
6932 system_prompt: "test".into(),
6933 tools: vec![],
6934 task: "test".into(),
6935 })
6936 .await;
6937
6938 let events = events.lock().unwrap();
6939 let spawned = events
6940 .iter()
6941 .any(|e| matches!(e, AgentEvent::AgentSpawned { spawned_name, .. } if spawned_name == "spawn:emitter"));
6942 assert!(spawned, "expected AgentSpawned event");
6943
6944 let completed = events
6945 .iter()
6946 .any(|e| matches!(e, AgentEvent::SubAgentCompleted { agent, success, .. } if agent == "spawn:emitter" && *success));
6947 assert!(completed, "expected SubAgentCompleted event");
6948 }
6949
6950 #[tokio::test]
6951 async fn spawn_agent_empty_tools() {
6952 let provider = Arc::new(MockProvider::new(vec![CompletionResponse {
6954 content: vec![ContentBlock::Text {
6955 text: "Pure reasoning response".into(),
6956 }],
6957 usage: TokenUsage::default(),
6958 stop_reason: StopReason::EndTurn,
6959 model: None,
6960 }]));
6961
6962 let tool = build_spawn_tool(provider, make_spawn_config(), vec![]);
6963
6964 let result = tool
6965 .spawn(SpawnAgentInput {
6966 name: "thinker".into(),
6967 system_prompt: "You are a reasoning agent.".into(),
6968 tools: vec![],
6969 task: "Think about this.".into(),
6970 })
6971 .await
6972 .unwrap();
6973 assert!(!result.is_error);
6974 assert!(result.content.contains("Pure reasoning response"));
6975 }
6976
6977 #[tokio::test]
6978 async fn spawn_agent_prompt_too_long() {
6979 let provider = Arc::new(MockProvider::new(vec![]));
6980 let tool = build_spawn_tool(provider, make_spawn_config(), vec![]);
6981
6982 let long_prompt = "x".repeat(SPAWN_MAX_PROMPT_BYTES + 1);
6983 let result = tool
6984 .spawn(SpawnAgentInput {
6985 name: "verbose".into(),
6986 system_prompt: long_prompt,
6987 tools: vec![],
6988 task: "test".into(),
6989 })
6990 .await
6991 .unwrap();
6992 assert!(result.is_error);
6993 assert!(result.content.contains("System prompt too long"));
6994 }
6995
6996 #[test]
6997 fn spawn_config_validation_rejects_zero_agents() {
6998 let toml_str = r#"
6999[provider]
7000name = "anthropic"
7001model = "claude-sonnet-4-20250514"
7002
7003[orchestrator.spawn]
7004max_spawned_agents = 0
7005"#;
7006 let err = crate::config::HeartbitConfig::from_toml(toml_str).unwrap_err();
7007 assert!(
7008 err.to_string()
7009 .contains("max_spawned_agents must be at least 1"),
7010 "err: {err}"
7011 );
7012 }
7013
7014 #[test]
7015 fn spawn_config_validation_rejects_zero_turns() {
7016 let toml_str = r#"
7017[provider]
7018name = "anthropic"
7019model = "claude-sonnet-4-20250514"
7020
7021[orchestrator.spawn]
7022max_turns = 0
7023"#;
7024 let err = crate::config::HeartbitConfig::from_toml(toml_str).unwrap_err();
7025 assert!(
7026 err.to_string().contains("max_turns must be at least 1"),
7027 "err: {err}"
7028 );
7029 }
7030
7031 #[test]
7032 fn spawn_config_from_toml() {
7033 let toml_str = r#"
7034[provider]
7035name = "anthropic"
7036model = "claude-sonnet-4-20250514"
7037
7038[orchestrator.spawn]
7039max_spawned_agents = 5
7040tool_allowlist = ["read", "grep", "bash"]
7041max_turns = 20
7042max_tokens = 8192
7043max_total_tokens = 100000
7044"#;
7045 let config: crate::config::HeartbitConfig = toml::from_str(toml_str).unwrap();
7046 let spawn = config.orchestrator.spawn.as_ref().unwrap();
7047 assert_eq!(spawn.max_spawned_agents, 5);
7048 assert_eq!(spawn.tool_allowlist, vec!["read", "grep", "bash"]);
7049 assert_eq!(spawn.max_turns, 20);
7050 assert_eq!(spawn.max_tokens, 8192);
7051 assert_eq!(spawn.max_total_tokens, 100_000);
7052 }
7053
7054 #[test]
7055 fn spawn_disabled_by_default() {
7056 let toml_str = r#"
7057[provider]
7058name = "anthropic"
7059model = "claude-sonnet-4-20250514"
7060"#;
7061 let config: crate::config::HeartbitConfig = toml::from_str(toml_str).unwrap();
7062 assert!(config.orchestrator.spawn.is_none());
7063 }
7064
7065 #[test]
7066 fn spawn_config_invalid_tool_rejected_at_build() {
7067 let provider = Arc::new(MockProvider::new(vec![]));
7068 let mut config = make_spawn_config();
7069 config.tool_allowlist = vec!["nonexistent_tool".into()];
7070
7071 let result = Orchestrator::builder(provider)
7072 .sub_agent("worker", "does work", "You work.")
7073 .spawn_config(config, vec![]) .build();
7075
7076 match result {
7077 Err(e) => {
7078 let msg = e.to_string();
7079 assert!(
7080 msg.contains("nonexistent_tool"),
7081 "expected error mentioning the bad tool, got: {msg}"
7082 );
7083 }
7084 Ok(_) => panic!("expected build error for invalid tool in allowlist"),
7085 }
7086 }
7087
7088 #[test]
7089 fn spawn_tool_definition_includes_allowlist() {
7090 let mut config = make_spawn_config();
7091 config.tool_allowlist = vec!["read".into(), "grep".into()];
7092 let def = SpawnAgentTool::build_definition(&config);
7093 assert_eq!(def.name, "spawn_agent");
7094 assert!(def.description.contains("read, grep"));
7095 assert!(def.description.contains("3 agents max"));
7096 }
7097
7098 #[tokio::test]
7099 async fn spawn_system_prompt_added_when_configured() {
7100 use std::sync::Arc;
7103
7104 let provider = Arc::new(MockProvider::new(vec![
7105 CompletionResponse {
7107 content: vec![ContentBlock::Text {
7108 text: "No delegation needed.".into(),
7109 }],
7110 usage: TokenUsage::default(),
7111 stop_reason: StopReason::EndTurn,
7112 model: None,
7113 },
7114 ]));
7115
7116 let config = make_spawn_config();
7117 let mut orchestrator = Orchestrator::builder(provider)
7118 .sub_agent("worker", "does work", "You work.")
7119 .spawn_config(config, vec![])
7120 .build()
7121 .unwrap();
7122
7123 let output = orchestrator.run("test task").await.unwrap();
7124 assert!(output.result.contains("No delegation needed."));
7126 }
7127
7128 #[tokio::test(flavor = "multi_thread")]
7152 async fn orchestrator_propagates_tenant_tracker_to_sub_agents() {
7153 use crate::agent::tenant_tracker::TenantTokenTracker;
7154 use crate::auth::TenantScope;
7155
7156 let tracker = Arc::new(TenantTokenTracker::new(1_000_000));
7157
7158 let scope = TenantScope::new("tenant-abc");
7161 drop(tracker.reserve(&scope, 100_000).unwrap());
7162
7163 let initial_snap = tracker.snapshot();
7165 assert_eq!(initial_snap.len(), 1, "entry must exist after reserve+drop");
7166 assert_eq!(initial_snap[0].1.in_flight, 0);
7167
7168 let provider = Arc::new(MockProvider::new(vec![
7172 CompletionResponse {
7174 content: vec![ContentBlock::ToolUse {
7175 id: "tt-call-1".into(),
7176 name: "delegate_task".into(),
7177 input: json!({
7178 "tasks": [{"agent": "worker", "task": "do work"}]
7179 }),
7180 }],
7181 stop_reason: StopReason::ToolUse,
7182 usage: TokenUsage {
7183 input_tokens: 5,
7184 output_tokens: 5,
7185 ..Default::default()
7186 },
7187 model: None,
7188 },
7189 CompletionResponse {
7191 content: vec![ContentBlock::Text {
7192 text: "Work done.".into(),
7193 }],
7194 stop_reason: StopReason::EndTurn,
7195 usage: TokenUsage {
7196 input_tokens: 400,
7197 output_tokens: 100,
7198 ..Default::default()
7199 },
7200 model: None,
7201 },
7202 CompletionResponse {
7204 content: vec![ContentBlock::Text {
7205 text: "All done.".into(),
7206 }],
7207 stop_reason: StopReason::EndTurn,
7208 usage: TokenUsage {
7209 input_tokens: 2,
7210 output_tokens: 2,
7211 ..Default::default()
7212 },
7213 model: None,
7214 },
7215 ]));
7216
7217 let mut orch = Orchestrator::builder(provider)
7220 .audit_user_context("user-1", "tenant-abc")
7221 .sub_agent("worker", "Does work", "You work.")
7222 .tenant_tracker(tracker.clone())
7223 .build()
7224 .unwrap();
7225
7226 orch.run("do work").await.unwrap();
7227
7228 let snap_mid = tracker.snapshot();
7231 let state_mid = snap_mid
7232 .iter()
7233 .find(|(tid, _)| tid == "tenant-abc")
7234 .map(|(_, s)| s.clone())
7235 .expect("entry for 'tenant-abc' should still exist after the run");
7236
7237 assert!(
7238 state_mid.high_water >= 500,
7239 "high_water ({}) must be >= 500 (sub-agent's 500 tokens); \
7240 if it is < 500, the sub-agent runner is not propagating the tracker",
7241 state_mid.high_water
7242 );
7243
7244 drop(orch);
7247
7248 let snap_final = tracker.snapshot();
7249 let state_final = snap_final
7250 .iter()
7251 .find(|(tid, _)| tid == "tenant-abc")
7252 .map(|(_, s)| s.clone())
7253 .expect("entry should still exist after drop");
7254
7255 assert_eq!(
7256 state_final.in_flight, 0,
7257 "in_flight should return to 0 after all runners are dropped"
7258 );
7259 }
7260}