openhawk_core/
orchestrator.rs

1use std::collections::HashMap;
2
3use crate::error::HawkError;
4
5pub type Result<T> = std::result::Result<T, HawkError>;
6
7#[derive(Debug, Clone, PartialEq)]
8pub enum SubTaskStatus {
9    Pending,
10    Running,
11    Completed,
12    Failed(String),
13}
14
15#[derive(Debug, Clone)]
16pub struct SubTask {
17    pub description: String,
18    pub assigned_agent: Option<u32>,
19    pub status: SubTaskStatus,
20    pub required_capabilities: Vec<String>,
21}
22
23#[derive(Debug, Clone)]
24pub struct OrchestrationPlan {
25    pub task_description: String,
26    pub subtasks: Vec<SubTask>,
27    /// (dependency_idx, dependent_idx): dependency must complete before dependent
28    pub dependencies: Vec<(usize, usize)>,
29}
30
31#[derive(Debug)]
32pub struct OrchestrationReport {
33    pub plan: OrchestrationPlan,
34    pub success: bool,
35    pub summary: String,
36}
37
38#[derive(Debug, Clone)]
39pub struct AgentCapabilityRecord {
40    pub pid: u32,
41    pub name: String,
42    pub capabilities: Vec<String>,
43}
44
45#[derive(Debug)]
46pub enum OrchestratorError {
47    NoAgentsRegistered,
48    CyclicDependency,
49    Hawk(HawkError),
50}
51
52impl std::fmt::Display for OrchestratorError {
53    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
54        match self {
55            OrchestratorError::NoAgentsRegistered => write!(f, "no agents registered"),
56            OrchestratorError::CyclicDependency => write!(f, "cyclic dependency in plan"),
57            OrchestratorError::Hawk(e) => write!(f, "{e}"),
58        }
59    }
60}
61
62impl std::error::Error for OrchestratorError {}
63
64pub struct Orchestrator {
65    agents: Vec<AgentCapabilityRecord>,
66}
67
68impl Orchestrator {
69    pub fn new() -> Self {
70        Self { agents: Vec::new() }
71    }
72
73    pub fn register_agent(&mut self, pid: u32, name: impl Into<String>, capabilities: Vec<String>) {
74        self.agents.push(AgentCapabilityRecord { pid, name: name.into(), capabilities });
75    }
76
77    pub fn orchestrate(&self, task_description: &str) -> std::result::Result<OrchestrationPlan, OrchestratorError> {
78        let then_parts: Vec<&str> = task_description.split(" then ").collect();
79        let mut subtasks: Vec<SubTask> = Vec::new();
80        let mut dependencies: Vec<(usize, usize)> = Vec::new();
81
82        for (ti, then_part) in then_parts.iter().enumerate() {
83            let and_parts: Vec<&str> = then_part.split(" and ").collect();
84            let group_start = subtasks.len();
85
86            for and_part in &and_parts {
87                let trimmed = and_part.trim().to_string();
88                if trimmed.is_empty() { continue; }
89                let caps = infer_capabilities(&trimmed);
90                let assigned = best_agent(&self.agents, &caps).map(|r| r.pid);
91                subtasks.push(SubTask {
92                    description: trimmed,
93                    assigned_agent: assigned,
94                    status: SubTaskStatus::Pending,
95                    required_capabilities: caps,
96                });
97            }
98
99            // All subtasks in this group depend on all subtasks in the previous group
100            if ti > 0 {
101                let prev_group_end = group_start; // first index of current group
102                let prev_group_start = dependencies
103                    .last()
104                    .map(|&(_, dep)| dep)
105                    .unwrap_or(0);
106                // Find the range of the previous group
107                let prev_start = if ti == 1 { 0 } else { prev_group_start };
108                let prev_end = group_start;
109                for dep_idx in prev_start..prev_end {
110                    for cur_idx in group_start..subtasks.len() {
111                        dependencies.push((dep_idx, cur_idx));
112                    }
113                }
114                let _ = prev_group_end;
115            }
116        }
117
118        Ok(OrchestrationPlan { task_description: task_description.to_string(), subtasks, dependencies })
119    }
120
121    pub fn execute_plan(&self, mut plan: OrchestrationPlan) -> std::result::Result<OrchestrationReport, OrchestratorError> {
122        let order = topological_sort(plan.subtasks.len(), &plan.dependencies)
123            .ok_or(OrchestratorError::CyclicDependency)?;
124
125        let mut failed_count = 0usize;
126
127        for idx in order {
128            plan.subtasks[idx].status = SubTaskStatus::Running;
129
130            if simulate_execute(&plan.subtasks[idx]).is_ok() {
131                plan.subtasks[idx].status = SubTaskStatus::Completed;
132                continue;
133            }
134
135            // Retry once with same agent
136            if simulate_execute(&plan.subtasks[idx]).is_ok() {
137                plan.subtasks[idx].status = SubTaskStatus::Completed;
138                continue;
139            }
140
141            // Reassign to next best agent
142            let caps = plan.subtasks[idx].required_capabilities.clone();
143            let current_pid = plan.subtasks[idx].assigned_agent;
144            let next = self.agents.iter()
145                .filter(|a| Some(a.pid) != current_pid)
146                .max_by_key(|a| capability_overlap(&a.capabilities, &caps));
147
148            if let Some(agent) = next {
149                plan.subtasks[idx].assigned_agent = Some(agent.pid);
150                if simulate_execute(&plan.subtasks[idx]).is_ok() {
151                    plan.subtasks[idx].status = SubTaskStatus::Completed;
152                    continue;
153                }
154            }
155
156            plan.subtasks[idx].status = SubTaskStatus::Failed("no agent could complete task".to_string());
157            failed_count += 1;
158        }
159
160        let total = plan.subtasks.len();
161        let completed = total - failed_count;
162        let success = failed_count == 0;
163        let summary = if success {
164            format!("All {total} sub-tasks completed successfully.")
165        } else {
166            format!("{completed}/{total} sub-tasks completed; {failed_count} failed.")
167        };
168
169        Ok(OrchestrationReport { plan, success, summary })
170    }
171}
172
173impl Default for Orchestrator {
174    fn default() -> Self {
175        Self::new()
176    }
177}
178
179#[allow(dead_code)]
180fn split_task(desc: &str) -> Vec<(String, bool)> {
181    let then_parts: Vec<&str> = desc.split(" then ").collect();
182    let mut result = Vec::new();
183
184    for (ti, then_part) in then_parts.iter().enumerate() {
185        let and_parts: Vec<&str> = then_part.split(" and ").collect();
186        for (ai, and_part) in and_parts.iter().enumerate() {
187            let trimmed = and_part.trim().to_string();
188            if trimmed.is_empty() { continue; }
189            let is_sequential = ti > 0 && ai == 0;
190            result.push((trimmed, is_sequential));
191        }
192    }
193
194    if result.is_empty() {
195        result.push((desc.trim().to_string(), false));
196    }
197    result
198}
199
200fn infer_capabilities(desc: &str) -> Vec<String> {
201    let lower = desc.to_lowercase();
202    let mut caps = Vec::new();
203    let keywords: &[(&str, &str)] = &[
204        ("research", "research"),
205        ("search", "research"),
206        ("summar", "summarization"),
207        ("code", "coding"),
208        ("implement", "coding"),
209        ("write", "coding"),
210        ("review", "review"),
211        ("test", "testing"),
212        ("deploy", "deployment"),
213        ("analyz", "analysis"),
214        ("analys", "analysis"),
215        ("web", "web-search"),
216    ];
217    for (kw, cap) in keywords {
218        if lower.contains(kw) && !caps.contains(&cap.to_string()) {
219            caps.push(cap.to_string());
220        }
221    }
222    caps
223}
224
225pub fn capability_overlap(agent_caps: &[String], required: &[String]) -> usize {
226    required.iter().filter(|r| agent_caps.contains(r)).count()
227}
228
229fn best_agent<'a>(agents: &'a [AgentCapabilityRecord], required: &[String]) -> Option<&'a AgentCapabilityRecord> {
230    agents.iter().max_by_key(|a| capability_overlap(&a.capabilities, required))
231}
232
233pub fn topological_sort(n: usize, deps: &[(usize, usize)]) -> Option<Vec<usize>> {
234    let mut in_degree = vec![0usize; n];
235    let mut adj: HashMap<usize, Vec<usize>> = HashMap::new();
236
237    for &(dep, dependent) in deps {
238        if dep >= n || dependent >= n { return None; }
239        in_degree[dependent] += 1;
240        adj.entry(dep).or_default().push(dependent);
241    }
242
243    let mut queue: Vec<usize> = (0..n).filter(|&i| in_degree[i] == 0).collect();
244    let mut order = Vec::with_capacity(n);
245
246    while let Some(node) = queue.first().copied() {
247        queue.remove(0);
248        order.push(node);
249        if let Some(neighbors) = adj.get(&node) {
250            for &nb in neighbors {
251                in_degree[nb] -= 1;
252                if in_degree[nb] == 0 {
253                    queue.push(nb);
254                }
255            }
256        }
257    }
258
259    if order.len() == n { Some(order) } else { None }
260}
261
262fn simulate_execute(subtask: &SubTask) -> std::result::Result<(), String> {
263    if subtask.assigned_agent.is_none() {
264        return Err("no agent assigned".to_string());
265    }
266    Ok(())
267}
268
269#[cfg(test)]
270mod tests {
271    use super::*;
272
273    fn make_orchestrator() -> Orchestrator {
274        let mut o = Orchestrator::new();
275        o.register_agent(1, "research-agent", vec!["research".into(), "web-search".into()]);
276        o.register_agent(2, "coding-agent", vec!["coding".into(), "testing".into()]);
277        o.register_agent(3, "review-agent", vec!["review".into(), "analysis".into()]);
278        o
279    }
280
281    #[test]
282    fn single_task_produces_one_subtask() {
283        let o = make_orchestrator();
284        let plan = o.orchestrate("research quantum computing").unwrap();
285        assert_eq!(plan.subtasks.len(), 1);
286        assert!(plan.dependencies.is_empty());
287    }
288
289    #[test]
290    fn and_produces_parallel_subtasks_no_dependencies() {
291        let o = make_orchestrator();
292        let plan = o.orchestrate("research topic and write code").unwrap();
293        assert_eq!(plan.subtasks.len(), 2);
294        assert!(plan.dependencies.is_empty());
295    }
296
297    #[test]
298    fn then_produces_sequential_dependency() {
299        let o = make_orchestrator();
300        let plan = o.orchestrate("research topic then write code").unwrap();
301        assert_eq!(plan.subtasks.len(), 2);
302        assert_eq!(plan.dependencies.len(), 1);
303        assert_eq!(plan.dependencies[0], (0, 1));
304    }
305
306    #[test]
307    fn subtasks_have_non_empty_descriptions() {
308        let o = make_orchestrator();
309        let plan = o.orchestrate("research topic then write code then review changes").unwrap();
310        for st in &plan.subtasks {
311            assert!(!st.description.is_empty());
312        }
313    }
314
315    #[test]
316    fn plan_preserves_task_description() {
317        let o = make_orchestrator();
318        let desc = "research topic and write code";
319        let plan = o.orchestrate(desc).unwrap();
320        assert_eq!(plan.task_description, desc);
321    }
322
323    #[test]
324    fn research_task_assigned_to_research_agent() {
325        let o = make_orchestrator();
326        let plan = o.orchestrate("research quantum computing").unwrap();
327        assert_eq!(plan.subtasks[0].assigned_agent, Some(1));
328    }
329
330    #[test]
331    fn coding_task_assigned_to_coding_agent() {
332        let o = make_orchestrator();
333        let plan = o.orchestrate("implement the algorithm").unwrap();
334        assert_eq!(plan.subtasks[0].assigned_agent, Some(2));
335    }
336
337    #[test]
338    fn review_task_assigned_to_review_agent() {
339        let o = make_orchestrator();
340        let plan = o.orchestrate("review the changes").unwrap();
341        assert_eq!(plan.subtasks[0].assigned_agent, Some(3));
342    }
343
344    #[test]
345    fn all_subtasks_get_agent_assigned_when_agents_available() {
346        let o = make_orchestrator();
347        let plan = o.orchestrate("research topic and write code and review changes").unwrap();
348        for st in &plan.subtasks {
349            assert!(st.assigned_agent.is_some());
350        }
351    }
352
353    #[test]
354    fn no_agents_still_produces_plan_with_none_assigned() {
355        let o = Orchestrator::new();
356        let plan = o.orchestrate("research topic").unwrap();
357        assert_eq!(plan.subtasks[0].assigned_agent, None);
358    }
359
360    #[test]
361    fn independent_subtasks_all_complete() {
362        let o = make_orchestrator();
363        let plan = o.orchestrate("research topic and write code").unwrap();
364        let report = o.execute_plan(plan).unwrap();
365        assert!(report.success);
366        for st in &report.plan.subtasks {
367            assert_eq!(st.status, SubTaskStatus::Completed);
368        }
369    }
370
371    #[test]
372    fn sequential_subtasks_complete_in_order() {
373        let o = make_orchestrator();
374        let plan = o.orchestrate("research topic then write code").unwrap();
375        let report = o.execute_plan(plan).unwrap();
376        assert!(report.success);
377        assert_eq!(report.plan.subtasks[0].status, SubTaskStatus::Completed);
378        assert_eq!(report.plan.subtasks[1].status, SubTaskStatus::Completed);
379    }
380
381    #[test]
382    fn subtask_with_no_agent_fails_gracefully() {
383        let o = Orchestrator::new();
384        let plan = o.orchestrate("research topic").unwrap();
385        let report = o.execute_plan(plan).unwrap();
386        assert!(!report.success);
387        assert!(matches!(report.plan.subtasks[0].status, SubTaskStatus::Failed(_)));
388    }
389
390    #[test]
391    fn report_summary_reflects_failure_count() {
392        let o = Orchestrator::new();
393        let plan = o.orchestrate("research topic and write code").unwrap();
394        let report = o.execute_plan(plan).unwrap();
395        assert!(report.summary.contains("failed") || report.summary.contains("0/"));
396    }
397
398    #[test]
399    fn reassignment_uses_best_matching_agent() {
400        let mut o = Orchestrator::new();
401        o.register_agent(10, "generic-agent", vec!["generic".into()]);
402        o.register_agent(11, "research-agent", vec!["research".into()]);
403        let plan = o.orchestrate("research quantum computing").unwrap();
404        assert_eq!(plan.subtasks[0].assigned_agent, Some(11));
405    }
406
407    #[test]
408    fn successful_report_has_correct_summary() {
409        let o = make_orchestrator();
410        let plan = o.orchestrate("research topic").unwrap();
411        let report = o.execute_plan(plan).unwrap();
412        assert!(report.success);
413        assert!(report.summary.contains("1"));
414        assert!(report.summary.contains("completed"));
415    }
416
417    #[test]
418    fn report_contains_original_plan() {
419        let o = make_orchestrator();
420        let plan = o.orchestrate("research topic and write code").unwrap();
421        let report = o.execute_plan(plan).unwrap();
422        assert_eq!(report.plan.subtasks.len(), 2);
423    }
424
425    #[test]
426    fn topo_sort_no_deps_returns_all_nodes() {
427        let order = topological_sort(3, &[]).unwrap();
428        assert_eq!(order.len(), 3);
429    }
430
431    #[test]
432    fn topo_sort_linear_chain() {
433        let order = topological_sort(3, &[(0, 1), (1, 2)]).unwrap();
434        assert_eq!(order, vec![0, 1, 2]);
435    }
436
437    #[test]
438    fn topo_sort_cycle_returns_none() {
439        let result = topological_sort(2, &[(0, 1), (1, 0)]);
440        assert!(result.is_none());
441    }
442
443    #[test]
444    fn capability_overlap_full_match() {
445        let agent = vec!["research".into(), "web-search".into()];
446        let required = vec!["research".into(), "web-search".into()];
447        assert_eq!(capability_overlap(&agent, &required), 2);
448    }
449
450    #[test]
451    fn capability_overlap_no_match() {
452        let agent = vec!["coding".into()];
453        let required = vec!["research".into()];
454        assert_eq!(capability_overlap(&agent, &required), 0);
455    }
456
457    #[test]
458    fn capability_overlap_partial_match() {
459        let agent = vec!["research".into(), "coding".into()];
460        let required = vec!["research".into(), "web-search".into()];
461        assert_eq!(capability_overlap(&agent, &required), 1);
462    }
463}
openhawk_core/orchestrator.rs

openhawk_core/
orchestrator.rs
