Skip to main content

graphify_serve/
lib.rs

1//! MCP server for graph queries.
2//!
3//! Provides graph traversal and scoring functions used by the query
4//! engine and MCP protocol server. Port of Python query tools.
5
6pub mod mcp;
7
8use std::collections::{HashMap, HashSet, VecDeque};
9use std::path::Path;
10
11use graphify_core::graph::KnowledgeGraph;
12use serde_json::Value;
13use thiserror::Error;
14
15/// Errors from the server.
16#[derive(Debug, Error)]
17pub enum ServeError {
18    #[error("IO error: {0}")]
19    Io(#[from] std::io::Error),
20
21    #[error("graph load error: {0}")]
22    GraphLoad(String),
23
24    #[error("serialization error: {0}")]
25    Serialization(#[from] serde_json::Error),
26}
27
28/// Score nodes by relevance to search terms.
29///
30/// Returns `(score, node_id)` pairs sorted by descending score.
31/// Scoring: +2.0 for exact label match, +1.0 for label contains,
32/// +0.5 for id contains, plus a small degree-based boost.
33pub fn score_nodes(graph: &KnowledgeGraph, terms: &[String]) -> Vec<(f64, String)> {
34    let lower_terms: Vec<String> = terms.iter().map(|t| t.to_lowercase()).collect();
35
36    let mut scored = Vec::new();
37    for node_id in graph.node_ids() {
38        if let Some(node) = graph.get_node(&node_id) {
39            let label_lower = node.label.to_lowercase();
40            let id_lower = node.id.to_lowercase();
41
42            let mut score: f64 = 0.0;
43
44            for term in &lower_terms {
45                // Exact match in label
46                if label_lower == *term {
47                    score += 2.0;
48                } else if label_lower.contains(term.as_str()) {
49                    score += 1.0;
50                }
51
52                // Match in node ID
53                if id_lower.contains(term.as_str()) {
54                    score += 0.5;
55                }
56            }
57
58            if score > 0.0 {
59                // Boost by degree (well-connected nodes are more relevant)
60                let degree_boost = (graph.degree(&node_id) as f64).ln_1p() * 0.1;
61                score += degree_boost;
62                scored.push((score, node_id.clone()));
63            }
64        }
65    }
66
67    scored.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap_or(std::cmp::Ordering::Equal));
68    scored
69}
70
71/// BFS traversal from start nodes up to a maximum depth.
72///
73/// Returns `(visited_nodes, edges_traversed)` where edges are `(source, target)` pairs.
74pub fn bfs(
75    graph: &KnowledgeGraph,
76    start: &[String],
77    depth: usize,
78) -> (Vec<String>, Vec<(String, String)>) {
79    let mut visited: HashSet<String> = HashSet::new();
80    let mut edges: Vec<(String, String)> = Vec::new();
81    let mut queue: VecDeque<(String, usize)> = VecDeque::new();
82
83    for s in start {
84        if graph.get_node(s).is_some() {
85            visited.insert(s.clone());
86            queue.push_back((s.clone(), 0));
87        }
88    }
89
90    while let Some((current, current_depth)) = queue.pop_front() {
91        if current_depth >= depth {
92            continue;
93        }
94
95        for neighbor_id in graph.neighbor_ids(&current) {
96            edges.push((current.clone(), neighbor_id.clone()));
97
98            if !visited.contains(&neighbor_id) {
99                visited.insert(neighbor_id.clone());
100                queue.push_back((neighbor_id, current_depth + 1));
101            }
102        }
103    }
104
105    let visited_vec: Vec<String> = visited.into_iter().collect();
106    (visited_vec, edges)
107}
108
109/// DFS traversal from start nodes up to a maximum depth.
110///
111/// Returns `(visited_nodes, edges_traversed)` where edges are `(source, target)` pairs.
112pub fn dfs(
113    graph: &KnowledgeGraph,
114    start: &[String],
115    depth: usize,
116) -> (Vec<String>, Vec<(String, String)>) {
117    let mut visited: HashSet<String> = HashSet::new();
118    let mut edges: Vec<(String, String)> = Vec::new();
119    let mut stack: Vec<(String, usize)> = Vec::new();
120
121    for s in start {
122        if graph.get_node(s).is_some() {
123            visited.insert(s.clone());
124            stack.push((s.clone(), 0));
125        }
126    }
127
128    while let Some((current, current_depth)) = stack.pop() {
129        if current_depth >= depth {
130            continue;
131        }
132
133        for neighbor_id in graph.neighbor_ids(&current) {
134            edges.push((current.clone(), neighbor_id.clone()));
135
136            if !visited.contains(&neighbor_id) {
137                visited.insert(neighbor_id.clone());
138                stack.push((neighbor_id, current_depth + 1));
139            }
140        }
141    }
142
143    let visited_vec: Vec<String> = visited.into_iter().collect();
144    (visited_vec, edges)
145}
146
147/// Convert a subgraph (set of nodes and edges) to a text representation
148/// suitable for LLM context windows.
149///
150/// Respects a `token_budget` (approximate: 1 token ≈ 4 chars).
151pub fn subgraph_to_text(
152    graph: &KnowledgeGraph,
153    nodes: &[String],
154    edges: &[(String, String)],
155    token_budget: usize,
156) -> String {
157    let char_budget = token_budget * 4;
158    let mut output = String::with_capacity(char_budget.min(64 * 1024));
159
160    // Header
161    output.push_str(&format!(
162        "=== Knowledge Graph Context ({} nodes, {} edges) ===\n\n",
163        nodes.len(),
164        edges.len()
165    ));
166
167    // Nodes section
168    output.push_str("## Nodes\n\n");
169    for node_id in nodes {
170        if output.len() >= char_budget {
171            output.push_str("\n... (truncated due to token budget)\n");
172            break;
173        }
174
175        if let Some(node) = graph.get_node(node_id) {
176            output.push_str(&format!(
177                "- **{}** [{}] (type: {:?}",
178                node.label, node.id, node.node_type
179            ));
180            if let Some(community) = node.community {
181                output.push_str(&format!(", community: {}", community));
182            }
183            output.push_str(&format!(", file: {})\n", node.source_file));
184        }
185    }
186
187    // Edges section
188    if output.len() < char_budget {
189        output.push_str("\n## Relationships\n\n");
190
191        // Deduplicate edges for display
192        let mut seen: HashSet<(&str, &str)> = HashSet::new();
193        for (src, tgt) in edges {
194            if output.len() >= char_budget {
195                output.push_str("\n... (truncated due to token budget)\n");
196                break;
197            }
198
199            if seen.insert((src.as_str(), tgt.as_str())) {
200                let src_label = graph.get_node(src).map(|n| n.label.as_str()).unwrap_or(src);
201                let tgt_label = graph.get_node(tgt).map(|n| n.label.as_str()).unwrap_or(tgt);
202                output.push_str(&format!("- {} -> {}\n", src_label, tgt_label));
203            }
204        }
205    }
206
207    output
208}
209
210// ---------------------------------------------------------------------------
211// Smart graph summarization
212// ---------------------------------------------------------------------------
213
214/// Abstraction level for graph summaries.
215#[derive(Debug, Clone, Copy, PartialEq, Eq)]
216pub enum SummaryLevel {
217    /// Full detail: all nodes and edges within budget.
218    Detailed,
219    /// Community representatives + cross-community edges.
220    Community,
221    /// Directory-level super-nodes with aggregated dependencies.
222    Architecture,
223}
224
225/// Generate a multi-level graph summary within a token budget.
226///
227/// - `Detailed`: Equivalent to `subgraph_to_text` on the full graph.
228/// - `Community`: One representative node per community (highest degree) + cross-community edges.
229/// - `Architecture`: Groups nodes by directory into super-nodes, merges edges.
230pub fn smart_summary(
231    graph: &KnowledgeGraph,
232    communities: &HashMap<usize, Vec<String>>,
233    level: SummaryLevel,
234    token_budget: usize,
235) -> String {
236    match level {
237        SummaryLevel::Detailed => {
238            let all_nodes = graph.node_ids();
239            let all_edges: Vec<(String, String)> = graph
240                .edges_with_endpoints()
241                .iter()
242                .map(|(s, t, _)| (s.to_string(), t.to_string()))
243                .collect();
244            subgraph_to_text(graph, &all_nodes, &all_edges, token_budget)
245        }
246        SummaryLevel::Community => community_level_summary(graph, communities, token_budget),
247        SummaryLevel::Architecture => architecture_level_summary(graph, token_budget),
248    }
249}
250
251/// Community-level summary: one representative per community + cross-community edges.
252fn community_level_summary(
253    graph: &KnowledgeGraph,
254    communities: &HashMap<usize, Vec<String>>,
255    token_budget: usize,
256) -> String {
257    let char_budget = token_budget * 4;
258    let mut output = String::with_capacity(char_budget.min(64 * 1024));
259
260    // Pick representative (highest degree) for each community
261    let mut representatives: HashMap<usize, (&str, usize)> = HashMap::new();
262    for (&cid, members) in communities {
263        let best = members
264            .iter()
265            .map(|id| (id.as_str(), graph.degree(id)))
266            .max_by_key(|(_, d)| *d)
267            .unwrap_or(("", 0));
268        representatives.insert(cid, best);
269    }
270
271    output.push_str(&format!(
272        "=== Architecture Summary ({} communities, {} total nodes) ===\n\n",
273        communities.len(),
274        graph.node_count()
275    ));
276
277    // Node→community reverse map
278    let mut node_cid: HashMap<&str, usize> = HashMap::new();
279    for (&cid, members) in communities {
280        for m in members {
281            node_cid.insert(m.as_str(), cid);
282        }
283    }
284
285    // Communities
286    output.push_str("## Communities\n\n");
287    let mut sorted_cids: Vec<usize> = communities.keys().copied().collect();
288    sorted_cids.sort();
289    for cid in &sorted_cids {
290        if output.len() >= char_budget {
291            output.push_str("\n... (truncated)\n");
292            break;
293        }
294        let members = &communities[cid];
295        let (rep_id, rep_deg) = representatives[cid];
296        let rep_label = graph
297            .get_node(rep_id)
298            .map(|n| n.label.as_str())
299            .unwrap_or(rep_id);
300        output.push_str(&format!(
301            "- Community {} ({} nodes): representative **{}** (degree {})\n",
302            cid,
303            members.len(),
304            rep_label,
305            rep_deg
306        ));
307    }
308
309    // Cross-community edges
310    output.push_str("\n## Cross-Community Dependencies\n\n");
311    let mut cross_edges: HashMap<(usize, usize), usize> = HashMap::new();
312    for (src, tgt, _) in graph.edges_with_endpoints() {
313        let sc = node_cid.get(src).copied().unwrap_or(usize::MAX);
314        let tc = node_cid.get(tgt).copied().unwrap_or(usize::MAX);
315        if sc != tc && sc != usize::MAX && tc != usize::MAX {
316            let key = if sc < tc { (sc, tc) } else { (tc, sc) };
317            *cross_edges.entry(key).or_default() += 1;
318        }
319    }
320    let mut sorted_cross: Vec<_> = cross_edges.into_iter().collect();
321    sorted_cross.sort_by_key(|(_, count)| std::cmp::Reverse(*count));
322    for ((c1, c2), count) in sorted_cross.iter().take(20) {
323        if output.len() >= char_budget {
324            break;
325        }
326        output.push_str(&format!(
327            "- Community {} ↔ Community {}: {} edges\n",
328            c1, c2, count
329        ));
330    }
331
332    output
333}
334
335/// Architecture-level summary: group by directory, aggregate edges.
336fn architecture_level_summary(graph: &KnowledgeGraph, token_budget: usize) -> String {
337    let char_budget = token_budget * 4;
338    let mut output = String::with_capacity(char_budget.min(64 * 1024));
339
340    // Group nodes by directory
341    let mut dir_nodes: HashMap<String, Vec<&str>> = HashMap::new();
342    for node in graph.nodes() {
343        let dir = std::path::Path::new(&node.source_file)
344            .parent()
345            .and_then(|p| p.to_str())
346            .unwrap_or(".")
347            .to_string();
348        dir_nodes.entry(dir).or_default().push(&node.id);
349    }
350
351    // Node→dir mapping
352    let mut node_dir: HashMap<&str, &str> = HashMap::new();
353    for (dir, nodes) in &dir_nodes {
354        for &nid in nodes {
355            node_dir.insert(nid, dir.as_str());
356        }
357    }
358
359    output.push_str(&format!(
360        "=== Architecture Overview ({} packages/directories) ===\n\n",
361        dir_nodes.len()
362    ));
363
364    // Directory summaries
365    output.push_str("## Packages\n\n");
366    let mut sorted_dirs: Vec<_> = dir_nodes.iter().collect();
367    sorted_dirs.sort_by_key(|(_, nodes)| std::cmp::Reverse(nodes.len()));
368    for (dir, nodes) in sorted_dirs.iter().take(30) {
369        if output.len() >= char_budget {
370            output.push_str("\n... (truncated)\n");
371            break;
372        }
373        output.push_str(&format!("- **{}** ({} entities)\n", dir, nodes.len()));
374    }
375
376    // Inter-directory dependencies
377    output.push_str("\n## Dependencies\n\n");
378    let mut dir_edges: HashMap<(&str, &str), usize> = HashMap::new();
379    for (src, tgt, _) in graph.edges_with_endpoints() {
380        let sd = node_dir.get(src).copied().unwrap_or("?");
381        let td = node_dir.get(tgt).copied().unwrap_or("?");
382        if sd != td {
383            let key = if sd < td { (sd, td) } else { (td, sd) };
384            *dir_edges.entry(key).or_default() += 1;
385        }
386    }
387    let mut sorted_deps: Vec<_> = dir_edges.into_iter().collect();
388    sorted_deps.sort_by_key(|(_, count)| std::cmp::Reverse(*count));
389    for ((d1, d2), count) in sorted_deps.iter().take(20) {
390        if output.len() >= char_budget {
391            break;
392        }
393        output.push_str(&format!("- {} → {}: {} edges\n", d1, d2, count));
394    }
395
396    output
397}
398
399/// Load a knowledge graph from a JSON file.
400pub fn load_graph(graph_path: &Path) -> Result<KnowledgeGraph, ServeError> {
401    let content = std::fs::read_to_string(graph_path)?;
402    let value: Value = serde_json::from_str(&content)?;
403    KnowledgeGraph::from_node_link_json(&value).map_err(|e| ServeError::GraphLoad(e.to_string()))
404}
405
406/// Get basic statistics about the graph.
407pub fn graph_stats(graph: &KnowledgeGraph) -> HashMap<String, Value> {
408    let mut stats = HashMap::new();
409    stats.insert("node_count".to_string(), Value::from(graph.node_count()));
410    stats.insert("edge_count".to_string(), Value::from(graph.edge_count()));
411    stats.insert(
412        "community_count".to_string(),
413        Value::from(graph.communities.len()),
414    );
415
416    // Degree statistics
417    let node_ids = graph.node_ids();
418    if !node_ids.is_empty() {
419        let degrees: Vec<usize> = node_ids.iter().map(|id| graph.degree(id)).collect();
420        let max_degree = degrees.iter().copied().max().unwrap_or(0);
421        let avg_degree = degrees.iter().sum::<usize>() as f64 / degrees.len() as f64;
422        stats.insert("max_degree".to_string(), Value::from(max_degree));
423        stats.insert(
424            "avg_degree".to_string(),
425            Value::from(format!("{:.2}", avg_degree)),
426        );
427    }
428
429    stats
430}
431
432/// Start the MCP server over stdio (JSON-RPC 2.0).
433///
434/// Reads requests from stdin, writes responses to stdout.
435/// This is the entry point called by the CLI `serve` command.
436pub async fn start_server(graph_path: &Path) -> Result<(), ServeError> {
437    // Run the synchronous stdio loop; use spawn_blocking so we don't
438    // block the tokio runtime (though for stdio this is fine).
439    let path = graph_path.to_path_buf();
440    tokio::task::spawn_blocking(move || mcp::run_mcp_server(&path))
441        .await
442        .map_err(|e| ServeError::Io(std::io::Error::other(e)))??;
443    Ok(())
444}
445
446// ---------------------------------------------------------------------------
447// Advanced graph algorithms
448// ---------------------------------------------------------------------------
449
450/// Find all simple paths between `source` and `target` up to `max_length` edges.
451///
452/// Returns a vec of paths, each path being a vec of node IDs.
453/// Limits to at most 50 paths to prevent combinatorial explosion.
454pub fn all_simple_paths(
455    graph: &KnowledgeGraph,
456    source: &str,
457    target: &str,
458    max_length: usize,
459) -> Vec<Vec<String>> {
460    const MAX_PATHS: usize = 50;
461    let mut result: Vec<Vec<String>> = Vec::new();
462    let mut stack: Vec<(String, Vec<String>)> = Vec::new();
463
464    if graph.get_node(source).is_none() || graph.get_node(target).is_none() {
465        return result;
466    }
467
468    stack.push((source.to_string(), vec![source.to_string()]));
469
470    while let Some((current, path)) = stack.pop() {
471        if result.len() >= MAX_PATHS {
472            break;
473        }
474        if current == target && path.len() > 1 {
475            result.push(path);
476            continue;
477        }
478        if path.len() > max_length + 1 {
479            continue;
480        }
481
482        for neighbor_id in graph.neighbor_ids(&current) {
483            if !path.contains(&neighbor_id) {
484                let mut new_path = path.clone();
485                new_path.push(neighbor_id.clone());
486                stack.push((neighbor_id, new_path));
487            }
488        }
489    }
490
491    result.sort_by_key(|p| p.len());
492    result
493}
494
495/// Edge detail in a weighted path: (from_id, to_id, cost, relation).
496pub type EdgeDetail = (String, String, f64, String);
497
498/// Dijkstra shortest path using edge weights.
499///
500/// Cost = 1.0 / edge.weight (higher weight = shorter distance).
501/// Optionally filters edges below `min_confidence` score.
502/// Returns `(path, total_cost, edge_details)` or None if no path exists.
503pub fn dijkstra_path(
504    graph: &KnowledgeGraph,
505    source: &str,
506    target: &str,
507    min_confidence: f64,
508) -> Option<(Vec<String>, f64, Vec<EdgeDetail>)> {
509    use std::cmp::Ordering;
510    use std::collections::BinaryHeap;
511
512    if graph.get_node(source).is_none() || graph.get_node(target).is_none() {
513        return None;
514    }
515    if source == target {
516        return Some((vec![source.to_string()], 0.0, Vec::new()));
517    }
518
519    // Build adjacency with weights from edges
520    let mut adj: HashMap<String, Vec<(String, f64, String)>> = HashMap::new();
521    for (src, tgt, edge) in graph.edges_with_endpoints() {
522        if edge.confidence_score < min_confidence {
523            continue;
524        }
525        let cost = if edge.weight > 0.0 {
526            1.0 / edge.weight
527        } else {
528            f64::MAX
529        };
530        adj.entry(src.to_string()).or_default().push((
531            tgt.to_string(),
532            cost,
533            edge.relation.clone(),
534        ));
535        adj.entry(tgt.to_string()).or_default().push((
536            src.to_string(),
537            cost,
538            edge.relation.clone(),
539        ));
540    }
541
542    #[derive(PartialEq)]
543    struct State {
544        cost: f64,
545        node: String,
546    }
547    impl Eq for State {}
548    impl PartialOrd for State {
549        fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
550            Some(self.cmp(other))
551        }
552    }
553    impl Ord for State {
554        fn cmp(&self, other: &Self) -> Ordering {
555            other
556                .cost
557                .partial_cmp(&self.cost)
558                .unwrap_or(Ordering::Equal)
559        }
560    }
561
562    let mut dist: HashMap<String, f64> = HashMap::new();
563    let mut prev: HashMap<String, (String, f64, String)> = HashMap::new();
564    let mut heap = BinaryHeap::new();
565
566    dist.insert(source.to_string(), 0.0);
567    heap.push(State {
568        cost: 0.0,
569        node: source.to_string(),
570    });
571
572    while let Some(State { cost, node }) = heap.pop() {
573        if node == target {
574            break;
575        }
576        if cost > *dist.get(&node).unwrap_or(&f64::MAX) {
577            continue;
578        }
579        if let Some(neighbors) = adj.get(&node) {
580            for (next, edge_cost, relation) in neighbors {
581                let new_cost = cost + edge_cost;
582                if new_cost < *dist.get(next).unwrap_or(&f64::MAX) {
583                    dist.insert(next.clone(), new_cost);
584                    prev.insert(next.clone(), (node.clone(), *edge_cost, relation.clone()));
585                    heap.push(State {
586                        cost: new_cost,
587                        node: next.clone(),
588                    });
589                }
590            }
591        }
592    }
593
594    // Reconstruct path
595    if !prev.contains_key(target) {
596        return None;
597    }
598
599    let mut path = vec![target.to_string()];
600    let mut edge_details: Vec<(String, String, f64, String)> = Vec::new();
601    let mut current = target.to_string();
602    while let Some((from, cost, relation)) = prev.get(&current) {
603        edge_details.push((from.clone(), current.clone(), *cost, relation.clone()));
604        path.push(from.clone());
605        current = from.clone();
606    }
607    path.reverse();
608    edge_details.reverse();
609
610    let total_cost = *dist.get(target).unwrap_or(&f64::MAX);
611    Some((path, total_cost, edge_details))
612}
613
614// ---------------------------------------------------------------------------
615// Tests
616// ---------------------------------------------------------------------------
617
618#[cfg(test)]
619mod tests {
620    use super::*;
621    use graphify_core::confidence::Confidence;
622    use graphify_core::model::{GraphEdge, GraphNode, NodeType};
623
624    fn make_node(id: &str, label: &str) -> GraphNode {
625        GraphNode {
626            id: id.into(),
627            label: label.into(),
628            source_file: "test.rs".into(),
629            source_location: None,
630            node_type: NodeType::Class,
631            community: None,
632            extra: HashMap::new(),
633        }
634    }
635
636    fn make_edge(src: &str, tgt: &str) -> GraphEdge {
637        GraphEdge {
638            source: src.into(),
639            target: tgt.into(),
640            relation: "calls".into(),
641            confidence: Confidence::Extracted,
642            confidence_score: 1.0,
643            source_file: "test.rs".into(),
644            source_location: None,
645            weight: 1.0,
646            extra: HashMap::new(),
647        }
648    }
649
650    fn make_test_graph() -> KnowledgeGraph {
651        let mut g = KnowledgeGraph::new();
652        g.add_node(make_node("auth", "AuthService")).unwrap();
653        g.add_node(make_node("user", "UserManager")).unwrap();
654        g.add_node(make_node("db", "Database")).unwrap();
655        g.add_node(make_node("cache", "CacheLayer")).unwrap();
656        g.add_edge(make_edge("auth", "user")).unwrap();
657        g.add_edge(make_edge("auth", "db")).unwrap();
658        g.add_edge(make_edge("user", "db")).unwrap();
659        g.add_edge(make_edge("user", "cache")).unwrap();
660        g
661    }
662
663    #[test]
664    fn test_score_nodes_basic() {
665        let g = make_test_graph();
666        let results = score_nodes(&g, &["auth".to_string()]);
667        assert!(!results.is_empty());
668        // "auth" node should score highest
669        let top_id = &results[0].1;
670        assert_eq!(top_id, "auth");
671    }
672
673    #[test]
674    fn test_score_nodes_no_match() {
675        let g = make_test_graph();
676        let results = score_nodes(&g, &["nonexistent".to_string()]);
677        assert!(results.is_empty());
678    }
679
680    #[test]
681    fn test_score_nodes_multiple_terms() {
682        let g = make_test_graph();
683        let results = score_nodes(&g, &["user".to_string(), "manager".to_string()]);
684        assert!(!results.is_empty());
685        assert!(results.iter().any(|(_, id)| id == "user"));
686    }
687
688    #[test]
689    fn test_bfs_depth_0() {
690        let g = make_test_graph();
691        let (nodes, edges) = bfs(&g, &["auth".to_string()], 0);
692        assert_eq!(nodes.len(), 1);
693        assert!(edges.is_empty());
694    }
695
696    #[test]
697    fn test_bfs_depth_1() {
698        let g = make_test_graph();
699        let (nodes, edges) = bfs(&g, &["auth".to_string()], 1);
700        // auth -> user, auth -> db
701        assert!(nodes.len() >= 3); // auth, user, db
702        assert!(!edges.is_empty());
703    }
704
705    #[test]
706    fn test_bfs_depth_2() {
707        let g = make_test_graph();
708        let (nodes, _edges) = bfs(&g, &["auth".to_string()], 2);
709        // Should reach all 4 nodes
710        assert_eq!(nodes.len(), 4);
711    }
712
713    #[test]
714    fn test_dfs_depth_1() {
715        let g = make_test_graph();
716        let (nodes, edges) = dfs(&g, &["auth".to_string()], 1);
717        assert!(nodes.len() >= 3);
718        assert!(!edges.is_empty());
719    }
720
721    #[test]
722    fn test_bfs_nonexistent_start() {
723        let g = make_test_graph();
724        let (nodes, edges) = bfs(&g, &["nonexistent".to_string()], 3);
725        assert!(nodes.is_empty());
726        assert!(edges.is_empty());
727    }
728
729    #[test]
730    fn test_subgraph_to_text() {
731        let g = make_test_graph();
732        let nodes = vec!["auth".to_string(), "user".to_string()];
733        let edges = vec![("auth".to_string(), "user".to_string())];
734        let text = subgraph_to_text(&g, &nodes, &edges, 1000);
735
736        assert!(text.contains("Knowledge Graph Context"));
737        assert!(text.contains("AuthService"));
738        assert!(text.contains("UserManager"));
739        assert!(text.contains("Relationships"));
740    }
741
742    #[test]
743    fn test_subgraph_to_text_budget() {
744        let g = make_test_graph();
745        let nodes: Vec<String> = g.node_ids();
746        let edges = vec![
747            ("auth".to_string(), "user".to_string()),
748            ("auth".to_string(), "db".to_string()),
749        ];
750        // Very small budget
751        let text = subgraph_to_text(&g, &nodes, &edges, 10);
752        assert!(text.contains("truncated") || text.len() < 200);
753    }
754
755    #[test]
756    fn test_graph_stats() {
757        let g = make_test_graph();
758        let stats = graph_stats(&g);
759        assert_eq!(stats["node_count"], 4);
760        assert_eq!(stats["edge_count"], 4);
761    }
762
763    #[test]
764    fn test_bfs_multiple_starts() {
765        let g = make_test_graph();
766        let (nodes, _) = bfs(&g, &["auth".to_string(), "cache".to_string()], 1);
767        assert!(nodes.len() >= 4);
768    }
769
770    // -- all_simple_paths tests --
771
772    #[test]
773    fn test_all_simple_paths_direct() {
774        let g = make_test_graph();
775        let paths = all_simple_paths(&g, "auth", "user", 4);
776        assert!(!paths.is_empty());
777        // Direct edge exists: auth → user
778        assert!(paths.iter().any(|p| p.len() == 2));
779    }
780
781    #[test]
782    fn test_all_simple_paths_indirect() {
783        let g = make_test_graph();
784        // auth → db has direct path (len 2) and indirect auth → user → db (len 3)
785        let paths = all_simple_paths(&g, "auth", "db", 4);
786        assert!(
787            paths.len() >= 2,
788            "should find multiple paths, got {}",
789            paths.len()
790        );
791    }
792
793    #[test]
794    fn test_all_simple_paths_no_path() {
795        let mut g = KnowledgeGraph::new();
796        g.add_node(make_node("a", "A")).unwrap();
797        g.add_node(make_node("b", "B")).unwrap();
798        // No edge between a and b
799        let paths = all_simple_paths(&g, "a", "b", 4);
800        assert!(paths.is_empty());
801    }
802
803    #[test]
804    fn test_all_simple_paths_nonexistent_node() {
805        let g = make_test_graph();
806        let paths = all_simple_paths(&g, "auth", "nonexistent", 4);
807        assert!(paths.is_empty());
808    }
809
810    #[test]
811    fn test_all_simple_paths_sorted_by_length() {
812        let g = make_test_graph();
813        let paths = all_simple_paths(&g, "auth", "cache", 5);
814        for w in paths.windows(2) {
815            assert!(w[0].len() <= w[1].len(), "paths should be sorted by length");
816        }
817    }
818
819    // -- dijkstra_path tests --
820
821    #[test]
822    fn test_dijkstra_direct_path() {
823        let g = make_test_graph();
824        let result = dijkstra_path(&g, "auth", "user", 0.0);
825        assert!(result.is_some());
826        let (path, cost, edges) = result.unwrap();
827        assert_eq!(path.first().unwrap(), "auth");
828        assert_eq!(path.last().unwrap(), "user");
829        assert!(cost > 0.0);
830        assert!(!edges.is_empty());
831    }
832
833    #[test]
834    fn test_dijkstra_same_node() {
835        let g = make_test_graph();
836        let result = dijkstra_path(&g, "auth", "auth", 0.0);
837        assert!(result.is_some());
838        let (path, cost, _) = result.unwrap();
839        assert_eq!(path.len(), 1);
840        assert!((cost - 0.0).abs() < f64::EPSILON);
841    }
842
843    #[test]
844    fn test_dijkstra_no_path() {
845        let mut g = KnowledgeGraph::new();
846        g.add_node(make_node("a", "A")).unwrap();
847        g.add_node(make_node("b", "B")).unwrap();
848        let result = dijkstra_path(&g, "a", "b", 0.0);
849        assert!(result.is_none());
850    }
851
852    #[test]
853    fn test_dijkstra_nonexistent_node() {
854        let g = make_test_graph();
855        assert!(dijkstra_path(&g, "auth", "nonexistent", 0.0).is_none());
856    }
857
858    #[test]
859    fn test_dijkstra_min_confidence_filter() {
860        // Create graph with mixed confidence edges
861        let mut g = KnowledgeGraph::new();
862        g.add_node(make_node("a", "A")).unwrap();
863        g.add_node(make_node("b", "B")).unwrap();
864        g.add_node(make_node("c", "C")).unwrap();
865
866        // a→b: low confidence (0.3), a→c→b: high confidence (1.0)
867        let mut low_edge = make_edge("a", "b");
868        low_edge.confidence_score = 0.3;
869        g.add_edge(low_edge).unwrap();
870
871        let mut high1 = make_edge("a", "c");
872        high1.confidence_score = 1.0;
873        g.add_edge(high1).unwrap();
874
875        let mut high2 = make_edge("c", "b");
876        high2.confidence_score = 1.0;
877        g.add_edge(high2).unwrap();
878
879        // With min_confidence 0.5, should skip a→b and go a→c→b
880        let result = dijkstra_path(&g, "a", "b", 0.5);
881        assert!(result.is_some());
882        let (path, _, _) = result.unwrap();
883        assert_eq!(path.len(), 3, "should go through c, got path: {path:?}");
884    }
885}