1use std::collections::{HashMap, HashSet};
43
44use serde::{Deserialize, Serialize};
45
46#[derive(Debug, Clone, Serialize, Deserialize)]
55pub struct DependencyEdge {
56 pub from: String,
58 pub to: String,
60 pub confidence: f64,
62}
63
64impl DependencyEdge {
65 pub fn new(from: impl Into<String>, to: impl Into<String>, confidence: f64) -> Self {
66 Self {
67 from: from.into(),
68 to: to.into(),
69 confidence: confidence.clamp(0.0, 1.0),
70 }
71 }
72}
73
74#[derive(Debug, Clone, Default, Serialize, Deserialize)]
79pub struct DependencyGraph {
80 edges: Vec<DependencyEdge>,
82 start_nodes: HashSet<String>,
84 terminal_nodes: HashSet<String>,
86 task: String,
88 available_actions: Vec<String>,
90 #[serde(default)]
94 param_variants: HashMap<String, (String, Vec<String>)>,
95 #[serde(default)]
97 discover_order: Vec<String>,
98 #[serde(default)]
100 not_discover_order: Vec<String>,
101}
102
103impl DependencyGraph {
104 pub fn new() -> Self {
106 Self::default()
107 }
108
109 pub fn builder() -> DependencyGraphBuilder {
111 DependencyGraphBuilder::new()
112 }
113
114 pub fn valid_next_actions(&self, current_action: &str) -> Vec<String> {
122 let mut edges: Vec<_> = self
123 .edges
124 .iter()
125 .filter(|e| e.from == current_action)
126 .collect();
127
128 edges.sort_by(|a, b| {
130 b.confidence
131 .partial_cmp(&a.confidence)
132 .unwrap_or(std::cmp::Ordering::Equal)
133 });
134
135 edges.iter().map(|e| e.to.clone()).collect()
136 }
137
138 pub fn start_actions(&self) -> Vec<String> {
142 let mut actions: Vec<_> = self.start_nodes.iter().cloned().collect();
143 actions.sort();
144 actions
145 }
146
147 pub fn terminal_actions(&self) -> Vec<String> {
151 let mut actions: Vec<_> = self.terminal_nodes.iter().cloned().collect();
152 actions.sort();
153 actions
154 }
155
156 pub fn is_terminal(&self, action: &str) -> bool {
158 self.terminal_nodes.contains(action)
159 }
160
161 pub fn is_start(&self, action: &str) -> bool {
163 self.start_nodes.contains(action)
164 }
165
166 pub fn can_transition(&self, from: &str, to: &str) -> bool {
168 self.edges.iter().any(|e| e.from == from && e.to == to)
169 }
170
171 pub fn transition_confidence(&self, from: &str, to: &str) -> Option<f64> {
173 self.edges
174 .iter()
175 .find(|e| e.from == from && e.to == to)
176 .map(|e| e.confidence)
177 }
178
179 pub fn edges(&self) -> &[DependencyEdge] {
181 &self.edges
182 }
183
184 pub fn task(&self) -> &str {
186 &self.task
187 }
188
189 pub fn available_actions(&self) -> &[String] {
191 &self.available_actions
192 }
193
194 pub fn param_variants(&self, action: &str) -> Option<(&str, &[String])> {
196 self.param_variants
197 .get(action)
198 .map(|(key, values)| (key.as_str(), values.as_slice()))
199 }
200
201 pub fn all_param_variants(&self) -> &HashMap<String, (String, Vec<String>)> {
203 &self.param_variants
204 }
205
206 pub fn discover_order(&self) -> &[String] {
212 &self.discover_order
213 }
214
215 pub fn not_discover_order(&self) -> &[String] {
217 &self.not_discover_order
218 }
219
220 pub fn set_action_order(&mut self, discover: Vec<String>, not_discover: Vec<String>) {
222 self.discover_order = discover;
223 self.not_discover_order = not_discover;
224 }
225
226 pub fn has_action_order(&self) -> bool {
228 !self.discover_order.is_empty() || !self.not_discover_order.is_empty()
229 }
230
231 pub fn validate(&self) -> Result<(), DependencyGraphError> {
243 if self.start_nodes.is_empty() {
244 return Err(DependencyGraphError::NoStartNodes);
245 }
246
247 if self.terminal_nodes.is_empty() {
248 return Err(DependencyGraphError::NoTerminalNodes);
249 }
250
251 for node in &self.start_nodes {
253 if !self.available_actions.contains(node) {
254 return Err(DependencyGraphError::UnknownAction(node.clone()));
255 }
256 }
257
258 for node in &self.terminal_nodes {
260 if !self.available_actions.contains(node) {
261 return Err(DependencyGraphError::UnknownAction(node.clone()));
262 }
263 }
264
265 for edge in &self.edges {
267 if !self.available_actions.contains(&edge.from) {
268 return Err(DependencyGraphError::UnknownAction(edge.from.clone()));
269 }
270 if !self.available_actions.contains(&edge.to) {
271 return Err(DependencyGraphError::UnknownAction(edge.to.clone()));
272 }
273 }
274
275 Ok(())
276 }
277
278 pub fn to_mermaid(&self) -> String {
284 let mut lines = vec!["graph LR".to_string()];
285
286 for edge in &self.edges {
287 let label = format!("{:.0}%", edge.confidence * 100.0);
288 lines.push(format!(" {} -->|{}| {}", edge.from, label, edge.to));
289 }
290
291 for start in &self.start_nodes {
293 lines.push(format!(" style {} fill:#9f9", start));
294 }
295
296 for terminal in &self.terminal_nodes {
298 lines.push(format!(" style {} fill:#f99", terminal));
299 }
300
301 lines.join("\n")
302 }
303}
304
305#[derive(Debug, Clone, thiserror::Error)]
307pub enum DependencyGraphError {
308 #[error("No start nodes defined")]
309 NoStartNodes,
310
311 #[error("No terminal nodes defined")]
312 NoTerminalNodes,
313
314 #[error("Unknown action: {0}")]
315 UnknownAction(String),
316
317 #[error("Parse error: {0}")]
318 ParseError(String),
319
320 #[error("LLM error: {0}")]
321 LlmError(String),
322}
323
324pub trait DependencyGraphProvider: Send + Sync {
351 fn provide_graph(&self, task: &str, available_actions: &[String]) -> Option<DependencyGraph>;
361}
362
363#[derive(Debug, Clone, Default)]
369pub struct DependencyGraphBuilder {
370 edges: Vec<DependencyEdge>,
371 start_nodes: HashSet<String>,
372 terminal_nodes: HashSet<String>,
373 task: String,
374 available_actions: Vec<String>,
375 param_variants: HashMap<String, (String, Vec<String>)>,
376}
377
378impl DependencyGraphBuilder {
379 pub fn new() -> Self {
380 Self::default()
381 }
382
383 pub fn task(mut self, task: impl Into<String>) -> Self {
385 self.task = task.into();
386 self
387 }
388
389 pub fn available_actions<I, S>(mut self, actions: I) -> Self
391 where
392 I: IntoIterator<Item = S>,
393 S: Into<String>,
394 {
395 self.available_actions = actions.into_iter().map(|s| s.into()).collect();
396 self
397 }
398
399 pub fn edge(mut self, from: impl Into<String>, to: impl Into<String>, confidence: f64) -> Self {
401 self.edges.push(DependencyEdge::new(from, to, confidence));
402 self
403 }
404
405 pub fn start_node(mut self, action: impl Into<String>) -> Self {
407 self.start_nodes.insert(action.into());
408 self
409 }
410
411 pub fn start_nodes<I, S>(mut self, actions: I) -> Self
413 where
414 I: IntoIterator<Item = S>,
415 S: Into<String>,
416 {
417 self.start_nodes
418 .extend(actions.into_iter().map(|s| s.into()));
419 self
420 }
421
422 pub fn terminal_node(mut self, action: impl Into<String>) -> Self {
424 self.terminal_nodes.insert(action.into());
425 self
426 }
427
428 pub fn terminal_nodes<I, S>(mut self, actions: I) -> Self
430 where
431 I: IntoIterator<Item = S>,
432 S: Into<String>,
433 {
434 self.terminal_nodes
435 .extend(actions.into_iter().map(|s| s.into()));
436 self
437 }
438
439 pub fn param_variants<I, S>(
443 mut self,
444 action: impl Into<String>,
445 key: impl Into<String>,
446 values: I,
447 ) -> Self
448 where
449 I: IntoIterator<Item = S>,
450 S: Into<String>,
451 {
452 self.param_variants.insert(
453 action.into(),
454 (key.into(), values.into_iter().map(|s| s.into()).collect()),
455 );
456 self
457 }
458
459 pub fn build(self) -> DependencyGraph {
461 DependencyGraph {
462 edges: self.edges,
463 start_nodes: self.start_nodes,
464 terminal_nodes: self.terminal_nodes,
465 task: self.task,
466 available_actions: self.available_actions,
467 param_variants: self.param_variants,
468 discover_order: Vec::new(),
469 not_discover_order: Vec::new(),
470 }
471 }
472
473 pub fn build_validated(self) -> Result<DependencyGraph, DependencyGraphError> {
475 let graph = self.build();
476 graph.validate()?;
477 Ok(graph)
478 }
479}
480
481#[derive(Debug, Clone, Serialize, Deserialize)]
489pub struct LlmDependencyResponse {
490 pub edges: Vec<LlmEdge>,
492 pub start: Vec<String>,
494 pub terminal: Vec<String>,
496 #[serde(default)]
498 pub reasoning: Option<String>,
499}
500
501#[derive(Debug, Clone, Serialize, Deserialize)]
503pub struct LlmEdge {
504 pub from: String,
505 pub to: String,
506 pub confidence: f64,
507}
508
509impl LlmDependencyResponse {
510 pub fn into_graph(
512 self,
513 task: impl Into<String>,
514 available_actions: Vec<String>,
515 ) -> DependencyGraph {
516 let mut builder = DependencyGraphBuilder::new()
517 .task(task)
518 .available_actions(available_actions)
519 .start_nodes(self.start)
520 .terminal_nodes(self.terminal);
521
522 for edge in self.edges {
523 builder = builder.edge(edge.from, edge.to, edge.confidence);
524 }
525
526 builder.build()
527 }
528
529 pub fn parse(text: &str) -> Result<Self, DependencyGraphError> {
535 if let Some(response) = Self::parse_arrow_format(text) {
537 return Ok(response);
538 }
539
540 if let Ok(parsed) = serde_json::from_str(text) {
542 return Ok(parsed);
543 }
544
545 if let Some(json) = Self::extract_json(text) {
547 serde_json::from_str(&json).map_err(|e| DependencyGraphError::ParseError(e.to_string()))
548 } else {
549 Err(DependencyGraphError::ParseError(format!(
550 "No valid format found in response: {}",
551 text.chars().take(200).collect::<String>()
552 )))
553 }
554 }
555
556 fn parse_arrow_format(text: &str) -> Option<Self> {
563 if let Some(result) = Self::parse_arrow_only(text) {
565 return Some(result);
566 }
567
568 if let Some(result) = Self::parse_numbered_list(text) {
570 return Some(result);
571 }
572
573 None
574 }
575
576 fn parse_arrow_only(text: &str) -> Option<Self> {
578 let normalized = text.replace('→', "->");
579
580 let arrow_line = normalized.lines().find(|line| line.contains("->"))?;
582
583 let parts: Vec<&str> = arrow_line.split("->").collect();
584 if parts.len() < 2 {
585 return None;
586 }
587
588 let actions_in_order: Vec<String> = parts
590 .iter()
591 .filter_map(|part| {
592 let trimmed = part.trim();
593 let last_word = trimmed.split_whitespace().last()?;
595 let action: String = last_word.chars().filter(|c| c.is_alphabetic()).collect();
597 if action.is_empty() {
598 None
599 } else {
600 Some(action)
601 }
602 })
603 .collect();
604
605 if actions_in_order.len() < 2 {
606 return None;
607 }
608
609 Self::build_response(actions_in_order)
610 }
611
612 fn parse_numbered_list(text: &str) -> Option<Self> {
614 let mut actions_in_order: Vec<String> = Vec::new();
615
616 for i in 1..=10 {
618 let pattern = format!("{}.", i);
619 if let Some(pos) = text.find(&pattern) {
620 let after = &text[pos + pattern.len()..];
622 if let Some(word) = after.split_whitespace().next() {
623 let action: String = word.chars().filter(|c| c.is_alphabetic()).collect();
625 if !action.is_empty() && !actions_in_order.contains(&action) {
626 actions_in_order.push(action);
627 }
628 }
629 }
630 }
631
632 if actions_in_order.len() < 2 {
633 return None;
634 }
635
636 Self::build_response(actions_in_order)
637 }
638
639 fn build_response(actions_in_order: Vec<String>) -> Option<Self> {
641 let mut edges = Vec::new();
642 for window in actions_in_order.windows(2) {
643 edges.push(LlmEdge {
644 from: window[0].clone(),
645 to: window[1].clone(),
646 confidence: 0.9,
647 });
648 }
649
650 Some(Self {
651 edges,
652 start: vec![actions_in_order.first()?.clone()],
653 terminal: vec![actions_in_order.last()?.clone()],
654 reasoning: Some("Parsed from text format".to_string()),
655 })
656 }
657
658 fn extract_json(text: &str) -> Option<String> {
660 let start = text.find('{')?;
662 let chars: Vec<char> = text[start..].chars().collect();
663 let mut depth = 0;
664 let mut in_string = false;
665 let mut escape_next = false;
666
667 for (i, &ch) in chars.iter().enumerate() {
668 if escape_next {
669 escape_next = false;
670 continue;
671 }
672
673 match ch {
674 '\\' if in_string => escape_next = true,
675 '"' => in_string = !in_string,
676 '{' if !in_string => depth += 1,
677 '}' if !in_string => {
678 depth -= 1;
679 if depth == 0 {
680 return Some(chars[..=i].iter().collect());
681 }
682 }
683 _ => {}
684 }
685 }
686
687 None
688 }
689}
690
691pub trait DependencyPlanner: Send + Sync {
699 fn plan(
703 &self,
704 task: &str,
705 available_actions: &[String],
706 ) -> Result<DependencyGraph, DependencyGraphError>;
707
708 fn name(&self) -> &str;
710}
711
712#[derive(Debug, Clone, Default)]
717pub struct StaticDependencyPlanner {
718 patterns: HashMap<String, DependencyGraph>,
720 default_pattern: Option<String>,
722}
723
724impl StaticDependencyPlanner {
725 pub fn new() -> Self {
726 Self::default()
727 }
728
729 pub fn with_pattern(mut self, name: impl Into<String>, graph: DependencyGraph) -> Self {
731 let name = name.into();
732 if self.default_pattern.is_none() {
733 self.default_pattern = Some(name.clone());
734 }
735 self.patterns.insert(name, graph);
736 self
737 }
738
739 pub fn with_default_pattern(mut self, name: impl Into<String>) -> Self {
741 self.default_pattern = Some(name.into());
742 self
743 }
744
745 pub fn with_file_exploration_pattern(self) -> Self {
749 let graph = DependencyGraph::builder()
750 .task("File exploration")
751 .available_actions(["Grep", "List", "Read"])
752 .edge("Grep", "Read", 0.95)
753 .edge("List", "Grep", 0.60)
754 .edge("List", "Read", 0.40)
755 .start_nodes(["Grep", "List"])
756 .terminal_node("Read")
757 .build();
758
759 self.with_pattern("file_exploration", graph)
760 }
761
762 pub fn with_code_search_pattern(self) -> Self {
766 let graph = DependencyGraph::builder()
767 .task("Code search")
768 .available_actions(["Grep", "Read"])
769 .edge("Grep", "Read", 0.95)
770 .start_node("Grep")
771 .terminal_node("Read")
772 .build();
773
774 self.with_pattern("code_search", graph)
775 }
776}
777
778impl DependencyPlanner for StaticDependencyPlanner {
779 fn plan(
780 &self,
781 task: &str,
782 available_actions: &[String],
783 ) -> Result<DependencyGraph, DependencyGraphError> {
784 if let Some(pattern_name) = &self.default_pattern {
786 if let Some(graph) = self.patterns.get(pattern_name) {
787 let mut graph = graph.clone();
788 graph.task = task.to_string();
789 graph.available_actions = available_actions.to_vec();
790 return Ok(graph);
791 }
792 }
793
794 if available_actions.is_empty() {
797 return Err(DependencyGraphError::NoStartNodes);
798 }
799
800 let mut builder = DependencyGraphBuilder::new()
801 .task(task)
802 .available_actions(available_actions.to_vec())
803 .start_node(&available_actions[0]);
804
805 if available_actions.len() > 1 {
806 for window in available_actions.windows(2) {
807 builder = builder.edge(&window[0], &window[1], 0.80);
808 }
809 builder = builder.terminal_node(&available_actions[available_actions.len() - 1]);
810 } else {
811 builder = builder.terminal_node(&available_actions[0]);
812 }
813
814 Ok(builder.build())
815 }
816
817 fn name(&self) -> &str {
818 "StaticDependencyPlanner"
819 }
820}
821
822use crate::actions::ActionDef;
827
828pub struct DependencyPromptGenerator;
833
834impl DependencyPromptGenerator {
835 pub fn generate_prompt(task: &str, actions: &[ActionDef]) -> String {
839 let actions_list = actions
840 .iter()
841 .map(|a| a.name.as_str())
842 .collect::<Vec<_>>()
843 .join(", ");
844
845 format!(
846 r#"{task}
847Steps: {actions_list}
848The very first step is:"#
849 )
850 }
851
852 pub fn generate_first_prompt(_task: &str, actions: &[ActionDef]) -> String {
857 let mut sorted_actions: Vec<&ActionDef> = actions.iter().collect();
859 sorted_actions.sort_by(|a, b| a.name.cmp(&b.name));
860
861 let actions_list = sorted_actions
862 .iter()
863 .map(|a| a.name.as_str())
864 .collect::<Vec<_>>()
865 .join(", ");
866
867 let descriptions: Vec<String> = sorted_actions
869 .iter()
870 .map(|a| format!("- {}: {}", a.name, a.description))
871 .collect();
872 let descriptions_block = descriptions.join("\n");
873
874 let first_verb = sorted_actions
876 .first()
877 .map(|a| Self::extract_verb(&a.description))
878 .unwrap_or_else(|| "CHECK".to_string());
879
880 format!(
881 r#"Steps: {actions_list}
882{descriptions_block}
883Which step {first_verb}S first?
884Answer:"#
885 )
886 }
887
888 pub fn generate_last_prompt(_task: &str, actions: &[ActionDef]) -> String {
893 let mut sorted_actions: Vec<&ActionDef> = actions.iter().collect();
895 sorted_actions.sort_by(|a, b| a.name.cmp(&b.name));
896
897 let actions_list = sorted_actions
898 .iter()
899 .map(|a| a.name.as_str())
900 .collect::<Vec<_>>()
901 .join(", ");
902
903 let descriptions: Vec<String> = sorted_actions
904 .iter()
905 .map(|a| format!("- {}: {}", a.name, a.description))
906 .collect();
907 let descriptions_block = descriptions.join("\n");
908
909 format!(
910 r#"Steps: {actions_list}
911{descriptions_block}
912Which step should be done last?
913Answer:"#
914 )
915 }
916
917 pub fn generate_pair_prompt(task: &str, action_a: &str, action_b: &str) -> String {
919 format!(
920 r#"For {task}, which comes first: {action_a} or {action_b}?
921Answer (one word):"#
922 )
923 }
924
925 fn extract_verb(description: &str) -> String {
930 description
931 .split_whitespace()
932 .next()
933 .map(|w| {
934 let word = w.trim_end_matches('s').trim_end_matches('S');
936 word.to_uppercase()
937 })
938 .unwrap_or_else(|| "CHECK".to_string())
939 }
940}
941
942#[derive(Debug, Clone)]
951pub struct GraphNavigator {
952 graph: DependencyGraph,
953 completed_actions: HashSet<String>,
955}
956
957impl GraphNavigator {
958 pub fn new(graph: DependencyGraph) -> Self {
959 Self {
960 graph,
961 completed_actions: HashSet::new(),
962 }
963 }
964
965 pub fn mark_completed(&mut self, action: &str) {
967 self.completed_actions.insert(action.to_string());
968 }
969
970 pub fn suggest_next(&self) -> Vec<String> {
976 if self.completed_actions.is_empty() {
977 return self.graph.start_actions();
979 }
980
981 let mut candidates = Vec::new();
983 for completed in &self.completed_actions {
984 for next in self.graph.valid_next_actions(completed) {
985 if !self.completed_actions.contains(&next) && !candidates.contains(&next) {
986 candidates.push(next);
987 }
988 }
989 }
990
991 candidates
992 }
993
994 pub fn is_task_complete(&self) -> bool {
998 self.graph
999 .terminal_actions()
1000 .iter()
1001 .any(|t| self.completed_actions.contains(t))
1002 }
1003
1004 pub fn progress(&self) -> f64 {
1006 if self.graph.available_actions.is_empty() {
1007 return 0.0;
1008 }
1009 self.completed_actions.len() as f64 / self.graph.available_actions.len() as f64
1010 }
1011
1012 pub fn graph(&self) -> &DependencyGraph {
1014 &self.graph
1015 }
1016}
1017
1018use crate::learn::offline::LearnedActionOrder;
1023
1024#[derive(Debug, Clone)]
1045pub struct LearnedDependencyProvider {
1046 action_order: LearnedActionOrder,
1047}
1048
1049impl LearnedDependencyProvider {
1050 pub fn new(action_order: LearnedActionOrder) -> Self {
1052 Self { action_order }
1053 }
1054
1055 pub fn action_order(&self) -> &LearnedActionOrder {
1057 &self.action_order
1058 }
1059}
1060
1061impl DependencyGraphProvider for LearnedDependencyProvider {
1062 fn provide_graph(&self, task: &str, available_actions: &[String]) -> Option<DependencyGraph> {
1063 if !self.action_order.matches_actions(available_actions) {
1065 tracing::debug!(
1066 expected_hash = self.action_order.action_set_hash,
1067 actual_hash = LearnedActionOrder::compute_hash(available_actions),
1068 "Action set mismatch, falling back to LLM"
1069 );
1070 return None;
1071 }
1072
1073 let sorted_discover = &self.action_order.discover;
1075 let sorted_not_discover = &self.action_order.not_discover;
1076
1077 if sorted_discover.is_empty() && sorted_not_discover.is_empty() {
1079 tracing::warn!("Both discover and not_discover are empty, cannot build graph");
1080 return None;
1081 }
1082
1083 tracing::info!(
1084 discover = ?self.action_order.discover,
1085 not_discover = ?self.action_order.not_discover,
1086 "Using learned action order (LLM skipped)"
1087 );
1088
1089 let mut builder = DependencyGraphBuilder::new()
1090 .task(task)
1091 .available_actions(available_actions.iter().cloned());
1092
1093 if !sorted_discover.is_empty() {
1095 builder = builder.start_node(&sorted_discover[0]);
1096 } else if !sorted_not_discover.is_empty() {
1097 builder = builder.start_node(&sorted_not_discover[0]);
1098 }
1099
1100 if let Some(last) = sorted_not_discover.last() {
1102 builder = builder.terminal_node(last);
1103 } else if !sorted_discover.is_empty() {
1104 builder = builder.terminal_node(sorted_discover.last().unwrap());
1105 }
1106
1107 for window in sorted_discover.windows(2) {
1109 builder = builder.edge(&window[0], &window[1], 0.9);
1110 }
1111
1112 if !sorted_discover.is_empty() && !sorted_not_discover.is_empty() {
1114 builder = builder.edge(
1115 sorted_discover.last().unwrap(),
1116 &sorted_not_discover[0],
1117 0.9,
1118 );
1119 }
1120
1121 for window in sorted_not_discover.windows(2) {
1123 builder = builder.edge(&window[0], &window[1], 0.9);
1124 }
1125
1126 builder.build_validated().ok()
1128 }
1129}
1130
1131#[cfg(test)]
1136mod tests {
1137 use super::*;
1138
1139 #[test]
1140 fn test_dependency_graph_builder() {
1141 let graph = DependencyGraph::builder()
1142 .task("Find auth function")
1143 .available_actions(["Grep", "List", "Read"])
1144 .edge("Grep", "Read", 0.95)
1145 .edge("List", "Grep", 0.60)
1146 .start_nodes(["Grep", "List"])
1147 .terminal_node("Read")
1148 .build();
1149
1150 assert_eq!(graph.task(), "Find auth function");
1151 assert!(graph.is_start("Grep"));
1152 assert!(graph.is_start("List"));
1153 assert!(graph.is_terminal("Read"));
1154 assert!(graph.can_transition("Grep", "Read"));
1155 assert!(!graph.can_transition("Read", "Grep"));
1156 }
1157
1158 #[test]
1159 fn test_valid_next_actions() {
1160 let graph = DependencyGraph::builder()
1161 .available_actions(["Grep", "List", "Read"])
1162 .edge("Grep", "Read", 0.95)
1163 .edge("List", "Grep", 0.60)
1164 .edge("List", "Read", 0.40)
1165 .start_nodes(["Grep", "List"])
1166 .terminal_node("Read")
1167 .build();
1168
1169 let next = graph.valid_next_actions("Grep");
1171 assert_eq!(next, vec!["Read"]);
1172
1173 let next = graph.valid_next_actions("List");
1175 assert_eq!(next, vec!["Grep", "Read"]);
1176
1177 let next = graph.valid_next_actions("Read");
1179 assert!(next.is_empty());
1180 }
1181
1182 #[test]
1183 fn test_static_planner_file_exploration() {
1184 let planner = StaticDependencyPlanner::new().with_file_exploration_pattern();
1185
1186 let graph = planner
1187 .plan("Find auth.rs", &["Grep".to_string(), "Read".to_string()])
1188 .unwrap();
1189
1190 assert!(graph.is_start("Grep"));
1191 assert!(graph.is_terminal("Read"));
1192 }
1193
1194 #[test]
1195 fn test_graph_navigator() {
1196 let graph = DependencyGraph::builder()
1197 .available_actions(["Grep", "Read"])
1198 .edge("Grep", "Read", 0.95)
1199 .start_node("Grep")
1200 .terminal_node("Read")
1201 .build();
1202
1203 let mut nav = GraphNavigator::new(graph);
1204
1205 assert_eq!(nav.suggest_next(), vec!["Grep"]);
1207 assert!(!nav.is_task_complete());
1208
1209 nav.mark_completed("Grep");
1211 assert_eq!(nav.suggest_next(), vec!["Read"]);
1212 assert!(!nav.is_task_complete());
1213
1214 nav.mark_completed("Read");
1216 assert!(nav.is_task_complete());
1217 assert!(nav.suggest_next().is_empty());
1218 }
1219
1220 #[test]
1221 fn test_llm_response_parsing() {
1222 let json = r#"{
1223 "edges": [
1224 {"from": "Grep", "to": "Read", "confidence": 0.95}
1225 ],
1226 "start": ["Grep"],
1227 "terminal": ["Read"],
1228 "reasoning": "Search first, then read"
1229 }"#;
1230
1231 let response = LlmDependencyResponse::parse(json).unwrap();
1232 assert_eq!(response.edges.len(), 1);
1233 assert_eq!(response.start, vec!["Grep"]);
1234 assert_eq!(response.terminal, vec!["Read"]);
1235 assert!(response.reasoning.is_some());
1236
1237 let graph = response.into_graph(
1238 "Find function",
1239 vec!["Grep".to_string(), "Read".to_string()],
1240 );
1241 assert!(graph.can_transition("Grep", "Read"));
1242 }
1243
1244 #[test]
1245 fn test_mermaid_output() {
1246 let graph = DependencyGraph::builder()
1247 .available_actions(["Grep", "List", "Read"])
1248 .edge("Grep", "Read", 0.95)
1249 .edge("List", "Grep", 0.60)
1250 .start_nodes(["Grep", "List"])
1251 .terminal_node("Read")
1252 .build();
1253
1254 let mermaid = graph.to_mermaid();
1255 assert!(mermaid.contains("graph LR"));
1256 assert!(mermaid.contains("Grep -->|95%| Read"));
1257 assert!(mermaid.contains("style Read fill:#f99"));
1258 }
1259
1260 #[test]
1265 fn test_learned_action_order_hash() {
1266 let actions = vec![
1267 "Grep".to_string(),
1268 "Read".to_string(),
1269 "Restart".to_string(),
1270 ];
1271
1272 let order = LearnedActionOrder::new(
1273 vec!["Grep".to_string(), "Read".to_string()],
1274 vec!["Restart".to_string()],
1275 &actions,
1276 );
1277
1278 let actions_reordered = vec![
1280 "Restart".to_string(),
1281 "Grep".to_string(),
1282 "Read".to_string(),
1283 ];
1284 assert!(order.matches_actions(&actions_reordered));
1285
1286 let actions_different = vec!["Grep".to_string(), "Read".to_string()];
1288 assert!(!order.matches_actions(&actions_different));
1289 }
1290
1291 #[test]
1292 fn test_learned_dependency_provider_cache_hit() {
1293 let actions = vec![
1294 "Grep".to_string(),
1295 "Read".to_string(),
1296 "Restart".to_string(),
1297 ];
1298
1299 let order = LearnedActionOrder::new(
1300 vec!["Grep".to_string(), "Read".to_string()],
1301 vec!["Restart".to_string()],
1302 &actions,
1303 );
1304
1305 let provider = LearnedDependencyProvider::new(order);
1306
1307 let graph = provider.provide_graph("troubleshooting", &actions);
1309 assert!(graph.is_some());
1310
1311 let graph = graph.unwrap();
1312 assert!(graph.is_start("Grep"));
1313 assert!(graph.is_terminal("Restart"));
1314 assert!(graph.can_transition("Grep", "Read"));
1315 assert!(graph.can_transition("Read", "Restart"));
1316 }
1317
1318 #[test]
1319 fn test_learned_dependency_provider_cache_miss() {
1320 let original_actions = vec![
1321 "Grep".to_string(),
1322 "Read".to_string(),
1323 "Restart".to_string(),
1324 ];
1325
1326 let order = LearnedActionOrder::new(
1327 vec!["Grep".to_string(), "Read".to_string()],
1328 vec!["Restart".to_string()],
1329 &original_actions,
1330 );
1331
1332 let provider = LearnedDependencyProvider::new(order);
1333
1334 let different_actions = vec!["Grep".to_string(), "Read".to_string()];
1336 let graph = provider.provide_graph("troubleshooting", &different_actions);
1337 assert!(graph.is_none());
1338 }
1339
1340 #[test]
1341 fn test_learned_dependency_provider_discover_only() {
1342 let actions = vec!["Grep".to_string(), "Read".to_string()];
1343
1344 let order = LearnedActionOrder::new(
1345 vec!["Grep".to_string(), "Read".to_string()],
1346 vec![], &actions,
1348 );
1349
1350 let provider = LearnedDependencyProvider::new(order);
1351 let graph = provider.provide_graph("search task", &actions);
1352 assert!(graph.is_some());
1353
1354 let graph = graph.unwrap();
1355 assert!(graph.is_start("Grep"));
1356 assert!(graph.is_terminal("Read")); assert!(graph.can_transition("Grep", "Read"));
1358 }
1359
1360 #[test]
1361 fn test_learned_dependency_provider_not_discover_only() {
1362 let actions = vec!["Restart".to_string(), "CheckStatus".to_string()];
1363
1364 let order = LearnedActionOrder::new(
1365 vec![], vec!["Restart".to_string(), "CheckStatus".to_string()],
1367 &actions,
1368 );
1369
1370 let provider = LearnedDependencyProvider::new(order);
1371 let graph = provider.provide_graph("ops task", &actions);
1372 assert!(graph.is_some());
1373
1374 let graph = graph.unwrap();
1375 assert!(graph.is_start("Restart")); assert!(graph.is_terminal("CheckStatus"));
1377 assert!(graph.can_transition("Restart", "CheckStatus"));
1378 }
1379
1380 #[test]
1381 fn test_learned_dependency_provider_empty_lists() {
1382 let actions = vec!["Grep".to_string(), "Read".to_string()];
1383
1384 let order = LearnedActionOrder::new(
1385 vec![], vec![], &actions,
1388 );
1389
1390 let provider = LearnedDependencyProvider::new(order);
1391 let graph = provider.provide_graph("empty task", &actions);
1393 assert!(graph.is_none());
1394 }
1395
1396 #[test]
1401 fn test_extract_json_simple() {
1402 let text = r#"Here is the result: {"edges": [], "start": ["A"], "terminal": ["B"]}"#;
1403 let json = LlmDependencyResponse::extract_json(text);
1404 assert!(json.is_some());
1405 let json = json.unwrap();
1406 assert!(json.starts_with('{'));
1407 assert!(json.ends_with('}'));
1408 }
1409
1410 #[test]
1411 fn test_extract_json_nested() {
1412 let text = r#"Result: {"edges": [{"from": "A", "to": "B", "confidence": 0.9}], "start": ["A"], "terminal": ["B"]}"#;
1413 let json = LlmDependencyResponse::extract_json(text);
1414 assert!(json.is_some());
1415
1416 let parsed: Result<LlmDependencyResponse, _> = serde_json::from_str(&json.unwrap());
1418 assert!(parsed.is_ok());
1419 }
1420
1421 #[test]
1422 fn test_extract_json_with_string_braces() {
1423 let text =
1425 r#"{"edges": [], "start": ["A"], "terminal": ["B"], "reasoning": "Use {pattern}"}"#;
1426 let json = LlmDependencyResponse::extract_json(text);
1427 assert!(json.is_some());
1428 assert_eq!(json.unwrap(), text);
1429 }
1430
1431 #[test]
1432 fn test_extract_json_no_json() {
1433 let text = "This is just plain text without JSON";
1434 let json = LlmDependencyResponse::extract_json(text);
1435 assert!(json.is_none());
1436 }
1437
1438 #[test]
1443 fn test_validate_unknown_start_node() {
1444 let graph = DependencyGraph::builder()
1445 .available_actions(["Grep", "Read"])
1446 .start_node("Unknown") .terminal_node("Read")
1448 .build();
1449
1450 let result = graph.validate();
1451 assert!(result.is_err());
1452 assert!(matches!(
1453 result.unwrap_err(),
1454 DependencyGraphError::UnknownAction(name) if name == "Unknown"
1455 ));
1456 }
1457
1458 #[test]
1459 fn test_validate_unknown_terminal_node() {
1460 let graph = DependencyGraph::builder()
1461 .available_actions(["Grep", "Read"])
1462 .start_node("Grep")
1463 .terminal_node("Unknown") .build();
1465
1466 let result = graph.validate();
1467 assert!(result.is_err());
1468 assert!(matches!(
1469 result.unwrap_err(),
1470 DependencyGraphError::UnknownAction(name) if name == "Unknown"
1471 ));
1472 }
1473
1474 #[test]
1475 fn test_validate_valid_graph() {
1476 let graph = DependencyGraph::builder()
1477 .available_actions(["Grep", "Read"])
1478 .edge("Grep", "Read", 0.9)
1479 .start_node("Grep")
1480 .terminal_node("Read")
1481 .build();
1482
1483 assert!(graph.validate().is_ok());
1484 }
1485
1486 #[test]
1491 fn test_start_actions_sorted() {
1492 let graph = DependencyGraph::builder()
1493 .available_actions(["Zebra", "Apple", "Mango"])
1494 .start_nodes(["Zebra", "Apple", "Mango"])
1495 .terminal_node("Zebra")
1496 .build();
1497
1498 let actions = graph.start_actions();
1499 assert_eq!(actions, vec!["Apple", "Mango", "Zebra"]);
1501 }
1502
1503 #[test]
1504 fn test_terminal_actions_sorted() {
1505 let graph = DependencyGraph::builder()
1506 .available_actions(["Zebra", "Apple", "Mango"])
1507 .start_node("Apple")
1508 .terminal_nodes(["Zebra", "Apple", "Mango"])
1509 .build();
1510
1511 let actions = graph.terminal_actions();
1512 assert_eq!(actions, vec!["Apple", "Mango", "Zebra"]);
1514 }
1515}