1use serde::{Deserialize, Serialize};
9
10use super::events::AgentEvent;
11use crate::error::Error;
12
13pub trait RoutingStrategy: Send + Sync {
38 fn route(&self, task: &str, agents: &[AgentCapability])
40 -> (RoutingDecision, ComplexitySignals);
41}
42
43pub struct KeywordRoutingStrategy;
50
51impl RoutingStrategy for KeywordRoutingStrategy {
52 fn route(
53 &self,
54 task: &str,
55 agents: &[AgentCapability],
56 ) -> (RoutingDecision, ComplexitySignals) {
57 let analyzer = TaskComplexityAnalyzer::new(agents);
58 analyzer.analyze(task)
59 }
60}
61
62#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Deserialize, Serialize)]
64#[serde(rename_all = "snake_case")]
65pub enum RoutingMode {
66 #[default]
68 Auto,
69 AlwaysOrchestrate,
71 SingleAgent,
73}
74
75#[derive(Debug, Clone, PartialEq)]
77pub enum RoutingDecision {
78 SingleAgent {
80 agent_index: usize,
82 reason: &'static str,
84 },
85 Orchestrate {
87 reason: &'static str,
89 },
90}
91
92#[derive(Debug, Clone, Default)]
94pub struct ComplexitySignals {
95 pub word_count: usize,
97 pub step_markers: usize,
99 pub domain_signals: Vec<String>,
101 pub explicit_delegation: bool,
103 pub names_multiple_agents: bool,
105 pub covering_agents: Vec<usize>,
107 pub complexity_score: f32,
109}
110
111#[derive(Debug, Clone)]
113pub struct AgentCapability {
114 pub name: String,
116 pub description_lower: String,
118 pub tool_names: Vec<String>,
120 pub domains: Vec<String>,
122}
123
124const CODE_KEYWORDS: &[&str] = &[
127 "code",
128 "implement",
129 "refactor",
130 "debug",
131 "compile",
132 "function",
133 "class",
134 "module",
135 "rust",
136 "python",
137 "javascript",
138 "typescript",
139 "java",
140 "golang",
141 "programming",
142 "syntax",
143 "bug",
144 "fix",
145 "test",
146 "unit test",
147];
148
149const RESEARCH_KEYWORDS: &[&str] = &[
150 "research",
151 "investigate",
152 "analyze",
153 "study",
154 "survey",
155 "find",
156 "search",
157 "explore",
158 "review",
159 "literature",
160 "paper",
161 "arxiv",
162];
163
164const DATABASE_KEYWORDS: &[&str] = &[
165 "database",
166 "sql",
167 "query",
168 "table",
169 "schema",
170 "migration",
171 "postgres",
172 "mysql",
173 "sqlite",
174 "mongodb",
175 "redis",
176 "index",
177];
178
179const FRONTEND_KEYWORDS: &[&str] = &[
180 "frontend",
181 "ui",
182 "ux",
183 "component",
184 "react",
185 "vue",
186 "angular",
187 "css",
188 "html",
189 "layout",
190 "responsive",
191 "design",
192 "button",
193 "form",
194 "modal",
195];
196
197const BACKEND_KEYWORDS: &[&str] = &[
198 "backend",
199 "api",
200 "endpoint",
201 "server",
202 "rest",
203 "graphql",
204 "middleware",
205 "authentication",
206 "authorization",
207 "route",
208 "handler",
209];
210
211const DEVOPS_KEYWORDS: &[&str] = &[
212 "devops",
213 "deploy",
214 "docker",
215 "kubernetes",
216 "ci/cd",
217 "pipeline",
218 "infrastructure",
219 "terraform",
220 "ansible",
221 "nginx",
222 "monitoring",
223 "logging",
224];
225
226const WRITING_KEYWORDS: &[&str] = &[
227 "write",
228 "document",
229 "documentation",
230 "readme",
231 "blog",
232 "article",
233 "report",
234 "summary",
235 "copywriting",
236 "content",
237 "draft",
238 "edit text",
239];
240
241const SECURITY_KEYWORDS: &[&str] = &[
242 "security",
243 "vulnerability",
244 "audit",
245 "penetration",
246 "encryption",
247 "auth",
248 "cors",
249 "xss",
250 "injection",
251 "firewall",
252 "certificate",
253 "tls",
254];
255
256const DOMAIN_LISTS: &[(&str, &[&str])] = &[
257 ("code", CODE_KEYWORDS),
258 ("research", RESEARCH_KEYWORDS),
259 ("database", DATABASE_KEYWORDS),
260 ("frontend", FRONTEND_KEYWORDS),
261 ("backend", BACKEND_KEYWORDS),
262 ("devops", DEVOPS_KEYWORDS),
263 ("writing", WRITING_KEYWORDS),
264 ("security", SECURITY_KEYWORDS),
265];
266
267const STEP_MARKERS: &[&str] = &[
270 "first,",
271 "second,",
272 "third,",
273 "then,",
274 "finally,",
275 "next,",
276 "after that",
277 "step 1",
278 "step 2",
279 "step 3",
280 "step 4",
281 "step 5",
282];
283
284const DELEGATION_PHRASES: &[&str] = &[
287 "delegate",
288 "have them",
289 "coordinate between",
290 "coordinate with",
291 "team up",
292 "work together",
293 "collaborate",
294 "assign to",
295 "hand off",
296 "pass to",
297 "ask the ",
303 "have the ",
304 "use the ",
305 "tell the ",
306 "instruct the ",
307 "let the ",
308 "get the ",
309];
310
311const SINGLE_AGENT_THRESHOLD: f32 = 0.30;
314const ORCHESTRATE_THRESHOLD: f32 = 0.55;
315
316const WEIGHT_SIMPLE_QUESTION: f32 = -0.30;
319const WEIGHT_WORD_COUNT_HIGH: f32 = 0.10;
320const WEIGHT_STEP_MARKERS: f32 = 0.25;
321const WEIGHT_DELEGATION: f32 = 0.30;
322const WEIGHT_NAMES_AGENTS: f32 = 0.40;
323const WEIGHT_DOMAIN_DIVERSITY: f32 = 0.20;
324
325impl AgentCapability {
326 pub fn from_config(name: &str, description: &str, tool_names: &[String]) -> Self {
328 let description_lower = description.to_lowercase();
329 let tool_lower: Vec<String> = tool_names.iter().map(|t| t.to_lowercase()).collect();
330
331 let combined = format!("{} {}", description_lower, tool_lower.join(" "));
333 let mut domains = Vec::new();
334 for &(domain, keywords) in DOMAIN_LISTS {
335 if keywords.iter().any(|kw| contains_keyword(&combined, kw)) {
336 domains.push(domain.to_string());
337 }
338 }
339
340 Self {
341 name: name.to_lowercase(),
342 description_lower,
343 tool_names: tool_lower,
344 domains,
345 }
346 }
347}
348
349pub struct TaskComplexityAnalyzer<'a> {
351 agents: &'a [AgentCapability],
352}
353
354impl<'a> TaskComplexityAnalyzer<'a> {
355 pub fn new(agents: &'a [AgentCapability]) -> Self {
357 Self { agents }
358 }
359
360 pub fn analyze(&self, task: &str) -> (RoutingDecision, ComplexitySignals) {
362 let mut signals = self.heuristic_signals(task);
363
364 if signals.complexity_score < SINGLE_AGENT_THRESHOLD {
372 let agent_index = if signals.domain_signals.is_empty() {
373 0
374 } else {
375 best_covering_agent(&signals.domain_signals, self.agents).unwrap_or(0)
376 };
377 let reason = if agent_index == 0 && signals.domain_signals.is_empty() {
378 "heuristic score below single-agent threshold (no domain signals)"
379 } else if agent_index == 0 {
380 "heuristic score below single-agent threshold (no agent matched detected domains)"
381 } else {
382 "heuristic score below single-agent threshold (matched by domain coverage)"
383 };
384 return (
385 RoutingDecision::SingleAgent {
386 agent_index,
387 reason,
388 },
389 signals,
390 );
391 }
392 if signals.complexity_score > ORCHESTRATE_THRESHOLD {
393 return (
394 RoutingDecision::Orchestrate {
395 reason: "heuristic score above orchestrate threshold",
396 },
397 signals,
398 );
399 }
400
401 let decision = self.capability_match(&signals.domain_signals, &mut signals.covering_agents);
403 (decision, signals)
404 }
405
406 pub fn heuristic_signals(&self, task: &str) -> ComplexitySignals {
408 let task_lower = task.to_lowercase();
409 let words: Vec<&str> = task.split_whitespace().collect();
410 let word_count = words.len();
411
412 let simple_question = is_simple_question(&task_lower, &words);
414
415 let step_markers = count_step_markers(&task_lower);
417 let numbered_items = words.iter().filter(|w| is_numbered_step_marker(w)).count();
421 let total_step_markers = step_markers + numbered_items;
422
423 let explicit_delegation = DELEGATION_PHRASES.iter().any(|p| task_lower.contains(p));
425
426 let domain_signals = detect_domains(&task_lower);
428
429 let names_multiple_agents = if self.agents.len() >= 2 {
431 let matching = self
432 .agents
433 .iter()
434 .filter(|a| task_lower.contains(&a.name))
435 .count();
436 matching >= 2
437 } else {
438 false
439 };
440
441 let mut score: f32 = 0.0;
443 if simple_question {
444 score += WEIGHT_SIMPLE_QUESTION;
445 }
446 if word_count > 100 {
447 score += WEIGHT_WORD_COUNT_HIGH;
448 }
449 if total_step_markers >= 2 {
450 score += WEIGHT_STEP_MARKERS;
451 }
452 if explicit_delegation {
453 score += WEIGHT_DELEGATION;
454 }
455 if names_multiple_agents {
456 score += WEIGHT_NAMES_AGENTS;
457 }
458 if domain_signals.len() >= 3 {
459 score += WEIGHT_DOMAIN_DIVERSITY;
460 }
461
462 score = score.clamp(0.0, 1.0);
464
465 ComplexitySignals {
466 word_count,
467 step_markers: total_step_markers,
468 domain_signals,
469 explicit_delegation,
470 names_multiple_agents,
471 covering_agents: Vec::new(),
472 complexity_score: score,
473 }
474 }
475
476 fn capability_match(
478 &self,
479 task_domains: &[String],
480 covering_agents: &mut Vec<usize>,
481 ) -> RoutingDecision {
482 if task_domains.is_empty() {
483 return RoutingDecision::SingleAgent {
485 agent_index: 0,
486 reason: "no domains detected, defaulting to single agent",
487 };
488 }
489
490 for (i, agent) in self.agents.iter().enumerate() {
491 let covers_all = task_domains.iter().all(|d| agent.domains.contains(d));
492 if covers_all {
493 covering_agents.push(i);
494 }
495 }
496
497 if covering_agents.is_empty() {
498 RoutingDecision::Orchestrate {
499 reason: "no single agent covers all detected domains",
500 }
501 } else {
502 RoutingDecision::SingleAgent {
504 agent_index: covering_agents[0],
505 reason: "single agent covers all detected domains",
506 }
507 }
508 }
509}
510
511fn is_simple_question(task_lower: &str, words: &[&str]) -> bool {
513 let question_starters = [
514 "what", "how", "why", "explain", "describe", "who", "when", "where",
515 ];
516 let starts_with_question = words
520 .first()
521 .is_some_and(|w| question_starters.iter().any(|q| w.starts_with(q)));
522 let has_step_markers = count_step_markers(task_lower) >= 2;
523 starts_with_question && !has_step_markers
524}
525
526fn count_step_markers(task_lower: &str) -> usize {
528 STEP_MARKERS
529 .iter()
530 .filter(|marker| task_lower.contains(*marker))
531 .count()
532}
533
534fn is_numbered_step_marker(word: &str) -> bool {
539 let trimmed = word.trim_end_matches([';', ',', ':']);
540 if trimmed.len() < 2 {
541 return false;
542 }
543 if let Some(prefix) = trimmed.strip_suffix('.')
545 && !prefix.is_empty()
546 && prefix.chars().all(|c| c.is_ascii_digit())
547 {
548 return true;
549 }
550 if let Some(inner) = trimmed.strip_prefix('(').and_then(|s| s.strip_suffix(')'))
552 && !inner.is_empty()
553 && inner.chars().all(|c| c.is_ascii_digit())
554 {
555 return true;
556 }
557 if let Some(prefix) = trimmed.strip_suffix(')')
559 && !prefix.is_empty()
560 && prefix.chars().all(|c| c.is_ascii_digit())
561 {
562 return true;
563 }
564 false
565}
566
567fn contains_keyword(text: &str, keyword: &str) -> bool {
570 if keyword.contains(' ') {
571 return text.contains(keyword);
573 }
574 for (start, _) in text.match_indices(keyword) {
576 let end = start + keyword.len();
577 let before_ok = start == 0 || !text.as_bytes()[start - 1].is_ascii_alphanumeric();
578 let after_ok = end == text.len() || !text.as_bytes()[end].is_ascii_alphanumeric();
579 if before_ok && after_ok {
580 return true;
581 }
582 }
583 false
584}
585
586fn best_covering_agent(task_domains: &[String], agents: &[AgentCapability]) -> Option<usize> {
597 if task_domains.is_empty() || agents.is_empty() {
598 return None;
599 }
600 let mut best: Option<(usize, usize)> = None;
601 for (i, agent) in agents.iter().enumerate() {
602 let count = task_domains
603 .iter()
604 .filter(|d| agent.domains.contains(d))
605 .count();
606 if count == 0 {
607 continue;
608 }
609 match best {
610 Some((_, c)) if c >= count => {}
611 _ => best = Some((i, count)),
612 }
613 }
614 best.map(|(i, _)| i)
615}
616
617fn detect_domains(task_lower: &str) -> Vec<String> {
619 let mut domains = Vec::new();
620 for &(domain, keywords) in DOMAIN_LISTS {
621 if keywords.iter().any(|kw| contains_keyword(task_lower, kw)) {
622 domains.push(domain.to_string());
623 }
624 }
625 domains
626}
627
628pub fn should_escalate(error: &Error, events: &[AgentEvent]) -> bool {
630 let inner = match error {
632 Error::WithPartialUsage { source, .. } => source.as_ref(),
633 other => other,
634 };
635
636 if matches!(inner, Error::MaxTurnsExceeded(_) | Error::RunTimeout(_)) {
638 return true;
639 }
640
641 if events
643 .iter()
644 .any(|e| matches!(e, AgentEvent::DoomLoopDetected { .. }))
645 {
646 return true;
647 }
648
649 let compaction_count = events
651 .iter()
652 .filter(|e| matches!(e, AgentEvent::AutoCompactionTriggered { .. }))
653 .count();
654 if compaction_count >= 2 {
655 return true;
656 }
657
658 false
659}
660
661pub fn resolve_routing_mode(config_mode: RoutingMode) -> RoutingMode {
663 match std::env::var("HEARTBIT_ROUTING").ok() {
664 Some(val) => match val.to_lowercase().as_str() {
665 "auto" => RoutingMode::Auto,
666 "always_orchestrate" => RoutingMode::AlwaysOrchestrate,
667 "single_agent" => RoutingMode::SingleAgent,
668 _ => {
669 tracing::warn!(
670 value = %val,
671 "unknown HEARTBIT_ROUTING value, falling back to config"
672 );
673 config_mode
674 }
675 },
676 None => config_mode,
677 }
678}
679
680#[cfg(test)]
681mod tests {
682 use super::*;
683
684 fn make_agents() -> Vec<AgentCapability> {
685 vec![
686 AgentCapability::from_config(
687 "coder",
688 "A code implementation agent that writes and debugs software",
689 &["bash".into(), "read_file".into(), "write_file".into()],
690 ),
691 AgentCapability::from_config(
692 "researcher",
693 "A research agent that investigates and analyzes topics",
694 &["web_search".into(), "read_file".into()],
695 ),
696 ]
697 }
698
699 #[test]
702 fn agent_capability_extracts_domains_from_description() {
703 let cap = AgentCapability::from_config(
704 "fullstack",
705 "Handles frontend React components and backend API endpoints with database queries",
706 &[],
707 );
708 assert!(cap.domains.contains(&"frontend".to_string()));
709 assert!(cap.domains.contains(&"backend".to_string()));
710 assert!(cap.domains.contains(&"database".to_string()));
711 }
712
713 #[test]
714 fn agent_capability_extracts_domains_from_tools() {
715 let cap = AgentCapability::from_config(
716 "devops-agent",
717 "Manages infrastructure",
718 &["docker_build".into(), "deploy_k8s".into()],
719 );
720 assert!(cap.domains.contains(&"devops".to_string()));
721 }
722
723 #[test]
726 fn simple_question_scores_below_threshold() {
727 let agents = make_agents();
728 let analyzer = TaskComplexityAnalyzer::new(&agents);
729 let (decision, signals) = analyzer.analyze("What is the capital of France?");
730 assert!(
731 signals.complexity_score < SINGLE_AGENT_THRESHOLD,
732 "score {} should be < {}",
733 signals.complexity_score,
734 SINGLE_AGENT_THRESHOLD
735 );
736 assert!(matches!(decision, RoutingDecision::SingleAgent { .. }));
737 }
738
739 #[test]
740 fn multi_step_multi_domain_routes_to_orchestrate() {
741 let agents = make_agents();
742 let analyzer = TaskComplexityAnalyzer::new(&agents);
743 let task = "First, research the best database schema for user authentication. \
744 Then, implement the backend API endpoints in Rust. \
745 Finally, write frontend React components for the login form.";
746 let (decision, signals) = analyzer.analyze(task);
747 assert!(
748 signals.step_markers >= 2,
749 "step_markers: {}",
750 signals.step_markers
751 );
752 assert!(
753 signals.domain_signals.len() >= 3,
754 "domains: {:?}",
755 signals.domain_signals
756 );
757 assert!(
760 matches!(decision, RoutingDecision::Orchestrate { .. }),
761 "decision: {decision:?}, score: {}",
762 signals.complexity_score
763 );
764 }
765
766 #[test]
767 fn delegation_language_boosts_score() {
768 let agents = make_agents();
769 let analyzer = TaskComplexityAnalyzer::new(&agents);
770 let task = "Delegate the research task to the researcher and coordinate with the coder";
771 let signals = analyzer.heuristic_signals(task);
772 assert!(signals.explicit_delegation);
773 assert!(
775 signals.complexity_score > ORCHESTRATE_THRESHOLD,
776 "score: {}",
777 signals.complexity_score
778 );
779 }
780
781 #[test]
782 fn naming_multiple_agents_boosts_score() {
783 let agents = make_agents();
784 let analyzer = TaskComplexityAnalyzer::new(&agents);
785 let task = "Have coder write the code and researcher find the documentation";
786 let signals = analyzer.heuristic_signals(task);
787 assert!(signals.names_multiple_agents);
788 assert!(
789 signals.complexity_score >= WEIGHT_NAMES_AGENTS,
790 "score: {}",
791 signals.complexity_score
792 );
793 }
794
795 #[test]
796 fn word_count_above_100_adds_weight() {
797 let agents = make_agents();
798 let analyzer = TaskComplexityAnalyzer::new(&agents);
799 let task = (0..110)
801 .map(|i| format!("word{i}"))
802 .collect::<Vec<_>>()
803 .join(" ");
804 let signals = analyzer.heuristic_signals(&task);
805 assert!(signals.word_count > 100);
806 assert!(
807 signals.complexity_score >= WEIGHT_WORD_COUNT_HIGH,
808 "score: {}",
809 signals.complexity_score
810 );
811 }
812
813 #[test]
814 fn numbered_list_detected_as_step_markers() {
815 let agents = make_agents();
816 let analyzer = TaskComplexityAnalyzer::new(&agents);
817 let task = "1. Set up the database. 2. Write the API. 3. Test everything.";
818 let signals = analyzer.heuristic_signals(task);
819 assert!(
820 signals.step_markers >= 2,
821 "step_markers: {}",
822 signals.step_markers
823 );
824 }
825
826 #[test]
827 fn score_clamped_to_zero_one() {
828 let agents = make_agents();
829 let analyzer = TaskComplexityAnalyzer::new(&agents);
830
831 let signals = analyzer.heuristic_signals("What is 2+2?");
833 assert!(
834 signals.complexity_score >= 0.0,
835 "score: {}",
836 signals.complexity_score
837 );
838
839 let task = "First, delegate to coder and researcher. Then step 1 deploy the docker \
841 kubernetes infrastructure with database schema, frontend React components, \
842 backend API endpoints, security audit, research papers, and write documentation. \
843 Finally, coordinate the team. ".repeat(5);
844 let signals = analyzer.heuristic_signals(&task);
845 assert!(
846 signals.complexity_score <= 1.0,
847 "score: {}",
848 signals.complexity_score
849 );
850 }
851
852 #[test]
855 fn one_agent_covers_all_domains_routes_to_single() {
856 let agents = vec![AgentCapability::from_config(
857 "fullstack",
858 "Handles code implementation, database queries, and backend API endpoints",
859 &[],
860 )];
861 let analyzer = TaskComplexityAnalyzer::new(&agents);
862 let task = "Update the database query and fix the backend API endpoint bug";
864 let (decision, signals) = analyzer.analyze(task);
865 match &decision {
867 RoutingDecision::SingleAgent { agent_index, .. } => {
868 assert_eq!(*agent_index, 0);
869 }
870 RoutingDecision::Orchestrate { reason } => {
871 assert!(
873 signals.complexity_score > ORCHESTRATE_THRESHOLD,
874 "unexpected orchestrate: {reason}"
875 );
876 }
877 }
878 }
879
880 #[test]
881 fn split_coverage_routes_to_orchestrate() {
882 let agents = vec![
883 AgentCapability::from_config(
884 "frontend-dev",
885 "Builds frontend React components and CSS layouts",
886 &[],
887 ),
888 AgentCapability::from_config(
889 "backend-dev",
890 "Builds backend API endpoints and database schemas",
891 &[],
892 ),
893 ];
894 let analyzer = TaskComplexityAnalyzer::new(&agents);
895 let task = "Build a React form that submits to a new backend API endpoint and stores data in the database";
897 let mut signals = analyzer.heuristic_signals(task);
898 let mut covering = Vec::new();
899 let decision = analyzer.capability_match(&signals.domain_signals, &mut covering);
900 signals.covering_agents = covering;
901 assert!(
902 matches!(decision, RoutingDecision::Orchestrate { .. }),
903 "expected Orchestrate, got: {decision:?}"
904 );
905 assert!(signals.covering_agents.is_empty());
906 }
907
908 #[test]
909 fn no_domains_defaults_to_single_agent() {
910 let agents = make_agents();
911 let analyzer = TaskComplexityAnalyzer::new(&agents);
912 let mut covering = Vec::new();
913 let decision = analyzer.capability_match(&[], &mut covering);
914 assert!(matches!(
915 decision,
916 RoutingDecision::SingleAgent { agent_index: 0, .. }
917 ));
918 }
919
920 #[test]
923 fn escalate_on_max_turns_exceeded() {
924 let err = Error::MaxTurnsExceeded(10);
925 assert!(should_escalate(&err, &[]));
926 }
927
928 #[test]
929 fn escalate_on_max_turns_wrapped_in_partial_usage() {
930 use crate::llm::types::TokenUsage;
931 let err = Error::MaxTurnsExceeded(10).with_partial_usage(TokenUsage::default());
932 assert!(should_escalate(&err, &[]));
933 }
934
935 #[test]
936 fn escalate_on_run_timeout() {
937 let err = Error::RunTimeout(std::time::Duration::from_secs(60));
938 assert!(should_escalate(&err, &[]));
939 }
940
941 #[test]
942 fn escalate_on_doom_loop_event() {
943 let events = vec![AgentEvent::DoomLoopDetected {
944 agent: "a".into(),
945 turn: 5,
946 consecutive_count: 3,
947 tool_names: vec!["web_search".into()],
948 }];
949 let err = Error::Agent("generic error".into());
950 assert!(should_escalate(&err, &events));
951 }
952
953 #[test]
954 fn escalate_on_two_compactions() {
955 let events = vec![
956 AgentEvent::AutoCompactionTriggered {
957 agent: "a".into(),
958 turn: 2,
959 success: true,
960 usage: Default::default(),
961 },
962 AgentEvent::AutoCompactionTriggered {
963 agent: "a".into(),
964 turn: 5,
965 success: true,
966 usage: Default::default(),
967 },
968 ];
969 let err = Error::Agent("context overflow".into());
970 assert!(should_escalate(&err, &events));
971 }
972
973 #[test]
974 fn no_escalation_on_single_compaction() {
975 let events = vec![AgentEvent::AutoCompactionTriggered {
976 agent: "a".into(),
977 turn: 2,
978 success: true,
979 usage: Default::default(),
980 }];
981 let err = Error::Agent("some error".into());
982 assert!(!should_escalate(&err, &events));
983 }
984
985 #[test]
986 fn no_escalation_on_normal_error() {
987 let err = Error::Agent("tool failed".into());
988 assert!(!should_escalate(&err, &[]));
989 }
990
991 #[test]
994 fn routing_mode_default_is_auto() {
995 assert_eq!(RoutingMode::default(), RoutingMode::Auto);
996 }
997
998 #[test]
999 fn routing_mode_roundtrips_json() {
1000 for mode in [
1001 RoutingMode::Auto,
1002 RoutingMode::AlwaysOrchestrate,
1003 RoutingMode::SingleAgent,
1004 ] {
1005 let json = serde_json::to_string(&mode).unwrap();
1006 let back: RoutingMode = serde_json::from_str(&json).unwrap();
1007 assert_eq!(mode, back, "failed for {json}");
1008 }
1009 }
1010
1011 #[test]
1012 fn routing_mode_deserializes_from_toml_strings() {
1013 #[derive(Deserialize)]
1014 struct W {
1015 mode: RoutingMode,
1016 }
1017 let w: W = toml::from_str(r#"mode = "auto""#).unwrap();
1018 assert_eq!(w.mode, RoutingMode::Auto);
1019 let w: W = toml::from_str(r#"mode = "always_orchestrate""#).unwrap();
1020 assert_eq!(w.mode, RoutingMode::AlwaysOrchestrate);
1021 let w: W = toml::from_str(r#"mode = "single_agent""#).unwrap();
1022 assert_eq!(w.mode, RoutingMode::SingleAgent);
1023 }
1024
1025 #[test]
1028 fn analyze_simple_task_two_agents_routes_single() {
1029 let agents = make_agents();
1030 let analyzer = TaskComplexityAnalyzer::new(&agents);
1031 let (decision, _) = analyzer.analyze("How do I parse JSON in Rust?");
1032 assert!(
1033 matches!(decision, RoutingDecision::SingleAgent { .. }),
1034 "got: {decision:?}"
1035 );
1036 }
1037
1038 #[test]
1039 fn analyze_complex_multi_domain_task_routes_orchestrate() {
1040 let agents = make_agents();
1041 let analyzer = TaskComplexityAnalyzer::new(&agents);
1042 let task = "First, research the latest security vulnerabilities. \
1043 Then, implement a fix in the backend API code. \
1044 Finally, deploy the fix using Docker and update the documentation.";
1045 let (decision, signals) = analyzer.analyze(task);
1046 assert!(
1047 signals.complexity_score > ORCHESTRATE_THRESHOLD
1048 || matches!(decision, RoutingDecision::Orchestrate { .. }),
1049 "decision: {decision:?}, score: {}",
1050 signals.complexity_score
1051 );
1052 }
1053
1054 #[test]
1055 fn analyze_delegation_with_agent_names_routes_orchestrate() {
1056 let agents = make_agents();
1057 let analyzer = TaskComplexityAnalyzer::new(&agents);
1058 let task =
1059 "Delegate to coder to implement the feature and have researcher find best practices";
1060 let (decision, signals) = analyzer.analyze(task);
1061 assert!(
1062 matches!(decision, RoutingDecision::Orchestrate { .. }),
1063 "decision: {decision:?}, score: {}",
1064 signals.complexity_score
1065 );
1066 }
1067
1068 #[test]
1071 fn resolve_routing_mode_uses_config_when_no_env() {
1072 unsafe {
1075 std::env::remove_var("HEARTBIT_ROUTING");
1076 }
1077 assert_eq!(
1078 resolve_routing_mode(RoutingMode::AlwaysOrchestrate),
1079 RoutingMode::AlwaysOrchestrate
1080 );
1081 }
1082
1083 #[test]
1088 fn keyword_routing_strategy_routes_simple_to_single() {
1089 let agents = make_agents();
1090 let strategy = KeywordRoutingStrategy;
1091 let (decision, _) = strategy.route("What is Rust?", &agents);
1092 assert!(
1093 matches!(decision, RoutingDecision::SingleAgent { .. }),
1094 "got: {decision:?}"
1095 );
1096 }
1097
1098 #[test]
1099 fn keyword_routing_strategy_routes_complex_to_orchestrate() {
1100 let agents = make_agents();
1101 let strategy = KeywordRoutingStrategy;
1102 let task =
1103 "Delegate to coder to implement the feature and have researcher find best practices";
1104 let (decision, _) = strategy.route(task, &agents);
1105 assert!(
1106 matches!(decision, RoutingDecision::Orchestrate { .. }),
1107 "got: {decision:?}"
1108 );
1109 }
1110
1111 #[test]
1112 fn custom_routing_strategy() {
1113 struct AlwaysOrchestrate;
1114 impl RoutingStrategy for AlwaysOrchestrate {
1115 fn route(
1116 &self,
1117 _task: &str,
1118 _agents: &[AgentCapability],
1119 ) -> (RoutingDecision, ComplexitySignals) {
1120 (
1121 RoutingDecision::Orchestrate {
1122 reason: "custom: always orchestrate",
1123 },
1124 ComplexitySignals::default(),
1125 )
1126 }
1127 }
1128
1129 let agents = make_agents();
1130 let strategy = AlwaysOrchestrate;
1131 let (decision, _) = strategy.route("What is 2+2?", &agents);
1132 assert!(
1133 matches!(decision, RoutingDecision::Orchestrate { .. }),
1134 "got: {decision:?}"
1135 );
1136 }
1137
1138 #[test]
1139 fn custom_routing_strategy_with_domain_matching() {
1140 struct PricingRouter;
1141 impl RoutingStrategy for PricingRouter {
1142 fn route(
1143 &self,
1144 task: &str,
1145 agents: &[AgentCapability],
1146 ) -> (RoutingDecision, ComplexitySignals) {
1147 let task_lower = task.to_lowercase();
1148 if task_lower.contains("pricing") || task_lower.contains("quote") {
1149 let idx = agents.iter().position(|a| a.name == "quoter").unwrap_or(0);
1150 return (
1151 RoutingDecision::SingleAgent {
1152 agent_index: idx,
1153 reason: "pricing domain detected",
1154 },
1155 ComplexitySignals::default(),
1156 );
1157 }
1158 KeywordRoutingStrategy.route(task, agents)
1160 }
1161 }
1162
1163 let agents = vec![
1164 AgentCapability::from_config("miner", "Finds sales leads", &[]),
1165 AgentCapability::from_config("quoter", "Generates pricing quotes", &[]),
1166 ];
1167 let strategy = PricingRouter;
1168
1169 let (decision, _) = strategy.route("Generate a pricing quote for the client", &agents);
1171 match decision {
1172 RoutingDecision::SingleAgent { agent_index, .. } => assert_eq!(agent_index, 1),
1173 other => panic!("expected SingleAgent, got: {other:?}"),
1174 }
1175
1176 let (decision, _) = strategy.route("What is 2+2?", &agents);
1178 assert!(matches!(decision, RoutingDecision::SingleAgent { .. }));
1179 }
1180
1181 #[test]
1182 fn routing_strategy_dyn_dispatch() {
1183 let strategy: std::sync::Arc<dyn RoutingStrategy> =
1185 std::sync::Arc::new(KeywordRoutingStrategy);
1186 let agents = make_agents();
1187 let (decision, _) = strategy.route("What is Rust?", &agents);
1188 assert!(matches!(decision, RoutingDecision::SingleAgent { .. }));
1189 }
1190
1191 #[test]
1192 fn missing_routing_field_defaults_to_auto() {
1193 #[derive(Deserialize)]
1194 struct TestConfig {
1195 #[serde(default)]
1196 routing: RoutingMode,
1197 }
1198 let config: TestConfig = toml::from_str("").unwrap();
1199 assert_eq!(config.routing, RoutingMode::Auto);
1200 }
1201
1202 #[test]
1205 fn contains_keyword_word_boundary() {
1206 assert!(!contains_keyword("builds backend api", "ui"));
1208 assert!(contains_keyword("the ui is broken", "ui"));
1210 assert!(contains_keyword("ui components", "ui"));
1212 assert!(contains_keyword("fix the ui", "ui"));
1214 assert!(contains_keyword("the api endpoint", "api"));
1216 assert!(!contains_keyword("the capital city", "api"));
1217 }
1218
1219 #[test]
1220 fn contains_keyword_multi_word() {
1221 assert!(contains_keyword("run the unit test suite", "unit test"));
1222 assert!(!contains_keyword("run the unittest suite", "unit test"));
1223 }
1224
1225 #[test]
1226 fn contains_keyword_adjacent_to_punctuation() {
1227 assert!(contains_keyword("fix the api.", "api"));
1229 assert!(contains_keyword("(api) endpoint", "api"));
1230 assert!(contains_keyword("api/rest", "api"));
1231 }
1232
1233 #[test]
1236 fn detect_domains_finds_multiple() {
1237 let domains = detect_domains("implement the api endpoint and write database migration");
1238 assert!(domains.contains(&"code".to_string())); assert!(domains.contains(&"backend".to_string())); assert!(domains.contains(&"database".to_string())); }
1242
1243 #[test]
1244 fn detect_domains_empty_for_generic_text() {
1245 let domains = detect_domains("hello world how are you");
1246 assert!(domains.is_empty());
1247 }
1248
1249 #[test]
1252 fn empty_task_routes_single_agent() {
1253 let agents = make_agents();
1254 let analyzer = TaskComplexityAnalyzer::new(&agents);
1255 let (decision, signals) = analyzer.analyze("");
1256 assert_eq!(signals.complexity_score, 0.0);
1257 assert!(matches!(decision, RoutingDecision::SingleAgent { .. }));
1258 }
1259
1260 #[test]
1261 fn single_agent_list_always_routes_single() {
1262 let agents = vec![AgentCapability::from_config("solo", "Does everything", &[])];
1263 let analyzer = TaskComplexityAnalyzer::new(&agents);
1264 let task = "Delegate the complex multi-step task that involves code, database, frontend, backend, security, devops, research, and writing";
1266 let (decision, signals) = analyzer.analyze(task);
1267 assert!(!signals.names_multiple_agents);
1269 match decision {
1271 RoutingDecision::SingleAgent { .. } | RoutingDecision::Orchestrate { .. } => {}
1272 }
1273 }
1274
1275 fn issue9_agents() -> Vec<AgentCapability> {
1290 vec![
1291 AgentCapability::from_config(
1292 "researcher",
1293 "Investigates a topic, gathers facts.",
1294 &["web_search".into()],
1295 ),
1296 AgentCapability::from_config(
1297 "coder",
1298 "Writes and refactors Rust code.",
1299 &["read".into(), "write".into(), "edit".into()],
1300 ),
1301 ]
1302 }
1303
1304 #[test]
1305 fn gh9_tier1_routes_to_best_covering_agent_when_domains_present() {
1306 let agents = issue9_agents();
1311 let analyzer = TaskComplexityAnalyzer::new(&agents);
1312 let task = "Ask the coder to refactor the parse_args function and write tests.";
1313 let (decision, signals) = analyzer.analyze(task);
1314 assert!(
1315 signals.domain_signals.contains(&"code".to_string()),
1316 "expected `code` in domain_signals, got: {:?}",
1317 signals.domain_signals
1318 );
1319 match decision {
1320 RoutingDecision::SingleAgent { agent_index, .. } => assert_eq!(
1321 agent_index, 1,
1322 "Tier 1 should route to coder (idx 1), got idx {agent_index}",
1323 ),
1324 RoutingDecision::Orchestrate { .. } => {}
1327 }
1328 }
1329
1330 #[test]
1331 fn gh9_parenthesised_numbers_count_as_step_markers() {
1332 let agents = make_agents();
1334 let analyzer = TaskComplexityAnalyzer::new(&agents);
1335 let task = "We need to: (1) investigate Rust async; (2) compare tokio vs smol; \
1336 (3) write a benchmark; (4) implement it; (5) review the code; \
1337 (6) summarize findings.";
1338 let signals = analyzer.heuristic_signals(task);
1339 assert!(
1340 signals.step_markers >= 4,
1341 "expected step_markers >= 4 for `(1) (2) ... (6)` pattern, got {}",
1342 signals.step_markers
1343 );
1344 }
1345
1346 #[test]
1347 fn gh9_right_paren_numbers_count_as_step_markers() {
1348 let agents = make_agents();
1350 let analyzer = TaskComplexityAnalyzer::new(&agents);
1351 let task = "Plan: 1) gather data; 2) draft outline; 3) review; 4) publish.";
1352 let signals = analyzer.heuristic_signals(task);
1353 assert!(
1354 signals.step_markers >= 3,
1355 "expected step_markers >= 3 for `1) 2) 3) 4)` pattern, got {}",
1356 signals.step_markers
1357 );
1358 }
1359
1360 #[test]
1361 fn gh9_ask_the_phrasing_triggers_delegation() {
1362 let agents = issue9_agents();
1363 let analyzer = TaskComplexityAnalyzer::new(&agents);
1364 let signals =
1365 analyzer.heuristic_signals("Ask the coder to refactor the parse_args function.");
1366 assert!(
1367 signals.explicit_delegation,
1368 "`ask the X to ...` should trigger explicit_delegation",
1369 );
1370 }
1371
1372 #[test]
1373 fn gh9_have_the_phrasing_triggers_delegation() {
1374 let agents = issue9_agents();
1375 let analyzer = TaskComplexityAnalyzer::new(&agents);
1376 let signals = analyzer.heuristic_signals("Have the researcher gather data on Rust async.");
1377 assert!(
1378 signals.explicit_delegation,
1379 "`have the X ...` should trigger explicit_delegation",
1380 );
1381 }
1382
1383 #[test]
1384 fn gh9_no_domains_still_falls_back_to_agent_zero() {
1385 let agents = issue9_agents();
1388 let analyzer = TaskComplexityAnalyzer::new(&agents);
1389 let (decision, signals) = analyzer.analyze("Hello, are you there?");
1390 assert!(
1391 signals.domain_signals.is_empty(),
1392 "domains should be empty for plain greeting, got: {:?}",
1393 signals.domain_signals
1394 );
1395 assert!(matches!(
1396 decision,
1397 RoutingDecision::SingleAgent { agent_index: 0, .. }
1398 ));
1399 }
1400}