recursive/
multi.rs

1//! Multi-agent orchestration: agent pool, role definitions, and message bus.
2
3use crate::agent::PlanningMode;
4use crate::kernel::{AgentKernel, TurnContext, TurnOutcome};
5use crate::message::Message;
6use crate::permissions::PermissionMode;
7use crate::{Config, LlmProvider};
8use std::collections::HashMap;
9use std::sync::atomic::{AtomicU64, Ordering};
10use std::sync::Arc;
11use std::time::{SystemTime, UNIX_EPOCH};
12use tokio::sync::{broadcast, RwLock};
13
14/// Shared memory store for multi-agent coordination.
15#[derive(Clone)]
16pub struct SharedMemory {
17    store: Arc<RwLock<HashMap<String, MemoryEntry>>>,
18}
19
20/// A single entry in the shared memory store.
21#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
22pub struct MemoryEntry {
23    pub key: String,
24    pub value: String,
25    pub author: String,
26    pub timestamp: u64,
27}
28
29impl SharedMemory {
30    pub fn new() -> Self {
31        Self {
32            store: Arc::new(RwLock::new(HashMap::new())),
33        }
34    }
35
36    pub async fn set(&self, key: String, value: String, author: String) {
37        let timestamp = SystemTime::now()
38            .duration_since(UNIX_EPOCH)
39            .unwrap_or_default()
40            .as_secs();
41        let entry = MemoryEntry {
42            key: key.clone(),
43            value,
44            author,
45            timestamp,
46        };
47        self.store.write().await.insert(key, entry);
48    }
49
50    pub async fn get(&self, key: &str) -> Option<MemoryEntry> {
51        self.store.read().await.get(key).cloned()
52    }
53
54    pub async fn keys(&self) -> Vec<String> {
55        self.store.read().await.keys().cloned().collect()
56    }
57
58    pub async fn all(&self) -> Vec<MemoryEntry> {
59        self.store.read().await.values().cloned().collect()
60    }
61
62    pub async fn remove(&self, key: &str) -> bool {
63        self.store.write().await.remove(key).is_some()
64    }
65
66    pub async fn to_context_string(&self) -> String {
67        let store = self.store.read().await;
68        if store.is_empty() {
69            return String::new();
70        }
71        let mut lines = vec!["[Shared Memory]".to_string()];
72        for entry in store.values() {
73            lines.push(format!(
74                "- {} = {} (by {})",
75                entry.key, entry.value, entry.author
76            ));
77        }
78        lines.join("\n")
79    }
80
81    pub async fn len(&self) -> usize {
82        self.store.read().await.len()
83    }
84
85    pub async fn is_empty(&self) -> bool {
86        self.store.read().await.is_empty()
87    }
88}
89
90impl Default for SharedMemory {
91    fn default() -> Self {
92        Self::new()
93    }
94}
95
96// --- Inter-agent messaging ---
97
98/// Message type for inter-agent communication.
99#[derive(Clone, Debug, serde::Serialize, serde::Deserialize, PartialEq)]
100pub enum MessageType {
101    Task,
102    Result,
103    Question,
104    Feedback,
105    Broadcast,
106}
107
108/// A message exchanged between agents via the message bus.
109#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)]
110pub struct AgentMessage {
111    pub id: String,
112    pub from: String,
113    pub to: String,
114    pub content: String,
115    pub msg_type: MessageType,
116    pub timestamp: u64,
117}
118
119/// Generate a unique message ID using blake3 hash of timestamp + atomic counter.
120fn generate_message_id() -> String {
121    static COUNTER: AtomicU64 = AtomicU64::new(0);
122
123    let count = COUNTER.fetch_add(1, Ordering::Relaxed);
124    let now = SystemTime::now()
125        .duration_since(UNIX_EPOCH)
126        .unwrap_or_default();
127    let input = format!("msg-{}-{}", now.as_nanos(), count);
128    let hash = blake3::hash(input.as_bytes());
129    hash.to_hex()[..16].to_string()
130}
131
132/// Get current timestamp as seconds since UNIX epoch.
133fn now_timestamp() -> u64 {
134    SystemTime::now()
135        .duration_since(UNIX_EPOCH)
136        .unwrap_or_default()
137        .as_secs()
138}
139
140/// An inter-agent message bus supporting publish/subscribe and history.
141#[derive(Clone)]
142pub struct MessageBus {
143    messages: Arc<RwLock<Vec<AgentMessage>>>,
144    subscribers: Arc<RwLock<HashMap<String, broadcast::Sender<AgentMessage>>>>,
145}
146
147impl MessageBus {
148    pub fn new() -> Self {
149        Self {
150            messages: Arc::new(RwLock::new(Vec::new())),
151            subscribers: Arc::new(RwLock::new(HashMap::new())),
152        }
153    }
154
155    /// Send a message. Stores in history and notifies relevant subscribers.
156    pub async fn send(&self, msg: AgentMessage) {
157        self.messages.write().await.push(msg.clone());
158        let subs = self.subscribers.read().await;
159        if msg.to == "broadcast" {
160            for tx in subs.values() {
161                let _ = tx.send(msg.clone());
162            }
163        } else if let Some(tx) = subs.get(&msg.to) {
164            let _ = tx.send(msg);
165        }
166    }
167
168    /// Subscribe to messages for a given role. Returns a broadcast receiver.
169    pub async fn subscribe(&self, role: &str) -> broadcast::Receiver<AgentMessage> {
170        let mut subs = self.subscribers.write().await;
171        let tx = subs.entry(role.to_string()).or_insert_with(|| {
172            let (tx, _) = broadcast::channel(64);
173            tx
174        });
175        tx.subscribe()
176    }
177
178    /// Get all messages addressed to this role (including broadcasts).
179    pub async fn inbox(&self, role: &str) -> Vec<AgentMessage> {
180        self.messages
181            .read()
182            .await
183            .iter()
184            .filter(|m| m.to == role || m.to == "broadcast")
185            .cloned()
186            .collect()
187    }
188
189    /// Get all messages sent by this role.
190    pub async fn outbox(&self, role: &str) -> Vec<AgentMessage> {
191        self.messages
192            .read()
193            .await
194            .iter()
195            .filter(|m| m.from == role)
196            .cloned()
197            .collect()
198    }
199
200    /// Get the full message history.
201    pub async fn history(&self) -> Vec<AgentMessage> {
202        self.messages.read().await.clone()
203    }
204
205    /// Clear all stored messages.
206    pub async fn clear(&self) {
207        self.messages.write().await.clear();
208    }
209}
210
211impl Default for MessageBus {
212    fn default() -> Self {
213        Self::new()
214    }
215}
216
217/// Definition of an agent role.
218#[derive(Clone, Debug)]
219pub struct AgentRole {
220    pub name: String,
221    pub system_prompt: String,
222    pub max_steps: usize,
223    pub allowed_tools: Vec<String>,
224}
225
226/// An agent pool manages multiple agents with different roles.
227pub struct AgentPool {
228    roles: HashMap<String, AgentRole>,
229    provider: Arc<dyn LlmProvider>,
230    #[allow(dead_code)]
231    config: Config,
232    memory: SharedMemory,
233    bus: MessageBus,
234}
235
236impl AgentPool {
237    pub fn new(provider: Arc<dyn LlmProvider>, config: Config) -> Self {
238        Self {
239            roles: HashMap::new(),
240            provider,
241            config,
242            memory: SharedMemory::new(),
243            bus: MessageBus::new(),
244        }
245    }
246
247    pub fn memory(&self) -> &SharedMemory {
248        &self.memory
249    }
250
251    pub fn bus(&self) -> &MessageBus {
252        &self.bus
253    }
254
255    pub fn add_role(&mut self, role: AgentRole) {
256        self.roles.insert(role.name.clone(), role);
257    }
258
259    pub fn get_role(&self, name: &str) -> Option<&AgentRole> {
260        self.roles.get(name)
261    }
262
263    pub fn role_names(&self) -> Vec<&str> {
264        self.roles.keys().map(|s| s.as_str()).collect()
265    }
266
267    pub fn role_count(&self) -> usize {
268        self.roles.len()
269    }
270
271    /// Remove a role from the pool.  Returns `true` if the role existed.
272    pub fn remove_role(&mut self, name: &str) -> bool {
273        self.roles.remove(name).is_some()
274    }
275
276    pub async fn run_with_role(
277        &self,
278        role_name: &str,
279        goal: &str,
280    ) -> Result<TurnOutcome, crate::Error> {
281        let role = self
282            .roles
283            .get(role_name)
284            .ok_or_else(|| crate::Error::Config {
285                message: format!("unknown role: {role_name}"),
286            })?;
287
288        let memory_ctx = self.memory.to_context_string().await;
289        let system_prompt = if memory_ctx.is_empty() {
290            role.system_prompt.clone()
291        } else {
292            format!("{}\n\n{}", role.system_prompt, memory_ctx)
293        };
294
295        let kernel = AgentKernel::builder()
296            .llm(self.provider.clone())
297            .max_steps(role.max_steps)
298            .build()?;
299
300        let ctx = TurnContext {
301            messages: vec![
302                Message::system(system_prompt),
303                Message::user(goal.to_string()),
304            ],
305            step_events_tx: None,
306            plan_confirmed: false,
307            plan_buffer: None,
308            tool_specs: kernel.tools().specs(),
309            streaming: false,
310            permission_hook: None,
311            planning_mode: PlanningMode::default(),
312            exploring_plan_mode: std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false)),
313            permission_mode: PermissionMode::Default,
314            mailbox: None,
315        };
316
317        kernel.run(ctx).await
318    }
319
320    /// Send a task message from one agent role to another.
321    pub async fn send_task(&self, from: &str, to: &str, content: &str) {
322        self.bus
323            .send(AgentMessage {
324                id: generate_message_id(),
325                from: from.to_string(),
326                to: to.to_string(),
327                content: content.to_string(),
328                msg_type: MessageType::Task,
329                timestamp: now_timestamp(),
330            })
331            .await;
332    }
333
334    /// Send a result message from one agent role to another.
335    pub async fn send_result(&self, from: &str, to: &str, content: &str) {
336        self.bus
337            .send(AgentMessage {
338                id: generate_message_id(),
339                from: from.to_string(),
340                to: to.to_string(),
341                content: content.to_string(),
342                msg_type: MessageType::Result,
343                timestamp: now_timestamp(),
344            })
345            .await;
346    }
347}
348
349/// A pipeline chains multiple agent roles in sequence.
350pub struct Pipeline {
351    stages: Vec<String>,
352}
353
354impl Pipeline {
355    pub fn new(stages: Vec<String>) -> Self {
356        Self { stages }
357    }
358
359    /// Execute the pipeline. Each stage's output becomes next stage's input.
360    pub async fn execute(
361        &self,
362        pool: &AgentPool,
363        initial_goal: &str,
364    ) -> Result<PipelineResult, crate::Error> {
365        let mut current_input = initial_goal.to_string();
366        let mut stage_outcomes = Vec::new();
367
368        for role_name in &self.stages {
369            let outcome = pool.run_with_role(role_name, &current_input).await?;
370
371            let output = outcome
372                .new_messages
373                .iter()
374                .rev()
375                .find(|m| m.role == crate::message::Role::Assistant)
376                .map(|m| m.content.clone())
377                .unwrap_or_default();
378
379            stage_outcomes.push(StageOutcome {
380                role: role_name.clone(),
381                output: output.clone(),
382                steps: outcome.steps,
383            });
384
385            current_input = output;
386        }
387
388        Ok(PipelineResult {
389            stages: stage_outcomes,
390        })
391    }
392}
393
394/// The result of running a full pipeline.
395#[derive(Debug)]
396pub struct PipelineResult {
397    pub stages: Vec<StageOutcome>,
398}
399
400/// The outcome of a single pipeline stage.
401#[derive(Debug)]
402pub struct StageOutcome {
403    pub role: String,
404    pub output: String,
405    pub steps: usize,
406}
407
408impl PipelineResult {
409    pub fn final_output(&self) -> &str {
410        self.stages.last().map(|s| s.output.as_str()).unwrap_or("")
411    }
412    pub fn stage_count(&self) -> usize {
413        self.stages.len()
414    }
415}
416
417/// A team orchestrator uses a lead agent to dynamically assign work
418/// to specialist agents.
419pub struct TeamOrchestrator {
420    lead_role: String,
421    available_roles: Vec<String>,
422}
423
424impl TeamOrchestrator {
425    pub fn new(lead_role: String, available_roles: Vec<String>) -> Self {
426        Self {
427            lead_role,
428            available_roles,
429        }
430    }
431
432    /// Run orchestration: lead plans delegations, specialists execute, lead synthesizes.
433    pub async fn run(&self, pool: &AgentPool, goal: &str) -> Result<TeamResult, crate::Error> {
434        // Phase 1: Ask lead to plan
435        let delegation_prompt = format!(
436            "{}\n\nAvailable specialists: {}\n\nTo delegate, use: DELEGATE:<role>:<task>\nWhen done, provide your final answer.",
437            goal,
438            self.available_roles.join(", ")
439        );
440
441        let lead_outcome = pool
442            .run_with_role(&self.lead_role, &delegation_prompt)
443            .await?;
444        let lead_response = lead_outcome
445            .new_messages
446            .iter()
447            .rev()
448            .find(|m| m.role == crate::message::Role::Assistant)
449            .map(|m| m.content.clone())
450            .unwrap_or_default();
451
452        let delegations = parse_delegations(&lead_response);
453
454        // Phase 2: Execute delegations
455        let mut delegation_results = Vec::new();
456        for (role, task) in &delegations {
457            if self.available_roles.contains(role) {
458                match pool.run_with_role(role, task).await {
459                    Ok(outcome) => {
460                        let result = outcome
461                            .new_messages
462                            .iter()
463                            .rev()
464                            .find(|m| m.role == crate::message::Role::Assistant)
465                            .map(|m| m.content.clone())
466                            .unwrap_or_default();
467                        delegation_results.push(DelegationResult {
468                            role: role.clone(),
469                            task: task.clone(),
470                            output: result,
471                            success: true,
472                        });
473                    }
474                    Err(e) => {
475                        delegation_results.push(DelegationResult {
476                            role: role.clone(),
477                            task: task.clone(),
478                            output: format!("Error: {e}"),
479                            success: false,
480                        });
481                    }
482                }
483            }
484        }
485
486        // Phase 3: If delegations happened, synthesize
487        let final_output = if delegation_results.is_empty() {
488            lead_response
489        } else {
490            let results_summary = delegation_results
491                .iter()
492                .map(|r| {
493                    format!(
494                        "- {} ({}): {}",
495                        r.role,
496                        if r.success { "ok" } else { "failed" },
497                        r.output
498                    )
499                })
500                .collect::<Vec<_>>()
501                .join("\n");
502
503            let synthesis_prompt = format!(
504                "Results from delegated tasks:\n\n{}\n\nProvide a final synthesis.",
505                results_summary
506            );
507
508            let synthesis = pool
509                .run_with_role(&self.lead_role, &synthesis_prompt)
510                .await?;
511            synthesis
512                .new_messages
513                .iter()
514                .rev()
515                .find(|m| m.role == crate::message::Role::Assistant)
516                .map(|m| m.content.clone())
517                .unwrap_or_default()
518        };
519
520        Ok(TeamResult {
521            delegations: delegation_results,
522            final_output,
523        })
524    }
525}
526
527/// Parse "DELEGATE:<role>:<task>" lines from text.
528pub fn parse_delegations(text: &str) -> Vec<(String, String)> {
529    text.lines()
530        .filter_map(|line| {
531            let trimmed = line.trim();
532            if let Some(rest) = trimmed.strip_prefix("DELEGATE:") {
533                let parts: Vec<&str> = rest.splitn(2, ':').collect();
534                if parts.len() == 2 {
535                    Some((parts[0].to_string(), parts[1].to_string()))
536                } else {
537                    None
538                }
539            } else {
540                None
541            }
542        })
543        .collect()
544}
545
546/// The result of a team orchestration run.
547#[derive(Debug)]
548pub struct TeamResult {
549    pub delegations: Vec<DelegationResult>,
550    pub final_output: String,
551}
552
553/// The result of a single delegation to a specialist.
554#[derive(Debug)]
555pub struct DelegationResult {
556    pub role: String,
557    pub task: String,
558    pub output: String,
559    pub success: bool,
560}
561
562/// Default role set for common multi-agent patterns.
563pub fn default_roles() -> Vec<AgentRole> {
564    vec![
565        AgentRole {
566            name: "planner".into(),
567            system_prompt: "You are a planning agent. Analyze the task, break it into steps, \
568                            and output a structured plan. Do not execute — only plan."
569                .into(),
570            max_steps: 10,
571            allowed_tools: vec![],
572        },
573        AgentRole {
574            name: "coder".into(),
575            system_prompt: "You are a coding agent. Implement the task using the available \
576                            tools. Write code, run tests, fix errors."
577                .into(),
578            max_steps: 50,
579            allowed_tools: vec![],
580        },
581        AgentRole {
582            name: "reviewer".into(),
583            system_prompt: "You are a code review agent. Read the code changes, identify \
584                            issues, suggest improvements. Do not modify files."
585                .into(),
586            max_steps: 20,
587            allowed_tools: vec!["read_file".into(), "search_files".into()],
588        },
589    ]
590}
591
592#[cfg(test)]
593mod tests {
594    use super::*;
595    use crate::llm::{Completion, MockProvider};
596    use std::path::PathBuf;
597
598    fn test_config() -> Config {
599        Config {
600            workspace: PathBuf::from("."),
601            api_base: String::new(),
602            api_key: None,
603            model: String::new(),
604            provider_type: "openai".into(),
605            preset: None,
606            max_steps: 32,
607            temperature: 0.2,
608            system_prompt: String::new(),
609            retry_max: 2,
610            retry_initial_backoff_secs: 1,
611            retry_max_backoff_secs: 8,
612            shell_timeout_secs: 300,
613            headless: false,
614            memory_summary_limit: 5,
615            thinking_budget: None,
616            session_name: None,
617        }
618    }
619
620    #[test]
621    fn new_pool_is_empty() {
622        let provider = Arc::new(MockProvider::new(vec![]));
623        let pool = AgentPool::new(provider, test_config());
624        assert_eq!(pool.role_count(), 0);
625    }
626
627    #[test]
628    fn add_role_and_get_role() {
629        let provider = Arc::new(MockProvider::new(vec![]));
630        let mut pool = AgentPool::new(provider, test_config());
631
632        let role = AgentRole {
633            name: "tester".into(),
634            system_prompt: "You test things.".into(),
635            max_steps: 5,
636            allowed_tools: vec!["run_shell".into()],
637        };
638        pool.add_role(role.clone());
639
640        let retrieved = pool.get_role("tester").unwrap();
641        assert_eq!(retrieved.name, "tester");
642        assert_eq!(retrieved.system_prompt, "You test things.");
643        assert_eq!(retrieved.max_steps, 5);
644        assert_eq!(retrieved.allowed_tools, vec!["run_shell"]);
645    }
646
647    #[test]
648    fn role_names_returns_all_registered() {
649        let provider = Arc::new(MockProvider::new(vec![]));
650        let mut pool = AgentPool::new(provider, test_config());
651
652        pool.add_role(AgentRole {
653            name: "alpha".into(),
654            system_prompt: "A".into(),
655            max_steps: 1,
656            allowed_tools: vec![],
657        });
658        pool.add_role(AgentRole {
659            name: "beta".into(),
660            system_prompt: "B".into(),
661            max_steps: 2,
662            allowed_tools: vec![],
663        });
664
665        let mut names = pool.role_names();
666        names.sort();
667        assert_eq!(names, vec!["alpha", "beta"]);
668        assert_eq!(pool.role_count(), 2);
669    }
670
671    #[tokio::test]
672    async fn run_with_unknown_role_returns_error() {
673        let provider = Arc::new(MockProvider::new(vec![]));
674        let pool = AgentPool::new(provider, test_config());
675
676        let result = pool.run_with_role("nonexistent", "do something").await;
677        assert!(result.is_err());
678        let err = result.unwrap_err();
679        assert!(err.to_string().contains("unknown role"));
680    }
681
682    #[tokio::test]
683    async fn run_with_role_succeeds_with_mock() {
684        let provider = Arc::new(MockProvider::new(vec![Completion {
685            content: "Plan: step 1, step 2, step 3".into(),
686            tool_calls: vec![],
687            finish_reason: Some("stop".into()),
688            usage: None,
689            reasoning_content: None,
690        }]));
691
692        let mut pool = AgentPool::new(provider, test_config());
693        pool.add_role(AgentRole {
694            name: "planner".into(),
695            system_prompt: "You are a planner.".into(),
696            max_steps: 5,
697            allowed_tools: vec![],
698        });
699
700        let outcome = pool.run_with_role("planner", "plan a task").await.unwrap();
701        assert_eq!(
702            outcome.finish_reason,
703            crate::agent::FinishReason::NoMoreToolCalls
704        );
705        assert!(outcome.final_text.unwrap().contains("Plan:"));
706    }
707
708    #[test]
709    fn default_roles_returns_three_roles() {
710        let roles = default_roles();
711        assert_eq!(roles.len(), 3);
712
713        let names: Vec<&str> = roles.iter().map(|r| r.name.as_str()).collect();
714        assert!(names.contains(&"planner"));
715        assert!(names.contains(&"coder"));
716        assert!(names.contains(&"reviewer"));
717    }
718
719    #[tokio::test]
720    async fn shared_memory_set_and_get() {
721        let mem = SharedMemory::new();
722        mem.set("goal".into(), "build feature X".into(), "planner".into())
723            .await;
724
725        let entry = mem.get("goal").await.unwrap();
726        assert_eq!(entry.key, "goal");
727        assert_eq!(entry.value, "build feature X");
728        assert_eq!(entry.author, "planner");
729        assert!(entry.timestamp > 0);
730    }
731
732    #[tokio::test]
733    async fn shared_memory_keys() {
734        let mem = SharedMemory::new();
735        mem.set("a".into(), "1".into(), "agent1".into()).await;
736        mem.set("b".into(), "2".into(), "agent2".into()).await;
737
738        let mut keys = mem.keys().await;
739        keys.sort();
740        assert_eq!(keys, vec!["a", "b"]);
741        assert_eq!(mem.len().await, 2);
742        assert!(!mem.is_empty().await);
743    }
744
745    #[tokio::test]
746    async fn shared_memory_remove() {
747        let mem = SharedMemory::new();
748        mem.set("tmp".into(), "val".into(), "x".into()).await;
749        assert!(mem.get("tmp").await.is_some());
750
751        let removed = mem.remove("tmp").await;
752        assert!(removed);
753        assert!(mem.get("tmp").await.is_none());
754
755        // Removing non-existent key returns false
756        let removed_again = mem.remove("tmp").await;
757        assert!(!removed_again);
758    }
759
760    #[tokio::test]
761    async fn shared_memory_to_context_string() {
762        let mem = SharedMemory::new();
763        mem.set("status".into(), "in-progress".into(), "coder".into())
764            .await;
765
766        let ctx = mem.to_context_string().await;
767        assert!(ctx.contains("[Shared Memory]"));
768        assert!(ctx.contains("status = in-progress (by coder)"));
769    }
770
771    #[tokio::test]
772    async fn shared_memory_empty_context_returns_empty() {
773        let mem = SharedMemory::new();
774        let ctx = mem.to_context_string().await;
775        assert!(ctx.is_empty());
776    }
777
778    #[tokio::test]
779    async fn agent_pool_includes_memory_context() {
780        let provider = Arc::new(MockProvider::new(vec![Completion {
781            content: "I see the shared memory context.".into(),
782            tool_calls: vec![],
783            finish_reason: Some("stop".into()),
784            usage: None,
785            reasoning_content: None,
786        }]));
787
788        let mut pool = AgentPool::new(provider, test_config());
789        pool.add_role(AgentRole {
790            name: "worker".into(),
791            system_prompt: "You are a worker.".into(),
792            max_steps: 5,
793            allowed_tools: vec![],
794        });
795
796        // Set memory before running
797        pool.memory()
798            .set("plan".into(), "step 1 done".into(), "planner".into())
799            .await;
800
801        let outcome = pool.run_with_role("worker", "continue work").await.unwrap();
802        assert_eq!(
803            outcome.finish_reason,
804            crate::agent::FinishReason::NoMoreToolCalls
805        );
806        // The run succeeded with memory context injected — no error means integration works
807        assert!(outcome.final_text.is_some());
808    }
809
810    // --- MessageBus tests ---
811
812    fn make_msg(from: &str, to: &str, content: &str, msg_type: MessageType) -> AgentMessage {
813        AgentMessage {
814            id: generate_message_id(),
815            from: from.to_string(),
816            to: to.to_string(),
817            content: content.to_string(),
818            msg_type,
819            timestamp: now_timestamp(),
820        }
821    }
822
823    #[tokio::test]
824    async fn message_bus_send_and_inbox() {
825        let bus = MessageBus::new();
826        let msg = make_msg("planner", "coder", "implement feature X", MessageType::Task);
827        bus.send(msg).await;
828
829        let inbox = bus.inbox("coder").await;
830        assert_eq!(inbox.len(), 1);
831        assert_eq!(inbox[0].content, "implement feature X");
832        assert_eq!(inbox[0].from, "planner");
833        assert_eq!(inbox[0].msg_type, MessageType::Task);
834
835        // Other roles see empty inbox
836        let empty = bus.inbox("reviewer").await;
837        assert!(empty.is_empty());
838    }
839
840    #[tokio::test]
841    async fn message_bus_outbox() {
842        let bus = MessageBus::new();
843        bus.send(make_msg(
844            "coder",
845            "reviewer",
846            "done coding",
847            MessageType::Result,
848        ))
849        .await;
850        bus.send(make_msg(
851            "coder",
852            "planner",
853            "need clarification",
854            MessageType::Question,
855        ))
856        .await;
857        bus.send(make_msg(
858            "planner",
859            "coder",
860            "here is the plan",
861            MessageType::Task,
862        ))
863        .await;
864
865        let outbox = bus.outbox("coder").await;
866        assert_eq!(outbox.len(), 2);
867        assert!(outbox.iter().all(|m| m.from == "coder"));
868
869        let planner_outbox = bus.outbox("planner").await;
870        assert_eq!(planner_outbox.len(), 1);
871    }
872
873    #[tokio::test]
874    async fn message_bus_broadcast_reaches_all() {
875        let bus = MessageBus::new();
876        bus.send(make_msg(
877            "admin",
878            "broadcast",
879            "system update",
880            MessageType::Broadcast,
881        ))
882        .await;
883
884        let coder_inbox = bus.inbox("coder").await;
885        let reviewer_inbox = bus.inbox("reviewer").await;
886        let planner_inbox = bus.inbox("planner").await;
887
888        assert_eq!(coder_inbox.len(), 1);
889        assert_eq!(reviewer_inbox.len(), 1);
890        assert_eq!(planner_inbox.len(), 1);
891        assert_eq!(coder_inbox[0].content, "system update");
892    }
893
894    #[tokio::test]
895    async fn message_bus_subscribe_receives() {
896        let bus = MessageBus::new();
897        let mut rx = bus.subscribe("coder").await;
898
899        // Send after subscribing
900        let msg = make_msg("planner", "coder", "task for you", MessageType::Task);
901        bus.send(msg).await;
902
903        let received = rx.recv().await.unwrap();
904        assert_eq!(received.content, "task for you");
905        assert_eq!(received.from, "planner");
906    }
907
908    #[tokio::test]
909    async fn message_bus_history() {
910        let bus = MessageBus::new();
911        bus.send(make_msg("a", "b", "msg1", MessageType::Task))
912            .await;
913        bus.send(make_msg("b", "a", "msg2", MessageType::Result))
914            .await;
915        bus.send(make_msg("a", "broadcast", "msg3", MessageType::Broadcast))
916            .await;
917
918        let history = bus.history().await;
919        assert_eq!(history.len(), 3);
920        assert_eq!(history[0].content, "msg1");
921        assert_eq!(history[1].content, "msg2");
922        assert_eq!(history[2].content, "msg3");
923    }
924
925    #[tokio::test]
926    async fn message_bus_clear() {
927        let bus = MessageBus::new();
928        bus.send(make_msg("a", "b", "hello", MessageType::Task))
929            .await;
930        assert_eq!(bus.history().await.len(), 1);
931
932        bus.clear().await;
933        assert!(bus.history().await.is_empty());
934        assert!(bus.inbox("b").await.is_empty());
935    }
936
937    #[tokio::test]
938    async fn agent_pool_send_task_convenience() {
939        let provider = Arc::new(MockProvider::new(vec![]));
940        let pool = AgentPool::new(provider, test_config());
941
942        pool.send_task("planner", "coder", "build module Y").await;
943        pool.send_result("coder", "planner", "module Y complete")
944            .await;
945
946        let inbox = pool.bus().inbox("coder").await;
947        assert_eq!(inbox.len(), 1);
948        assert_eq!(inbox[0].content, "build module Y");
949        assert_eq!(inbox[0].msg_type, MessageType::Task);
950
951        let planner_inbox = pool.bus().inbox("planner").await;
952        assert_eq!(planner_inbox.len(), 1);
953        assert_eq!(planner_inbox[0].content, "module Y complete");
954        assert_eq!(planner_inbox[0].msg_type, MessageType::Result);
955
956        let history = pool.bus().history().await;
957        assert_eq!(history.len(), 2);
958    }
959
960    // --- Pipeline tests ---
961
962    #[tokio::test]
963    async fn pipeline_empty_returns_empty_result() {
964        let provider = Arc::new(MockProvider::new(vec![]));
965        let pool = AgentPool::new(provider, test_config());
966
967        let pipeline = Pipeline::new(vec![]);
968        let result = pipeline.execute(&pool, "hello").await.unwrap();
969        assert_eq!(result.stage_count(), 0);
970        assert_eq!(result.final_output(), "");
971    }
972
973    #[tokio::test]
974    async fn pipeline_single_stage() {
975        let provider = Arc::new(MockProvider::new(vec![Completion {
976            content: "stage one output".into(),
977            tool_calls: vec![],
978            finish_reason: Some("stop".into()),
979            usage: None,
980            reasoning_content: None,
981        }]));
982
983        let mut pool = AgentPool::new(provider, test_config());
984        pool.add_role(AgentRole {
985            name: "writer".into(),
986            system_prompt: "You write things.".into(),
987            max_steps: 5,
988            allowed_tools: vec![],
989        });
990
991        let pipeline = Pipeline::new(vec!["writer".into()]);
992        let result = pipeline.execute(&pool, "write something").await.unwrap();
993
994        assert_eq!(result.stage_count(), 1);
995        assert_eq!(result.stages[0].role, "writer");
996        assert_eq!(result.stages[0].output, "stage one output");
997        assert_eq!(result.final_output(), "stage one output");
998    }
999
1000    #[tokio::test]
1001    async fn pipeline_multi_stage_passes_output() {
1002        let provider = Arc::new(MockProvider::new(vec![
1003            Completion {
1004                content: "draft text".into(),
1005                tool_calls: vec![],
1006                finish_reason: Some("stop".into()),
1007                usage: None,
1008                reasoning_content: None,
1009            },
1010            Completion {
1011                content: "polished text".into(),
1012                tool_calls: vec![],
1013                finish_reason: Some("stop".into()),
1014                usage: None,
1015                reasoning_content: None,
1016            },
1017        ]));
1018
1019        let mock_ref = provider.clone();
1020        let mut pool = AgentPool::new(provider, test_config());
1021        pool.add_role(AgentRole {
1022            name: "drafter".into(),
1023            system_prompt: "You draft text.".into(),
1024            max_steps: 5,
1025            allowed_tools: vec![],
1026        });
1027        pool.add_role(AgentRole {
1028            name: "editor".into(),
1029            system_prompt: "You polish text.".into(),
1030            max_steps: 5,
1031            allowed_tools: vec![],
1032        });
1033
1034        let pipeline = Pipeline::new(vec!["drafter".into(), "editor".into()]);
1035        let result = pipeline.execute(&pool, "original goal").await.unwrap();
1036
1037        assert_eq!(result.stage_count(), 2);
1038        assert_eq!(result.stages[0].output, "draft text");
1039        assert_eq!(result.stages[1].output, "polished text");
1040        assert_eq!(result.final_output(), "polished text");
1041
1042        // Verify second stage received first stage's output as its goal
1043        let calls = mock_ref.calls();
1044        assert_eq!(calls.len(), 2);
1045        // The second call's user message should contain "draft text"
1046        let second_call_user_msg = calls[1]
1047            .iter()
1048            .find(|m| m.role == crate::message::Role::User);
1049        assert!(second_call_user_msg.unwrap().content.contains("draft text"));
1050    }
1051
1052    #[tokio::test]
1053    async fn pipeline_fails_on_unknown_role() {
1054        let provider = Arc::new(MockProvider::new(vec![]));
1055        let pool = AgentPool::new(provider, test_config());
1056
1057        let pipeline = Pipeline::new(vec!["nonexistent".into()]);
1058        let result = pipeline.execute(&pool, "hello").await;
1059
1060        assert!(result.is_err());
1061        let err = result.unwrap_err();
1062        assert!(err.to_string().contains("unknown role"));
1063    }
1064
1065    #[tokio::test]
1066    async fn pipeline_final_output() {
1067        let provider = Arc::new(MockProvider::new(vec![
1068            Completion {
1069                content: "first".into(),
1070                tool_calls: vec![],
1071                finish_reason: Some("stop".into()),
1072                usage: None,
1073                reasoning_content: None,
1074            },
1075            Completion {
1076                content: "second".into(),
1077                tool_calls: vec![],
1078                finish_reason: Some("stop".into()),
1079                usage: None,
1080                reasoning_content: None,
1081            },
1082            Completion {
1083                content: "final answer".into(),
1084                tool_calls: vec![],
1085                finish_reason: Some("stop".into()),
1086                usage: None,
1087                reasoning_content: None,
1088            },
1089        ]));
1090
1091        let mut pool = AgentPool::new(provider, test_config());
1092        pool.add_role(AgentRole {
1093            name: "a".into(),
1094            system_prompt: "A".into(),
1095            max_steps: 5,
1096            allowed_tools: vec![],
1097        });
1098        pool.add_role(AgentRole {
1099            name: "b".into(),
1100            system_prompt: "B".into(),
1101            max_steps: 5,
1102            allowed_tools: vec![],
1103        });
1104        pool.add_role(AgentRole {
1105            name: "c".into(),
1106            system_prompt: "C".into(),
1107            max_steps: 5,
1108            allowed_tools: vec![],
1109        });
1110
1111        let pipeline = Pipeline::new(vec!["a".into(), "b".into(), "c".into()]);
1112        let result = pipeline.execute(&pool, "start").await.unwrap();
1113
1114        assert_eq!(result.final_output(), "final answer");
1115        assert_eq!(result.stage_count(), 3);
1116    }
1117
1118    // --- TeamOrchestrator / parse_delegations tests ---
1119
1120    #[test]
1121    fn parse_delegations_extracts_role_and_task() {
1122        let text = "DELEGATE:coder:write hello";
1123        let result = parse_delegations(text);
1124        assert_eq!(
1125            result,
1126            vec![("coder".to_string(), "write hello".to_string())]
1127        );
1128    }
1129
1130    #[test]
1131    fn parse_delegations_ignores_non_delegation() {
1132        let text = "Here is my plan:\n- Think about it\nDELEGATE:coder:implement feature\nSome other text\nDELEGATE:reviewer:check code";
1133        let result = parse_delegations(text);
1134        assert_eq!(result.len(), 2);
1135        assert_eq!(
1136            result[0],
1137            ("coder".to_string(), "implement feature".to_string())
1138        );
1139        assert_eq!(
1140            result[1],
1141            ("reviewer".to_string(), "check code".to_string())
1142        );
1143    }
1144
1145    #[test]
1146    fn parse_delegations_handles_colons_in_task() {
1147        let text = "DELEGATE:coder:write file:test.rs";
1148        let result = parse_delegations(text);
1149        assert_eq!(
1150            result,
1151            vec![("coder".to_string(), "write file:test.rs".to_string())]
1152        );
1153    }
1154
1155    #[tokio::test]
1156    async fn orchestrator_no_delegations_returns_lead_response() {
1157        // Lead responds without any DELEGATE lines
1158        let provider = Arc::new(MockProvider::new(vec![Completion {
1159            content: "I will handle this myself. The answer is 42.".into(),
1160            tool_calls: vec![],
1161            finish_reason: Some("stop".into()),
1162            usage: None,
1163            reasoning_content: None,
1164        }]));
1165
1166        let mut pool = AgentPool::new(provider, test_config());
1167        pool.add_role(AgentRole {
1168            name: "lead".into(),
1169            system_prompt: "You are the lead.".into(),
1170            max_steps: 5,
1171            allowed_tools: vec![],
1172        });
1173        pool.add_role(AgentRole {
1174            name: "coder".into(),
1175            system_prompt: "You code.".into(),
1176            max_steps: 5,
1177            allowed_tools: vec![],
1178        });
1179
1180        let orchestrator = TeamOrchestrator::new("lead".into(), vec!["coder".into()]);
1181        let result = orchestrator
1182            .run(&pool, "What is the meaning of life?")
1183            .await
1184            .unwrap();
1185
1186        assert!(result.delegations.is_empty());
1187        assert_eq!(
1188            result.final_output,
1189            "I will handle this myself. The answer is 42."
1190        );
1191    }
1192
1193    #[tokio::test]
1194    async fn orchestrator_with_delegations_executes_them() {
1195        // Completions: 1) lead delegates, 2) specialist responds, 3) lead synthesizes
1196        let provider = Arc::new(MockProvider::new(vec![
1197            Completion {
1198                content: "Let me delegate this.\nDELEGATE:coder:write a hello world program".into(),
1199                tool_calls: vec![],
1200                finish_reason: Some("stop".into()),
1201                usage: None,
1202                reasoning_content: None,
1203            },
1204            Completion {
1205                content: "fn main() { println!(\"Hello, world!\"); }".into(),
1206                tool_calls: vec![],
1207                finish_reason: Some("stop".into()),
1208                usage: None,
1209                reasoning_content: None,
1210            },
1211            Completion {
1212                content: "The coder produced a working hello world program in Rust.".into(),
1213                tool_calls: vec![],
1214                finish_reason: Some("stop".into()),
1215                usage: None,
1216                reasoning_content: None,
1217            },
1218        ]));
1219
1220        let mut pool = AgentPool::new(provider, test_config());
1221        pool.add_role(AgentRole {
1222            name: "lead".into(),
1223            system_prompt: "You are the lead.".into(),
1224            max_steps: 5,
1225            allowed_tools: vec![],
1226        });
1227        pool.add_role(AgentRole {
1228            name: "coder".into(),
1229            system_prompt: "You write code.".into(),
1230            max_steps: 5,
1231            allowed_tools: vec![],
1232        });
1233
1234        let orchestrator = TeamOrchestrator::new("lead".into(), vec!["coder".into()]);
1235        let result = orchestrator
1236            .run(&pool, "Create a hello world program")
1237            .await
1238            .unwrap();
1239
1240        assert_eq!(result.delegations.len(), 1);
1241        assert_eq!(result.delegations[0].role, "coder");
1242        assert_eq!(result.delegations[0].task, "write a hello world program");
1243        assert!(result.delegations[0].success);
1244        assert!(result.delegations[0].output.contains("Hello, world!"));
1245        assert_eq!(
1246            result.final_output,
1247            "The coder produced a working hello world program in Rust."
1248        );
1249    }
1250}
recursive/multi.rs

recursive/
multi.rs
