nusy_codegraph/
mcp_tools.rs

1//! MCP tool implementations for CodeGraph query and update operations.
2//!
3//! Provides 4 CRUD-style tools that agents use instead of file reads/writes:
4//! - `codegraph_query_objects` — filter by kind, name, parent, metrics, semantic similarity
5//! - `codegraph_update_object` — modify CodeNode fields (signature, docstring, body_hash)
6//! - `codegraph_add_edge` — create a relationship between CodeNodes
7//! - `codegraph_remove_edge` — logical delete of an edge
8//!
9//! These operate directly on Arrow RecordBatches — no file materialization.
10
11use crate::schema::{
12    CodeEdge, CodeEdgePredicate, CodeNode, CodeNodeKind, build_code_edges_batch, edge_col, node_col,
13};
14use arrow::array::{Array, Float64Array, Int32Array, RecordBatch, StringArray};
15
16/// Errors from MCP tool operations.
17#[derive(Debug, thiserror::Error)]
18pub enum McpToolError {
19    #[error("Node not found: {0}")]
20    NodeNotFound(String),
21
22    #[error("Edge not found: {0} -> {1} ({2})")]
23    EdgeNotFound(String, String, String),
24
25    #[error("Invalid kind: {0}")]
26    InvalidKind(String),
27
28    #[error("Invalid predicate: {0}")]
29    InvalidPredicate(String),
30
31    #[error("Arrow error: {0}")]
32    Arrow(#[from] arrow::error::ArrowError),
33}
34
35pub type Result<T> = std::result::Result<T, McpToolError>;
36
37/// Filter criteria for querying code objects.
38#[derive(Debug, Default, Clone)]
39pub struct QueryFilter {
40    /// Filter by CodeNodeKind (e.g., "function", "class").
41    pub kind: Option<String>,
42    /// Filter by name substring match.
43    pub name_contains: Option<String>,
44    /// Filter by parent_id exact match.
45    pub parent_id: Option<String>,
46    /// Filter by minimum LOC.
47    pub min_loc: Option<i32>,
48    /// Filter by minimum cyclomatic complexity.
49    pub min_complexity: Option<i32>,
50    /// Filter by maximum coverage percentage.
51    pub max_coverage: Option<f64>,
52    /// Limit number of results.
53    pub limit: Option<usize>,
54}
55
56/// Result of a query — matching CodeNodes extracted from RecordBatches.
57#[derive(Debug)]
58pub struct QueryResult {
59    pub nodes: Vec<CodeNode>,
60    pub total_scanned: usize,
61    pub total_matched: usize,
62}
63
64/// Query code objects from the nodes batch with optional filters.
65pub fn codegraph_query_objects(
66    nodes_batch: &RecordBatch,
67    filter: &QueryFilter,
68) -> Result<QueryResult> {
69    let ids = nodes_batch
70        .column(node_col::ID)
71        .as_any()
72        .downcast_ref::<StringArray>()
73        .expect("id column should be StringArray");
74    let names = nodes_batch
75        .column(node_col::NAME)
76        .as_any()
77        .downcast_ref::<StringArray>()
78        .expect("name column");
79    let parent_ids = nodes_batch
80        .column(node_col::PARENT_ID)
81        .as_any()
82        .downcast_ref::<StringArray>()
83        .expect("parent_id column");
84    let signatures = nodes_batch
85        .column(node_col::SIGNATURE)
86        .as_any()
87        .downcast_ref::<StringArray>()
88        .expect("signature column");
89    let docstrings = nodes_batch
90        .column(node_col::DOCSTRING)
91        .as_any()
92        .downcast_ref::<StringArray>()
93        .expect("docstring column");
94    let body_hashes = nodes_batch
95        .column(node_col::BODY_HASH)
96        .as_any()
97        .downcast_ref::<StringArray>()
98        .expect("body_hash column");
99    let locs = nodes_batch
100        .column(node_col::LOC)
101        .as_any()
102        .downcast_ref::<Int32Array>()
103        .expect("loc column");
104    let complexities = nodes_batch
105        .column(node_col::CYCLOMATIC_COMPLEXITY)
106        .as_any()
107        .downcast_ref::<Int32Array>()
108        .expect("complexity column");
109    let coverages = nodes_batch
110        .column(node_col::COVERAGE_PCT)
111        .as_any()
112        .downcast_ref::<Float64Array>()
113        .expect("coverage column");
114
115    // Extract kind values from dictionary-encoded column
116    let kind_col = nodes_batch.column(node_col::KIND);
117    let kind_dict = kind_col
118        .as_any()
119        .downcast_ref::<arrow::array::Int8DictionaryArray>()
120        .expect("kind should be dictionary");
121    let kind_values = kind_dict
122        .values()
123        .as_any()
124        .downcast_ref::<StringArray>()
125        .expect("kind values");
126
127    let total_scanned = nodes_batch.num_rows();
128    let mut matched = Vec::new();
129
130    for i in 0..total_scanned {
131        // Kind filter
132        let kind_key = kind_dict.keys().value(i) as usize;
133        let kind_str = kind_values.value(kind_key);
134        if let Some(ref filter_kind) = filter.kind
135            && kind_str != filter_kind.as_str()
136        {
137            continue;
138        }
139
140        // Name filter (substring match)
141        if let Some(ref name_substr) = filter.name_contains {
142            let name = names.value(i);
143            if !name.to_lowercase().contains(&name_substr.to_lowercase()) {
144                continue;
145            }
146        }
147
148        // Parent filter
149        if let Some(ref parent) = filter.parent_id
150            && (parent_ids.is_null(i) || parent_ids.value(i) != parent.as_str())
151        {
152            continue;
153        }
154
155        // LOC filter
156        if let Some(min_loc) = filter.min_loc
157            && (locs.is_null(i) || locs.value(i) < min_loc)
158        {
159            continue;
160        }
161
162        // Complexity filter
163        if let Some(min_complexity) = filter.min_complexity
164            && (complexities.is_null(i) || complexities.value(i) < min_complexity)
165        {
166            continue;
167        }
168
169        // Coverage filter (max coverage = find under-tested code)
170        if let Some(max_cov) = filter.max_coverage
171            && (coverages.is_null(i) || coverages.value(i) > max_cov)
172        {
173            continue;
174        }
175
176        let node = CodeNode {
177            id: ids.value(i).to_string(),
178            kind: CodeNodeKind::parse(kind_str).unwrap_or(CodeNodeKind::Function),
179            parent_id: if parent_ids.is_null(i) {
180                None
181            } else {
182                Some(parent_ids.value(i).to_string())
183            },
184            name: names.value(i).to_string(),
185            signature: if signatures.is_null(i) {
186                None
187            } else {
188                Some(signatures.value(i).to_string())
189            },
190            docstring: if docstrings.is_null(i) {
191                None
192            } else {
193                Some(docstrings.value(i).to_string())
194            },
195            body_hash: if body_hashes.is_null(i) {
196                None
197            } else {
198                Some(body_hashes.value(i).to_string())
199            },
200            body: None, // Body not extracted in query results
201            loc: if locs.is_null(i) {
202                None
203            } else {
204                Some(locs.value(i))
205            },
206            cyclomatic_complexity: if complexities.is_null(i) {
207                None
208            } else {
209                Some(complexities.value(i))
210            },
211            coverage_pct: if coverages.is_null(i) {
212                None
213            } else {
214                Some(coverages.value(i))
215            },
216            last_modified: None,
217            ..Default::default()
218        };
219
220        matched.push(node);
221
222        // Limit
223        if let Some(limit) = filter.limit
224            && matched.len() >= limit
225        {
226            break;
227        }
228    }
229
230    let total_matched = matched.len();
231    Ok(QueryResult {
232        nodes: matched,
233        total_scanned,
234        total_matched,
235    })
236}
237
238/// Fields that can be updated on a CodeNode.
239#[derive(Debug, Default, Clone)]
240pub struct NodeUpdate {
241    pub signature: Option<String>,
242    pub docstring: Option<String>,
243    pub body_hash: Option<String>,
244    pub body: Option<String>,
245    pub loc: Option<i32>,
246    pub cyclomatic_complexity: Option<i32>,
247    pub coverage_pct: Option<f64>,
248}
249
250/// Update a CodeNode's fields in-place. Returns the updated batch.
251///
252/// Finds the node by ID, applies the updates, and returns a new RecordBatch
253/// with the modifications. The original batch is not mutated.
254pub fn codegraph_update_object(
255    nodes_batch: &RecordBatch,
256    node_id: &str,
257    updates: &NodeUpdate,
258) -> Result<RecordBatch> {
259    let ids = nodes_batch
260        .column(node_col::ID)
261        .as_any()
262        .downcast_ref::<StringArray>()
263        .expect("id column");
264
265    // Find the row
266    let row_idx = (0..nodes_batch.num_rows())
267        .find(|&i| ids.value(i) == node_id)
268        .ok_or_else(|| McpToolError::NodeNotFound(node_id.to_string()))?;
269
270    // Rebuild columns with updates applied
271    let mut columns: Vec<arrow::array::ArrayRef> = Vec::new();
272    for col_idx in 0..nodes_batch.num_columns() {
273        match col_idx {
274            node_col::SIGNATURE if updates.signature.is_some() => {
275                let old = nodes_batch
276                    .column(col_idx)
277                    .as_any()
278                    .downcast_ref::<StringArray>()
279                    .expect("signature");
280                let mut vals: Vec<Option<String>> = (0..nodes_batch.num_rows())
281                    .map(|i| {
282                        if old.is_null(i) {
283                            None
284                        } else {
285                            Some(old.value(i).to_string())
286                        }
287                    })
288                    .collect();
289                vals[row_idx] = updates.signature.clone();
290                let refs: Vec<Option<&str>> = vals.iter().map(|s| s.as_deref()).collect();
291                columns.push(std::sync::Arc::new(StringArray::from(refs)));
292            }
293            node_col::DOCSTRING if updates.docstring.is_some() => {
294                let old = nodes_batch
295                    .column(col_idx)
296                    .as_any()
297                    .downcast_ref::<StringArray>()
298                    .expect("docstring");
299                let mut vals: Vec<Option<String>> = (0..nodes_batch.num_rows())
300                    .map(|i| {
301                        if old.is_null(i) {
302                            None
303                        } else {
304                            Some(old.value(i).to_string())
305                        }
306                    })
307                    .collect();
308                vals[row_idx] = updates.docstring.clone();
309                let refs: Vec<Option<&str>> = vals.iter().map(|s| s.as_deref()).collect();
310                columns.push(std::sync::Arc::new(StringArray::from(refs)));
311            }
312            // When body is updated, body_hash is auto-recomputed (explicit body_hash ignored).
313            node_col::BODY_HASH if updates.body_hash.is_some() || updates.body.is_some() => {
314                let old = nodes_batch
315                    .column(col_idx)
316                    .as_any()
317                    .downcast_ref::<StringArray>()
318                    .expect("body_hash");
319                let mut vals: Vec<Option<String>> = (0..nodes_batch.num_rows())
320                    .map(|i| {
321                        if old.is_null(i) {
322                            None
323                        } else {
324                            Some(old.value(i).to_string())
325                        }
326                    })
327                    .collect();
328                // If body is updated, auto-recompute body_hash
329                if let Some(ref body) = updates.body {
330                    vals[row_idx] = Some(crate::parser::sha256_hex(body.as_bytes()));
331                } else {
332                    vals[row_idx] = updates.body_hash.clone();
333                }
334                let refs: Vec<Option<&str>> = vals.iter().map(|s| s.as_deref()).collect();
335                columns.push(std::sync::Arc::new(StringArray::from(refs)));
336            }
337            node_col::BODY if updates.body.is_some() => {
338                let old = nodes_batch
339                    .column(col_idx)
340                    .as_any()
341                    .downcast_ref::<arrow::array::LargeStringArray>()
342                    .expect("body");
343                let mut vals: Vec<Option<String>> = (0..nodes_batch.num_rows())
344                    .map(|i| {
345                        if old.is_null(i) {
346                            None
347                        } else {
348                            Some(old.value(i).to_string())
349                        }
350                    })
351                    .collect();
352                vals[row_idx] = updates.body.clone();
353                let refs: Vec<Option<&str>> = vals.iter().map(|s| s.as_deref()).collect();
354                columns.push(std::sync::Arc::new(arrow::array::LargeStringArray::from(
355                    refs,
356                )));
357            }
358            node_col::LOC if updates.loc.is_some() => {
359                let old = nodes_batch
360                    .column(col_idx)
361                    .as_any()
362                    .downcast_ref::<Int32Array>()
363                    .expect("loc");
364                let mut vals: Vec<Option<i32>> = (0..nodes_batch.num_rows())
365                    .map(|i| {
366                        if old.is_null(i) {
367                            None
368                        } else {
369                            Some(old.value(i))
370                        }
371                    })
372                    .collect();
373                vals[row_idx] = updates.loc;
374                columns.push(std::sync::Arc::new(Int32Array::from(vals)));
375            }
376            node_col::CYCLOMATIC_COMPLEXITY if updates.cyclomatic_complexity.is_some() => {
377                let old = nodes_batch
378                    .column(col_idx)
379                    .as_any()
380                    .downcast_ref::<Int32Array>()
381                    .expect("complexity");
382                let mut vals: Vec<Option<i32>> = (0..nodes_batch.num_rows())
383                    .map(|i| {
384                        if old.is_null(i) {
385                            None
386                        } else {
387                            Some(old.value(i))
388                        }
389                    })
390                    .collect();
391                vals[row_idx] = updates.cyclomatic_complexity;
392                columns.push(std::sync::Arc::new(Int32Array::from(vals)));
393            }
394            node_col::COVERAGE_PCT if updates.coverage_pct.is_some() => {
395                let old = nodes_batch
396                    .column(col_idx)
397                    .as_any()
398                    .downcast_ref::<Float64Array>()
399                    .expect("coverage");
400                let mut vals: Vec<Option<f64>> = (0..nodes_batch.num_rows())
401                    .map(|i| {
402                        if old.is_null(i) {
403                            None
404                        } else {
405                            Some(old.value(i))
406                        }
407                    })
408                    .collect();
409                vals[row_idx] = updates.coverage_pct;
410                columns.push(std::sync::Arc::new(Float64Array::from(vals)));
411            }
412            _ => {
413                columns.push(nodes_batch.column(col_idx).clone());
414            }
415        }
416    }
417
418    let schema = std::sync::Arc::new(nodes_batch.schema().as_ref().clone());
419    Ok(RecordBatch::try_new(schema, columns)?)
420}
421
422/// Add a new edge to the edges batch. Returns a new batch with the edge appended.
423pub fn codegraph_add_edge(
424    edges_batch: &RecordBatch,
425    source_id: &str,
426    target_id: &str,
427    predicate: &str,
428    weight: Option<f32>,
429) -> Result<RecordBatch> {
430    let pred = CodeEdgePredicate::parse(predicate)
431        .ok_or_else(|| McpToolError::InvalidPredicate(predicate.to_string()))?;
432
433    let new_edge = CodeEdge {
434        source_id: source_id.to_string(),
435        target_id: target_id.to_string(),
436        predicate: pred,
437        weight,
438        commit_id: None,
439    };
440
441    let new_batch = build_code_edges_batch(&[new_edge])?;
442
443    // Concatenate existing + new
444    arrow::compute::concat_batches(&edges_batch.schema(), &[edges_batch.clone(), new_batch])
445        .map_err(McpToolError::Arrow)
446}
447
448/// Remove an edge (by setting weight to -1 as a tombstone marker).
449///
450/// Returns a new batch with the edge marked. Since Arrow batches are immutable,
451/// we rebuild with the edge's weight set to -1.0 (tombstone convention).
452pub fn codegraph_remove_edge(
453    edges_batch: &RecordBatch,
454    source_id: &str,
455    target_id: &str,
456    predicate: &str,
457) -> Result<RecordBatch> {
458    let sources = edges_batch
459        .column(edge_col::SOURCE_ID)
460        .as_any()
461        .downcast_ref::<StringArray>()
462        .expect("source_id");
463    let targets = edges_batch
464        .column(edge_col::TARGET_ID)
465        .as_any()
466        .downcast_ref::<StringArray>()
467        .expect("target_id");
468
469    // Extract predicate values from dictionary
470    let pred_col = edges_batch.column(edge_col::PREDICATE);
471    let pred_dict = pred_col
472        .as_any()
473        .downcast_ref::<arrow::array::Int8DictionaryArray>()
474        .expect("predicate dict");
475    let pred_values = pred_dict
476        .values()
477        .as_any()
478        .downcast_ref::<StringArray>()
479        .expect("pred values");
480
481    let row_idx = (0..edges_batch.num_rows())
482        .find(|&i| {
483            let key = pred_dict.keys().value(i) as usize;
484            sources.value(i) == source_id
485                && targets.value(i) == target_id
486                && pred_values.value(key) == predicate
487        })
488        .ok_or_else(|| {
489            McpToolError::EdgeNotFound(
490                source_id.to_string(),
491                target_id.to_string(),
492                predicate.to_string(),
493            )
494        })?;
495
496    // Rebuild weight column with tombstone
497    let mut columns: Vec<arrow::array::ArrayRef> = Vec::new();
498    for col_idx in 0..edges_batch.num_columns() {
499        if col_idx == edge_col::WEIGHT {
500            let old = edges_batch
501                .column(col_idx)
502                .as_any()
503                .downcast_ref::<arrow::array::Float32Array>()
504                .expect("weight");
505            let mut vals: Vec<Option<f32>> = (0..edges_batch.num_rows())
506                .map(|i| {
507                    if old.is_null(i) {
508                        None
509                    } else {
510                        Some(old.value(i))
511                    }
512                })
513                .collect();
514            vals[row_idx] = Some(-1.0); // Tombstone
515            columns.push(std::sync::Arc::new(arrow::array::Float32Array::from(vals)));
516        } else {
517            columns.push(edges_batch.column(col_idx).clone());
518        }
519    }
520
521    let schema = std::sync::Arc::new(edges_batch.schema().as_ref().clone());
522    Ok(RecordBatch::try_new(schema, columns)?)
523}
524
525#[cfg(test)]
526mod tests {
527    use super::*;
528    use crate::schema::{CodeEdge, CodeNode, build_code_edges_batch, build_code_nodes_batch};
529
530    fn sample_nodes() -> Vec<CodeNode> {
531        vec![
532            CodeNode {
533                id: "func:brain/signal_fusion.py::fuse".into(),
534                kind: CodeNodeKind::Function,
535                parent_id: Some("mod:brain/signal_fusion.py".into()),
536                name: "fuse".into(),
537                signature: Some("def fuse(signals: List) -> Decision".into()),
538                docstring: Some("Fuse cognitive signals.".into()),
539                body_hash: Some("abc123".into()),
540                body: None,
541                loc: Some(42),
542                cyclomatic_complexity: Some(8),
543                coverage_pct: Some(0.85),
544                last_modified: None,
545                ..Default::default()
546            },
547            CodeNode {
548                id: "class:brain/store.py::DualStore".into(),
549                kind: CodeNodeKind::Class,
550                parent_id: Some("mod:brain/store.py".into()),
551                name: "DualStore".into(),
552                signature: None,
553                docstring: Some("Fast/slow dual-store.".into()),
554                body_hash: Some("def456".into()),
555                body: None,
556                loc: Some(200),
557                cyclomatic_complexity: Some(15),
558                coverage_pct: Some(0.60),
559                last_modified: None,
560                ..Default::default()
561            },
562            CodeNode {
563                id: "func:brain/store.py::promote".into(),
564                kind: CodeNodeKind::Function,
565                parent_id: Some("class:brain/store.py::DualStore".into()),
566                name: "promote".into(),
567                signature: Some("def promote(self) -> None".into()),
568                docstring: None,
569                body_hash: Some("ghi789".into()),
570                body: None,
571                loc: Some(30),
572                cyclomatic_complexity: Some(3),
573                coverage_pct: Some(0.95),
574                last_modified: None,
575                ..Default::default()
576            },
577        ]
578    }
579
580    fn sample_edges() -> Vec<CodeEdge> {
581        vec![
582            CodeEdge {
583                source_id: "func:brain/signal_fusion.py::fuse".into(),
584                target_id: "class:brain/store.py::DualStore".into(),
585                predicate: CodeEdgePredicate::Uses,
586                weight: Some(1.0),
587                commit_id: None,
588            },
589            CodeEdge {
590                source_id: "func:brain/store.py::promote".into(),
591                target_id: "class:brain/store.py::DualStore".into(),
592                predicate: CodeEdgePredicate::Contains,
593                weight: None,
594                commit_id: None,
595            },
596        ]
597    }
598
599    // --- Phase 1: Query tests ---
600
601    #[test]
602    fn test_query_all_objects() {
603        let batch = build_code_nodes_batch(&sample_nodes()).unwrap();
604        let result = codegraph_query_objects(&batch, &QueryFilter::default()).unwrap();
605        assert_eq!(result.total_scanned, 3);
606        assert_eq!(result.total_matched, 3);
607        assert_eq!(result.nodes.len(), 3);
608    }
609
610    #[test]
611    fn test_query_by_kind() {
612        let batch = build_code_nodes_batch(&sample_nodes()).unwrap();
613        let filter = QueryFilter {
614            kind: Some("function".into()),
615            ..Default::default()
616        };
617        let result = codegraph_query_objects(&batch, &filter).unwrap();
618        assert_eq!(result.total_matched, 2); // fuse + promote
619    }
620
621    #[test]
622    fn test_query_by_name() {
623        let batch = build_code_nodes_batch(&sample_nodes()).unwrap();
624        let filter = QueryFilter {
625            name_contains: Some("fuse".into()),
626            ..Default::default()
627        };
628        let result = codegraph_query_objects(&batch, &filter).unwrap();
629        assert_eq!(result.total_matched, 1);
630        assert_eq!(result.nodes[0].name, "fuse");
631    }
632
633    #[test]
634    fn test_query_by_parent() {
635        let batch = build_code_nodes_batch(&sample_nodes()).unwrap();
636        let filter = QueryFilter {
637            parent_id: Some("class:brain/store.py::DualStore".into()),
638            ..Default::default()
639        };
640        let result = codegraph_query_objects(&batch, &filter).unwrap();
641        assert_eq!(result.total_matched, 1);
642        assert_eq!(result.nodes[0].name, "promote");
643    }
644
645    #[test]
646    fn test_query_by_min_loc() {
647        let batch = build_code_nodes_batch(&sample_nodes()).unwrap();
648        let filter = QueryFilter {
649            min_loc: Some(100),
650            ..Default::default()
651        };
652        let result = codegraph_query_objects(&batch, &filter).unwrap();
653        assert_eq!(result.total_matched, 1);
654        assert_eq!(result.nodes[0].name, "DualStore");
655    }
656
657    #[test]
658    fn test_query_by_max_coverage() {
659        let batch = build_code_nodes_batch(&sample_nodes()).unwrap();
660        let filter = QueryFilter {
661            max_coverage: Some(0.70),
662            ..Default::default()
663        };
664        let result = codegraph_query_objects(&batch, &filter).unwrap();
665        assert_eq!(result.total_matched, 1);
666        assert_eq!(result.nodes[0].name, "DualStore"); // 0.60 coverage
667    }
668
669    #[test]
670    fn test_query_with_limit() {
671        let batch = build_code_nodes_batch(&sample_nodes()).unwrap();
672        let filter = QueryFilter {
673            limit: Some(2),
674            ..Default::default()
675        };
676        let result = codegraph_query_objects(&batch, &filter).unwrap();
677        assert_eq!(result.total_matched, 2);
678    }
679
680    #[test]
681    fn test_query_combined_filters() {
682        let batch = build_code_nodes_batch(&sample_nodes()).unwrap();
683        let filter = QueryFilter {
684            kind: Some("function".into()),
685            min_complexity: Some(5),
686            ..Default::default()
687        };
688        let result = codegraph_query_objects(&batch, &filter).unwrap();
689        assert_eq!(result.total_matched, 1); // only fuse has complexity >= 5 AND is function
690        assert_eq!(result.nodes[0].name, "fuse");
691    }
692
693    #[test]
694    fn test_query_empty_batch() {
695        let batch = build_code_nodes_batch(&[]).unwrap();
696        let result = codegraph_query_objects(&batch, &QueryFilter::default()).unwrap();
697        assert_eq!(result.total_matched, 0);
698    }
699
700    // --- Phase 1: Update tests ---
701
702    #[test]
703    fn test_update_docstring() {
704        let batch = build_code_nodes_batch(&sample_nodes()).unwrap();
705        let updated = codegraph_update_object(
706            &batch,
707            "func:brain/signal_fusion.py::fuse",
708            &NodeUpdate {
709                docstring: Some("Updated docstring.".into()),
710                ..Default::default()
711            },
712        )
713        .unwrap();
714
715        // Query the updated batch
716        let result = codegraph_query_objects(
717            &updated,
718            &QueryFilter {
719                name_contains: Some("fuse".into()),
720                ..Default::default()
721            },
722        )
723        .unwrap();
724        assert_eq!(result.nodes[0].docstring, Some("Updated docstring.".into()));
725    }
726
727    #[test]
728    fn test_update_multiple_fields() {
729        let batch = build_code_nodes_batch(&sample_nodes()).unwrap();
730        let updated = codegraph_update_object(
731            &batch,
732            "func:brain/store.py::promote",
733            &NodeUpdate {
734                loc: Some(50),
735                coverage_pct: Some(0.99),
736                ..Default::default()
737            },
738        )
739        .unwrap();
740
741        let result = codegraph_query_objects(
742            &updated,
743            &QueryFilter {
744                name_contains: Some("promote".into()),
745                ..Default::default()
746            },
747        )
748        .unwrap();
749        assert_eq!(result.nodes[0].loc, Some(50));
750        assert_eq!(result.nodes[0].coverage_pct, Some(0.99));
751    }
752
753    #[test]
754    fn test_update_nonexistent_node() {
755        let batch = build_code_nodes_batch(&sample_nodes()).unwrap();
756        let result = codegraph_update_object(
757            &batch,
758            "func:nonexistent::foo",
759            &NodeUpdate {
760                docstring: Some("nope".into()),
761                ..Default::default()
762            },
763        );
764        assert!(result.is_err());
765    }
766
767    #[test]
768    fn test_update_preserves_other_nodes() {
769        let batch = build_code_nodes_batch(&sample_nodes()).unwrap();
770        let updated = codegraph_update_object(
771            &batch,
772            "func:brain/signal_fusion.py::fuse",
773            &NodeUpdate {
774                docstring: Some("Changed.".into()),
775                ..Default::default()
776            },
777        )
778        .unwrap();
779
780        assert_eq!(updated.num_rows(), 3);
781        // Other nodes unchanged
782        let result = codegraph_query_objects(
783            &updated,
784            &QueryFilter {
785                name_contains: Some("DualStore".into()),
786                ..Default::default()
787            },
788        )
789        .unwrap();
790        assert_eq!(
791            result.nodes[0].docstring,
792            Some("Fast/slow dual-store.".into())
793        );
794    }
795
796    // --- Phase 1: Edge tests ---
797
798    #[test]
799    fn test_add_edge() {
800        let batch = build_code_edges_batch(&sample_edges()).unwrap();
801        let updated = codegraph_add_edge(
802            &batch,
803            "func:brain/signal_fusion.py::fuse",
804            "func:brain/store.py::promote",
805            "calls",
806            Some(1.0),
807        )
808        .unwrap();
809
810        assert_eq!(updated.num_rows(), 3); // 2 original + 1 new
811    }
812
813    #[test]
814    fn test_add_edge_invalid_predicate() {
815        let batch = build_code_edges_batch(&sample_edges()).unwrap();
816        let result = codegraph_add_edge(&batch, "a", "b", "nonexistent_pred", None);
817        assert!(result.is_err());
818    }
819
820    #[test]
821    fn test_remove_edge() {
822        let batch = build_code_edges_batch(&sample_edges()).unwrap();
823        let updated = codegraph_remove_edge(
824            &batch,
825            "func:brain/signal_fusion.py::fuse",
826            "class:brain/store.py::DualStore",
827            "uses",
828        )
829        .unwrap();
830
831        // Row count unchanged but weight is -1 (tombstone)
832        assert_eq!(updated.num_rows(), 2);
833        let weights = updated
834            .column(edge_col::WEIGHT)
835            .as_any()
836            .downcast_ref::<arrow::array::Float32Array>()
837            .unwrap();
838        assert_eq!(weights.value(0), -1.0);
839    }
840
841    #[test]
842    fn test_remove_edge_not_found() {
843        let batch = build_code_edges_batch(&sample_edges()).unwrap();
844        let result = codegraph_remove_edge(&batch, "a", "b", "calls");
845        assert!(result.is_err());
846    }
847
848    #[test]
849    fn test_query_update_query_round_trip() {
850        let batch = build_code_nodes_batch(&sample_nodes()).unwrap();
851
852        // Query original
853        let before = codegraph_query_objects(
854            &batch,
855            &QueryFilter {
856                name_contains: Some("fuse".into()),
857                ..Default::default()
858            },
859        )
860        .unwrap();
861        assert_eq!(
862            before.nodes[0].docstring,
863            Some("Fuse cognitive signals.".into())
864        );
865
866        // Update
867        let updated = codegraph_update_object(
868            &batch,
869            "func:brain/signal_fusion.py::fuse",
870            &NodeUpdate {
871                docstring: Some("Parallel weighted voting.".into()),
872                ..Default::default()
873            },
874        )
875        .unwrap();
876
877        // Query again
878        let after = codegraph_query_objects(
879            &updated,
880            &QueryFilter {
881                name_contains: Some("fuse".into()),
882                ..Default::default()
883            },
884        )
885        .unwrap();
886        assert_eq!(
887            after.nodes[0].docstring,
888            Some("Parallel weighted voting.".into())
889        );
890    }
891}
nusy_codegraph/mcp_tools.rs

nusy_codegraph/
mcp_tools.rs
