From e4e2aa8058370288c09f2e0dcfcdd5aa627c6259 Mon Sep 17 00:00:00 2001 From: rUv Date: Thu, 26 Feb 2026 16:23:58 +0000 Subject: [PATCH] fix(ruvector-gnn): replace panic with Result in MultiHeadAttention and RuvectorLayer constructors MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit MultiHeadAttention::new() and RuvectorLayer::new() used assert!() for input validation which caused fatal abort() when called from NAPI-RS/WASM bindings — unrecoverable by JavaScript callers. Both now return Result, and all WASM/NAPI wrappers propagate errors as catchable JS exceptions. Also fixes pre-existing mmap.rs test compilation error (grad_offset returns Option, not usize). Closes #216 Co-Authored-By: claude-flow --- .../src/graph.rs | 16 ++-- crates/ruvector-cli/src/mcp/gnn_cache.rs | 3 +- .../tests/gnn_performance_test.rs | 12 +-- crates/ruvector-crv/src/stage_iii.rs | 4 +- crates/ruvector-gnn-node/src/lib.rs | 22 ++---- crates/ruvector-gnn-wasm/src/lib.rs | 10 +-- crates/ruvector-gnn/src/layer.rs | 73 ++++++++++++++----- crates/ruvector-gnn/src/mmap.rs | 6 +- crates/ruvector-gnn/src/search.rs | 2 +- .../src/verified_training.rs | 16 ++-- .../tests/integration.rs | 4 +- 11 files changed, 100 insertions(+), 68 deletions(-) diff --git a/crates/ruvector-attention-unified-wasm/src/graph.rs b/crates/ruvector-attention-unified-wasm/src/graph.rs index ff091c566..50f1ab85b 100644 --- a/crates/ruvector-attention-unified-wasm/src/graph.rs +++ b/crates/ruvector-attention-unified-wasm/src/graph.rs @@ -43,14 +43,10 @@ impl WasmGNNLayer { heads: usize, dropout: f32, ) -> Result { - if dropout < 0.0 || dropout > 1.0 { - return Err(JsError::new("Dropout must be between 0.0 and 1.0")); - } + let inner = RuvectorLayer::new(input_dim, hidden_dim, heads, dropout) + .map_err(|e| JsError::new(&e.to_string()))?; - Ok(WasmGNNLayer { - inner: RuvectorLayer::new(input_dim, hidden_dim, heads, dropout), - hidden_dim, - }) + Ok(WasmGNNLayer { inner, hidden_dim }) } /// Forward pass through the GNN layer @@ -378,6 +374,12 @@ mod tests { assert!(layer.is_err()); } + #[wasm_bindgen_test] + fn test_gnn_layer_invalid_heads() { + let layer = WasmGNNLayer::new(4, 7, 3, 0.1); + assert!(layer.is_err()); + } + #[wasm_bindgen_test] fn test_tensor_compress_creation() { let compressor = WasmTensorCompress::new(); diff --git a/crates/ruvector-cli/src/mcp/gnn_cache.rs b/crates/ruvector-cli/src/mcp/gnn_cache.rs index c9b933dfd..5313f0411 100644 --- a/crates/ruvector-cli/src/mcp/gnn_cache.rs +++ b/crates/ruvector-cli/src/mcp/gnn_cache.rs @@ -189,7 +189,8 @@ impl GnnCache { } // Create new layer - let layer = RuvectorLayer::new(input_dim, hidden_dim, heads, dropout); + let layer = RuvectorLayer::new(input_dim, hidden_dim, heads, dropout) + .expect("GNN layer cache: invalid layer configuration"); // Cache it { diff --git a/crates/ruvector-cli/tests/gnn_performance_test.rs b/crates/ruvector-cli/tests/gnn_performance_test.rs index 4e8413bea..cb8ae24bb 100644 --- a/crates/ruvector-cli/tests/gnn_performance_test.rs +++ b/crates/ruvector-cli/tests/gnn_performance_test.rs @@ -25,7 +25,7 @@ mod gnn_cache_tests { #[test] fn test_layer_creation_latency() { let start = Instant::now(); - let _layer = RuvectorLayer::new(128, 256, 4, 0.1); + let _layer = RuvectorLayer::new(128, 256, 4, 0.1).unwrap(); let elapsed = start.elapsed(); // Layer creation: 100ms in release, ~2000ms in debug @@ -48,7 +48,7 @@ mod gnn_cache_tests { /// Test that forward pass has acceptable latency #[test] fn test_forward_pass_latency() { - let layer = RuvectorLayer::new(128, 256, 4, 0.1); + let layer = RuvectorLayer::new(128, 256, 4, 0.1).unwrap(); let node = vec![0.5f32; 128]; let neighbors = vec![vec![0.3f32; 128], vec![0.7f32; 128]]; let weights = vec![0.5f32, 0.5f32]; @@ -83,7 +83,7 @@ mod gnn_cache_tests { /// Test batch operations performance #[test] fn test_batch_operations_performance() { - let layer = RuvectorLayer::new(64, 128, 2, 0.1); + let layer = RuvectorLayer::new(64, 128, 2, 0.1).unwrap(); // Create batch of operations let batch_size = 100; @@ -139,7 +139,7 @@ mod gnn_cache_tests { for (input, hidden, heads) in sizes { // Measure creation let start = Instant::now(); - let layer = RuvectorLayer::new(input, hidden, heads, 0.1); + let layer = RuvectorLayer::new(input, hidden, heads, 0.1).unwrap(); let create_ms = start.elapsed().as_secs_f64() * 1000.0; // Measure forward @@ -216,7 +216,7 @@ mod gnn_cache_integration { // First: measure time including layer creation let start_cold = Instant::now(); - let layer = RuvectorLayer::new(128, 256, 4, 0.1); + let layer = RuvectorLayer::new(128, 256, 4, 0.1).unwrap(); let node = vec![0.5f32; 128]; let neighbors = vec![vec![0.3f32; 128], vec![0.7f32; 128]]; let weights = vec![0.5f32, 0.5f32]; @@ -262,7 +262,7 @@ mod gnn_cache_integration { // Create layer once let start = Instant::now(); - let layer = RuvectorLayer::new(64, 128, 2, 0.1); + let layer = RuvectorLayer::new(64, 128, 2, 0.1).unwrap(); let creation_time = start.elapsed(); let node = vec![0.5f32; 64]; diff --git a/crates/ruvector-crv/src/stage_iii.rs b/crates/ruvector-crv/src/stage_iii.rs index 4dd906089..d0dfcd0f8 100644 --- a/crates/ruvector-crv/src/stage_iii.rs +++ b/crates/ruvector-crv/src/stage_iii.rs @@ -31,7 +31,9 @@ impl StageIIIEncoder { pub fn new(config: &CrvConfig) -> Self { let dim = config.dimensions; // Single GNN layer: input_dim -> hidden_dim, 1 head - let gnn_layer = RuvectorLayer::new(dim, dim, 1, 0.0); + // heads=1 always divides any dim, and dropout=0.0 is always valid + let gnn_layer = RuvectorLayer::new(dim, dim, 1, 0.0) + .expect("dim is always divisible by 1 head"); Self { dim, gnn_layer } } diff --git a/crates/ruvector-gnn-node/src/lib.rs b/crates/ruvector-gnn-node/src/lib.rs index 97577e141..647c337fa 100644 --- a/crates/ruvector-gnn-node/src/lib.rs +++ b/crates/ruvector-gnn-node/src/lib.rs @@ -44,21 +44,15 @@ impl RuvectorLayer { /// ``` #[napi(constructor)] pub fn new(input_dim: u32, hidden_dim: u32, heads: u32, dropout: f64) -> Result { - if dropout < 0.0 || dropout > 1.0 { - return Err(Error::new( - Status::InvalidArg, - "Dropout must be between 0.0 and 1.0".to_string(), - )); - } + let inner = RustRuvectorLayer::new( + input_dim as usize, + hidden_dim as usize, + heads as usize, + dropout as f32, + ) + .map_err(|e| Error::new(Status::InvalidArg, e.to_string()))?; - Ok(Self { - inner: RustRuvectorLayer::new( - input_dim as usize, - hidden_dim as usize, - heads as usize, - dropout as f32, - ), - }) + Ok(Self { inner }) } /// Forward pass through the GNN layer diff --git a/crates/ruvector-gnn-wasm/src/lib.rs b/crates/ruvector-gnn-wasm/src/lib.rs index 2bf7c9dc3..6fad92a4a 100644 --- a/crates/ruvector-gnn-wasm/src/lib.rs +++ b/crates/ruvector-gnn-wasm/src/lib.rs @@ -81,14 +81,10 @@ impl JsRuvectorLayer { heads: usize, dropout: f32, ) -> Result { - if dropout < 0.0 || dropout > 1.0 { - return Err(JsValue::from_str("Dropout must be between 0.0 and 1.0")); - } + let inner = RuvectorLayer::new(input_dim, hidden_dim, heads, dropout) + .map_err(|e| JsValue::from_str(&e.to_string()))?; - Ok(JsRuvectorLayer { - inner: RuvectorLayer::new(input_dim, hidden_dim, heads, dropout), - hidden_dim, - }) + Ok(JsRuvectorLayer { inner, hidden_dim }) } /// Forward pass through the GNN layer diff --git a/crates/ruvector-gnn/src/layer.rs b/crates/ruvector-gnn/src/layer.rs index fa7db02a1..61afb4cbb 100644 --- a/crates/ruvector-gnn/src/layer.rs +++ b/crates/ruvector-gnn/src/layer.rs @@ -3,6 +3,7 @@ //! This module implements graph neural network layers that operate on HNSW graph structure, //! including attention mechanisms, normalization, and gated recurrent updates. +use crate::error::GnnError; use ndarray::{Array1, Array2, ArrayView1}; use rand::Rng; use rand_distr::{Distribution, Normal}; @@ -93,22 +94,27 @@ pub struct MultiHeadAttention { impl MultiHeadAttention { /// Create a new multi-head attention layer - pub fn new(embed_dim: usize, num_heads: usize) -> Self { - assert!( - embed_dim % num_heads == 0, - "Embedding dimension must be divisible by number of heads" - ); + /// + /// # Errors + /// Returns `GnnError::LayerConfig` if `embed_dim` is not divisible by `num_heads`. + pub fn new(embed_dim: usize, num_heads: usize) -> Result { + if embed_dim % num_heads != 0 { + return Err(GnnError::layer_config(format!( + "Embedding dimension ({}) must be divisible by number of heads ({})", + embed_dim, num_heads + ))); + } let head_dim = embed_dim / num_heads; - Self { + Ok(Self { num_heads, head_dim, q_linear: Linear::new(embed_dim, embed_dim), k_linear: Linear::new(embed_dim, embed_dim), v_linear: Linear::new(embed_dim, embed_dim), out_linear: Linear::new(embed_dim, embed_dim), - } + }) } /// Forward pass: compute multi-head attention @@ -334,20 +340,31 @@ impl RuvectorLayer { /// * `hidden_dim` - Dimension of hidden representations /// * `heads` - Number of attention heads /// * `dropout` - Dropout rate (0.0 to 1.0) - pub fn new(input_dim: usize, hidden_dim: usize, heads: usize, dropout: f32) -> Self { - assert!( - dropout >= 0.0 && dropout <= 1.0, - "Dropout must be between 0.0 and 1.0" - ); + /// + /// # Errors + /// Returns `GnnError::LayerConfig` if `dropout` is outside `[0.0, 1.0]` or + /// if `hidden_dim` is not divisible by `heads`. + pub fn new( + input_dim: usize, + hidden_dim: usize, + heads: usize, + dropout: f32, + ) -> Result { + if !(0.0..=1.0).contains(&dropout) { + return Err(GnnError::layer_config(format!( + "Dropout must be between 0.0 and 1.0, got {}", + dropout + ))); + } - Self { + Ok(Self { w_msg: Linear::new(input_dim, hidden_dim), w_agg: Linear::new(hidden_dim, hidden_dim), w_update: GRUCell::new(hidden_dim, hidden_dim), - attention: MultiHeadAttention::new(hidden_dim, heads), + attention: MultiHeadAttention::new(hidden_dim, heads)?, norm: LayerNorm::new(hidden_dim, 1e-5), dropout, - } + }) } /// Forward pass through the GNN layer @@ -464,7 +481,7 @@ mod tests { #[test] fn test_multihead_attention() { - let attention = MultiHeadAttention::new(8, 2); + let attention = MultiHeadAttention::new(8, 2).unwrap(); let query = vec![0.5; 8]; let keys = vec![vec![0.3; 8], vec![0.7; 8]]; let values = vec![vec![0.2; 8], vec![0.8; 8]]; @@ -473,6 +490,14 @@ mod tests { assert_eq!(output.len(), 8); } + #[test] + fn test_multihead_attention_invalid_dims() { + let result = MultiHeadAttention::new(10, 3); + assert!(result.is_err()); + let err = result.unwrap_err().to_string(); + assert!(err.contains("divisible")); + } + #[test] fn test_gru_cell() { let gru = GRUCell::new(4, 8); @@ -485,7 +510,7 @@ mod tests { #[test] fn test_ruvector_layer() { - let layer = RuvectorLayer::new(4, 8, 2, 0.1); + let layer = RuvectorLayer::new(4, 8, 2, 0.1).unwrap(); let node = vec![1.0, 2.0, 3.0, 4.0]; let neighbors = vec![vec![0.5, 1.0, 1.5, 2.0], vec![2.0, 3.0, 4.0, 5.0]]; @@ -497,7 +522,7 @@ mod tests { #[test] fn test_ruvector_layer_no_neighbors() { - let layer = RuvectorLayer::new(4, 8, 2, 0.1); + let layer = RuvectorLayer::new(4, 8, 2, 0.1).unwrap(); let node = vec![1.0, 2.0, 3.0, 4.0]; let neighbors: Vec> = vec![]; @@ -506,4 +531,16 @@ mod tests { let output = layer.forward(&node, &neighbors, &weights); assert_eq!(output.len(), 8); } + + #[test] + fn test_ruvector_layer_invalid_dropout() { + let result = RuvectorLayer::new(4, 8, 2, 1.5); + assert!(result.is_err()); + } + + #[test] + fn test_ruvector_layer_invalid_heads() { + let result = RuvectorLayer::new(4, 7, 3, 0.1); + assert!(result.is_err()); + } } diff --git a/crates/ruvector-gnn/src/mmap.rs b/crates/ruvector-gnn/src/mmap.rs index 38582bf02..559a58a59 100644 --- a/crates/ruvector-gnn/src/mmap.rs +++ b/crates/ruvector-gnn/src/mmap.rs @@ -883,9 +883,9 @@ mod tests { let accumulator = MmapGradientAccumulator::new(&path, 128, 100).unwrap(); - assert_eq!(accumulator.grad_offset(0), 0); - assert_eq!(accumulator.grad_offset(1), 128 * 4); // 128 floats * 4 bytes - assert_eq!(accumulator.grad_offset(5), 128 * 4 * 5); + assert_eq!(accumulator.grad_offset(0), Some(0)); + assert_eq!(accumulator.grad_offset(1), Some(128 * 4)); // 128 floats * 4 bytes + assert_eq!(accumulator.grad_offset(5), Some(128 * 4 * 5)); } #[test] diff --git a/crates/ruvector-gnn/src/search.rs b/crates/ruvector-gnn/src/search.rs index 00bbfde74..71239355c 100644 --- a/crates/ruvector-gnn/src/search.rs +++ b/crates/ruvector-gnn/src/search.rs @@ -237,7 +237,7 @@ mod tests { // Single GNN layer that maintains dimension let gnn_layers = vec![ - RuvectorLayer::new(2, 2, 1, 0.0), // input_dim, hidden_dim, heads, dropout + RuvectorLayer::new(2, 2, 1, 0.0).unwrap(), // input_dim, hidden_dim, heads, dropout ]; let result = hierarchical_forward(&query, &layer_embeddings, &gnn_layers); diff --git a/crates/ruvector-graph-transformer/src/verified_training.rs b/crates/ruvector-graph-transformer/src/verified_training.rs index 19620719c..8d9e8afff 100644 --- a/crates/ruvector-graph-transformer/src/verified_training.rs +++ b/crates/ruvector-graph-transformer/src/verified_training.rs @@ -959,7 +959,7 @@ mod tests { }, ]; let mut trainer = VerifiedTrainer::new(4, 8, config, invariants); - let layer = RuvectorLayer::new(4, 8, 2, 0.1); + let layer = RuvectorLayer::new(4, 8, 2, 0.1).unwrap(); let (features, neighbors, weights, targets) = test_data(); for step_num in 1..=10 { @@ -1015,7 +1015,7 @@ mod tests { }]; let mut trainer = VerifiedTrainer::new(4, 8, config, invariants); - let layer = RuvectorLayer::new(4, 8, 2, 0.1); + let layer = RuvectorLayer::new(4, 8, 2, 0.1).unwrap(); // Use features that produce large gradients let features = vec![vec![10.0, 10.0, 10.0, 10.0]]; @@ -1061,7 +1061,7 @@ mod tests { }]; let mut trainer = VerifiedTrainer::new(4, 8, config, invariants); - let layer = RuvectorLayer::new(4, 8, 2, 0.0); + let layer = RuvectorLayer::new(4, 8, 2, 0.0).unwrap(); let features = vec![vec![1.0, 2.0, 3.0, 4.0]]; let neighbors = vec![vec![vec![0.5, 1.0, 1.5, 2.0]]]; @@ -1111,7 +1111,7 @@ mod tests { }]; let mut trainer = VerifiedTrainer::new(4, 8, config, invariants); - let layer = RuvectorLayer::new(4, 8, 2, 0.0); + let layer = RuvectorLayer::new(4, 8, 2, 0.0).unwrap(); let (features, neighbors, weights_data, targets) = test_data(); // Run 3 steps @@ -1188,7 +1188,7 @@ mod tests { }]; let mut trainer = VerifiedTrainer::new(4, 8, config, invariants); - let layer = RuvectorLayer::new(4, 8, 2, 0.0); + let layer = RuvectorLayer::new(4, 8, 2, 0.0).unwrap(); let (features, neighbors, weights_data, targets) = test_data(); let result = trainer @@ -1310,7 +1310,7 @@ mod tests { rollback_strategy: RollbackStrategy::DeltaApply, }]; let mut trainer = VerifiedTrainer::new(4, 8, config, invariants); - let layer = RuvectorLayer::new(4, 8, 2, 0.0); + let layer = RuvectorLayer::new(4, 8, 2, 0.0).unwrap(); let (features, neighbors, weights, targets) = test_data(); let _ = trainer.train_step(&features, &neighbors, &weights, &targets, &layer); @@ -1342,7 +1342,7 @@ mod tests { }]; let mut trainer = VerifiedTrainer::new(4, 8, config, invariants); - let layer = RuvectorLayer::new(4, 8, 2, 0.0); + let layer = RuvectorLayer::new(4, 8, 2, 0.0).unwrap(); let (features, neighbors, weights, targets) = test_data(); // During warmup (steps 0..4), violations should NOT block commit @@ -1390,7 +1390,7 @@ mod tests { ]; let mut trainer = VerifiedTrainer::new(4, 8, config, invariants); - let layer = RuvectorLayer::new(4, 8, 2, 0.1); + let layer = RuvectorLayer::new(4, 8, 2, 0.1).unwrap(); let (features, neighbors, weights, targets) = test_data(); let result = trainer diff --git a/crates/ruvector-graph-transformer/tests/integration.rs b/crates/ruvector-graph-transformer/tests/integration.rs index 6effbb6de..f28c25cc6 100644 --- a/crates/ruvector-graph-transformer/tests/integration.rs +++ b/crates/ruvector-graph-transformer/tests/integration.rs @@ -342,7 +342,7 @@ mod verified_training_tests { ]; let mut trainer = VerifiedTrainer::new(4, 8, config, invariants); - let layer = RuvectorLayer::new(4, 8, 2, 0.0); + let layer = RuvectorLayer::new(4, 8, 2, 0.0).unwrap(); let features = vec![vec![1.0, 0.0, 0.0, 0.0]]; let neighbors = vec![vec![]]; let weights = vec![vec![]]; @@ -373,7 +373,7 @@ mod verified_training_tests { ]; let mut trainer = VerifiedTrainer::new(4, 8, config, invariants); - let layer = RuvectorLayer::new(4, 8, 2, 0.0); + let layer = RuvectorLayer::new(4, 8, 2, 0.0).unwrap(); for _ in 0..3 { let result = trainer.train_step(