diff --git a/docs/architecture.md b/docs/architecture.md index 20a7f3f5..51560a8c 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -108,6 +108,29 @@ Connections between nodes use hashed IDs defined in `ids.rs`: - `StringId`: ID for interned string values - `ReferenceId`: ID for constant or method reference occurrences (combines reference kind, URI, and offset) +## MCP Server + +The `rubydex-mcp` crate exposes rubydex's code intelligence as MCP tools over stdio JSON-RPC. The server indexes the codebase on startup, then serves tool requests against the immutable graph. + +### Pagination + +Tools that may return a high number of results accept `offset` and `limit` parameters and return a `total` count to support pagination. + +Pagination uses a two-pass approach: first collect all entries that pass filtering into a `Vec`, then apply `skip(offset).take(limit)`. This ensures `total` accurately reflects the number of results the caller can page through. + +### Result Ordering + +All collection-returning tools iterate over `IdentityHashMap` or `IdentityHashSet` structures. These use a deterministic hasher, so iteration order is fixed for a given map state. + +- **Within a server session**: Order is consistent between requests as long as the graph has not been re-indexed. Incremental re-indexing (e.g., after a file save) may change the graph between paginated requests, causing items to shift, appear, or disappear. Callers should not assume pagination stability across graph changes. +- **Across server restarts**: Order may change. Indexing is parallelized, so thread scheduling affects insertion order into the graph, which determines HashMap/HashSet bucket layout. + +### Key Files + +- `rubydex-mcp/src/server.rs`: Tool handler implementations and pagination logic +- `rubydex-mcp/src/tools.rs`: Parameter struct definitions with JSON Schema annotations +- `rubydex-mcp/tests/mcp.rs`: Integration tests (full MCP protocol over stdio) + ## FFI Layer The Rust crate exposes a C-compatible FFI API through `rubydex-sys`. The C extension in `ext/rubydex/` wraps this API for Ruby. diff --git a/rust/rubydex-mcp/src/main.rs b/rust/rubydex-mcp/src/main.rs index 258e2b7d..1113947f 100644 --- a/rust/rubydex-mcp/src/main.rs +++ b/rust/rubydex-mcp/src/main.rs @@ -1,5 +1,6 @@ use clap::Parser; +mod queries; mod server; mod tools; diff --git a/rust/rubydex-mcp/src/queries.rs b/rust/rubydex-mcp/src/queries.rs new file mode 100644 index 00000000..d47ea2f8 --- /dev/null +++ b/rust/rubydex-mcp/src/queries.rs @@ -0,0 +1,923 @@ +use std::collections::HashMap; +use std::path::{Path, PathBuf}; + +use rubydex::model::declaration::{Ancestor, Ancestors, Declaration, Namespace}; +use rubydex::model::graph::Graph; +use rubydex::model::ids::{DeclarationId, UriId}; +use serde::Serialize; +use url::Url; + +// --------------------------------------------------------------------------- +// Error types +// --------------------------------------------------------------------------- + +#[derive(Debug, PartialEq)] +pub enum QueryError { + NotFound { + name: String, + }, + InvalidKind { + name: String, + actual_kind: String, + tool_name: &'static str, + }, + InvalidPath { + path: String, + }, +} + +impl QueryError { + pub fn to_json_string(&self) -> String { + let (error, message, suggestion) = match self { + QueryError::NotFound { name } => ( + "not_found", + format!("Declaration '{name}' not found"), + "Try search_declarations with a partial name to find the correct FQN".to_string(), + ), + QueryError::InvalidKind { + name, + actual_kind, + tool_name, + } => ( + "invalid_kind", + format!("'{name}' is not a class or module (it is a {actual_kind})"), + format!("{tool_name} only works on classes and modules, not methods or constants"), + ), + QueryError::InvalidPath { path } => ( + "invalid_path", + format!("Cannot convert '{path}' to a file URI"), + "Use a relative path like 'app/models/user.rb' or an absolute path".to_string(), + ), + }; + serde_json::to_string(&serde_json::json!({ + "error": error, + "message": message, + "suggestion": suggestion, + })) + .unwrap_or_else(|_| "{}".to_string()) + } +} + +// --------------------------------------------------------------------------- +// Result types +// --------------------------------------------------------------------------- + +#[derive(Debug, Serialize, PartialEq)] +pub struct Location { + pub path: String, + pub line: u32, +} + +#[derive(Debug, Serialize, PartialEq)] +pub struct LocationWithColumn { + pub path: String, + pub line: u32, + pub column: u32, +} + +#[derive(Debug, Serialize, PartialEq)] +pub struct AncestorEntry { + pub name: String, + pub kind: &'static str, +} + +#[derive(Debug, Serialize, PartialEq)] +pub struct MemberEntry { + pub name: String, + pub kind: &'static str, + #[serde(skip_serializing_if = "Option::is_none")] + pub location: Option, +} + +#[derive(Debug, Serialize, PartialEq)] +pub struct DefinitionEntry { + pub path: String, + pub line: u32, + pub comments: Vec, +} + +#[derive(Debug, Serialize, PartialEq)] +pub struct SearchEntry { + pub name: String, + pub kind: &'static str, + pub locations: Vec, +} + +#[derive(Debug, Serialize, PartialEq)] +pub struct SearchDeclarationsResult { + pub results: Vec, + pub total: usize, +} + +#[derive(Debug, Serialize, PartialEq)] +pub struct DeclarationDetail { + pub name: String, + pub kind: &'static str, + pub definitions: Vec, + pub ancestors: Vec, + pub members: Vec, +} + +#[derive(Debug, Serialize, PartialEq)] +pub struct DescendantEntry { + pub name: String, + pub kind: &'static str, +} + +#[derive(Debug, Serialize, PartialEq)] +pub struct DescendantsResult { + pub name: String, + pub descendants: Vec, + pub total: usize, +} + +#[derive(Debug, Serialize, PartialEq)] +pub struct ReferencesResult { + pub name: String, + pub references: Vec, + pub total: usize, +} + +#[derive(Debug, Serialize, PartialEq)] +pub struct FileDeclarationEntry { + pub name: String, + pub kind: &'static str, + pub line: u32, +} + +#[derive(Debug, Serialize, PartialEq)] +pub struct FileDeclarationsResult { + pub file: String, + pub declarations: Vec, +} + +#[derive(Debug, Serialize, PartialEq)] +pub struct CodebaseStats { + pub files: usize, + pub declarations: usize, + pub definitions: usize, + pub constant_references: usize, + pub method_references: usize, + pub breakdown_by_kind: HashMap<&'static str, usize>, +} + +// --------------------------------------------------------------------------- +// Helpers +// --------------------------------------------------------------------------- + +fn uri_to_path(uri: &str) -> Option { + Url::parse(uri).ok()?.to_file_path().ok() +} + +pub fn format_path(uri: &str, root: &Path) -> String { + let Some(path) = uri_to_path(uri) else { + return uri.to_string(); + }; + path.strip_prefix(root) + .map_or_else(|_| path.display().to_string(), |rel| rel.display().to_string()) +} + +fn lookup_declaration<'g>( + graph: &'g Graph, + name: &str, +) -> Result<(DeclarationId, &'g Declaration), QueryError> { + let declaration_id = DeclarationId::from(name); + match graph.declarations().get(&declaration_id) { + Some(decl) => Ok((declaration_id, decl)), + None => Err(QueryError::NotFound { + name: name.to_string(), + }), + } +} + +fn require_namespace<'a>( + decl: &'a Declaration, + name: &str, + tool_name: &'static str, +) -> Result<&'a Namespace, QueryError> { + decl.as_namespace().ok_or_else(|| QueryError::InvalidKind { + name: name.to_string(), + actual_kind: decl.kind().to_string(), + tool_name, + }) +} + +fn collect_ancestors(graph: &Graph, ancestors: &Ancestors) -> Vec { + ancestors + .iter() + .filter_map(|ancestor| match ancestor { + Ancestor::Complete(id) => { + let decl = graph.declarations().get(id)?; + Some(AncestorEntry { + name: decl.name().to_string(), + kind: decl.kind(), + }) + } + Ancestor::Partial(name_id) => { + let name_ref = graph.names().get(name_id)?; + Some(AncestorEntry { + name: format!("{name_ref:?}"), + kind: "Unresolved", + }) + } + }) + .collect() +} + +// --------------------------------------------------------------------------- +// Query functions +// --------------------------------------------------------------------------- + +pub fn query_search_declarations( + graph: &Graph, + root: &Path, + query: &str, + kind: Option<&str>, + limit: usize, + offset: usize, +) -> SearchDeclarationsResult { + let ids = rubydex::query::declaration_search(graph, query); + + let matches_filter = |id: &&_| -> bool { + let Some(decl) = graph.declarations().get(id) else { + return false; + }; + if let Some(kind) = kind { + decl.kind().eq_ignore_ascii_case(kind) + } else { + true + } + }; + + let total = ids.iter().filter(matches_filter).count(); + let results = ids + .iter() + .filter(matches_filter) + .skip(offset) + .take(limit) + .filter_map(|id| { + let decl = graph.declarations().get(id)?; + let locations = decl + .definitions() + .iter() + .filter_map(|def_id| { + let def = graph.definitions().get(def_id)?; + let doc = graph.documents().get(def.uri_id())?; + let loc = def.offset().to_location(doc).to_presentation(); + Some(Location { + path: format_path(doc.uri(), root), + line: loc.start_line(), + }) + }) + .collect(); + Some(SearchEntry { + name: decl.name().to_string(), + kind: decl.kind(), + locations, + }) + }) + .collect(); + + SearchDeclarationsResult { results, total } +} + +pub fn query_get_declaration( + graph: &Graph, + root: &Path, + name: &str, +) -> Result { + let (_, decl) = lookup_declaration(graph, name)?; + + let definitions = decl + .definitions() + .iter() + .filter_map(|def_id| { + let def = graph.definitions().get(def_id)?; + let doc = graph.documents().get(def.uri_id())?; + let loc = def.offset().to_location(doc).to_presentation(); + let comments = def + .comments() + .iter() + .map(|c| { + c.string() + .as_str() + .strip_prefix("# ") + .unwrap_or(c.string().as_str()) + .to_string() + }) + .collect(); + Some(DefinitionEntry { + path: format_path(doc.uri(), root), + line: loc.start_line(), + comments, + }) + }) + .collect(); + + let namespace = decl.as_namespace(); + let ancestors = namespace + .map(|ns| collect_ancestors(graph, ns.ancestors())) + .unwrap_or_default(); + + let members = namespace + .map(|ns| { + ns.members() + .iter() + .filter_map(|(_, member_id)| { + let member_decl = graph.declarations().get(member_id)?; + let member_def = member_decl + .definitions() + .first() + .and_then(|def_id| graph.definitions().get(def_id)); + + let location = if let Some(def) = member_def + && let Some(doc) = graph.documents().get(def.uri_id()) + { + let loc = def.offset().to_location(doc).to_presentation(); + Some(Location { + path: format_path(doc.uri(), root), + line: loc.start_line(), + }) + } else { + None + }; + + Some(MemberEntry { + name: member_decl.name().to_string(), + kind: member_decl.kind(), + location, + }) + }) + .collect() + }) + .unwrap_or_default(); + + Ok(DeclarationDetail { + name: decl.name().to_string(), + kind: decl.kind(), + definitions, + ancestors, + members, + }) +} + +pub fn query_get_descendants( + graph: &Graph, + name: &str, + limit: usize, + offset: usize, +) -> Result { + let (_, decl) = lookup_declaration(graph, name)?; + let namespace = require_namespace(decl, name, "get_descendants")?; + + let exists = |id: &&_| graph.declarations().get(id).is_some(); + + let total = namespace.descendants().iter().filter(exists).count(); + let descendants = namespace + .descendants() + .iter() + .filter(exists) + .skip(offset) + .take(limit) + .filter_map(|id| { + let desc_decl = graph.declarations().get(id)?; + Some(DescendantEntry { + name: desc_decl.name().to_string(), + kind: desc_decl.kind(), + }) + }) + .collect(); + + Ok(DescendantsResult { + name: decl.name().to_string(), + descendants, + total, + }) +} + +pub fn query_find_constant_references( + graph: &Graph, + root: &Path, + name: &str, + limit: usize, + offset: usize, +) -> Result { + let (_, decl) = lookup_declaration(graph, name)?; + + let has_document = |ref_id: &&_| { + graph + .constant_references() + .get(ref_id) + .and_then(|r| graph.documents().get(&r.uri_id())) + .is_some() + }; + + let total = decl.references().iter().filter(has_document).count(); + let references = decl + .references() + .iter() + .filter(has_document) + .skip(offset) + .take(limit) + .filter_map(|ref_id| { + let const_ref = graph.constant_references().get(ref_id)?; + let doc = graph.documents().get(&const_ref.uri_id())?; + let loc = const_ref.offset().to_location(doc).to_presentation(); + Some(LocationWithColumn { + path: format_path(doc.uri(), root), + line: loc.start_line(), + column: loc.start_col(), + }) + }) + .collect(); + + Ok(ReferencesResult { + name: name.to_string(), + references, + total, + }) +} + +pub fn query_get_file_declarations( + graph: &Graph, + root: &Path, + canonical_path: &Path, +) -> Result { + let Ok(uri) = Url::from_file_path(canonical_path) else { + return Err(QueryError::InvalidPath { + path: canonical_path.display().to_string(), + }); + }; + + let uri_id = UriId::from(uri.as_str()); + let Some(doc) = graph.documents().get(&uri_id) else { + return Err(QueryError::NotFound { + name: canonical_path.display().to_string(), + }); + }; + + let declarations = doc + .definitions() + .iter() + .filter_map(|def_id| { + let def = graph.definitions().get(def_id)?; + let loc = def.offset().to_location(doc).to_presentation(); + let (name, kind) = graph + .definition_id_to_declaration_id(*def_id) + .and_then(|decl_id| graph.declarations().get(decl_id)) + .map(|decl| (decl.name().to_string(), decl.kind()))?; + Some(FileDeclarationEntry { name, kind, line: loc.start_line() }) + }) + .collect(); + + Ok(FileDeclarationsResult { + file: format_path(doc.uri(), root), + declarations, + }) +} + +pub fn query_codebase_stats(graph: &Graph) -> CodebaseStats { + let mut breakdown_by_kind: HashMap<&'static str, usize> = HashMap::new(); + for decl in graph.declarations().values() { + *breakdown_by_kind.entry(decl.kind()).or_default() += 1; + } + + CodebaseStats { + files: graph.documents().len(), + declarations: graph.declarations().len(), + definitions: graph.definitions().len(), + constant_references: graph.constant_references().len(), + method_references: graph.method_references().len(), + breakdown_by_kind, + } +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +#[cfg(test)] +mod tests { + use super::*; + use rubydex::test_utils::GraphTest; + use std::path::PathBuf; + + fn test_root() -> PathBuf { + if cfg!(windows) { + PathBuf::from("C:\\test") + } else { + PathBuf::from("/test") + } + } + + fn test_uri(filename: &str) -> String { + if cfg!(windows) { + format!("file:///C:/test/{filename}") + } else { + format!("file:///test/{filename}") + } + } + + fn graph_with_source(source: &str) -> Graph { + graph_with_sources(&[(&test_uri("test.rb"), source)]) + } + + fn graph_with_sources(sources: &[(&str, &str)]) -> Graph { + let mut gt = GraphTest::new(); + for (uri, source) in sources { + gt.index_uri(uri, source); + } + gt.resolve(); + gt.into_graph() + } + + // -- query_search_declarations -- + + #[test] + fn search_returns_matching_results() { + let graph = graph_with_source("class Dog; end"); + let root = test_root(); + let result = query_search_declarations(&graph, &root, "Dog", None, 50, 0); + + assert_eq!(result.total, 1); + assert_eq!(result.results[0].name, "Dog"); + assert_eq!(result.results[0].kind, "Class"); + assert!(result.results[0].locations[0].path.ends_with("test.rb")); + assert_eq!(result.results[0].locations[0].line, 1); + } + + #[test] + fn search_kind_filter() { + let graph = graph_with_source( + " + class Dog; end + module Walkable; end + ", + ); + let root = test_root(); + + let result = query_search_declarations(&graph, &root, "Dog", Some("Class"), 50, 0); + assert_eq!(result.results.len(), 1); + assert_eq!(result.results[0].name, "Dog"); + + let result = query_search_declarations(&graph, &root, "Dog", Some("Module"), 50, 0); + assert!(result.results.is_empty()); + + // Case-insensitive kind filter + let result = query_search_declarations(&graph, &root, "Dog", Some("class"), 50, 0); + assert_eq!(result.results.len(), 1); + + // Case-insensitive query + let result = query_search_declarations(&graph, &root, "dog", None, 50, 0); + assert!(!result.results.is_empty()); + } + + #[test] + fn search_no_match() { + let graph = graph_with_source("class Dog; end"); + let root = test_root(); + let result = query_search_declarations(&graph, &root, "Zzzzzzzzz", None, 50, 0); + assert!(result.results.is_empty()); + assert_eq!(result.total, 0); + } + + #[test] + fn search_pagination() { + let graph = graph_with_source( + " + class A; end + class B; end + class C; end + ", + ); + let root = test_root(); + + let result = query_search_declarations(&graph, &root, "", None, 2, 0); + assert_eq!(result.results.len(), 2); + let total = result.total; + + let result = query_search_declarations(&graph, &root, "", None, 2, 9999); + assert!(result.results.is_empty()); + assert_eq!(result.total, total); + + // Consecutive pages return different items + let page1 = query_search_declarations(&graph, &root, "", None, 1, 0); + let page2 = query_search_declarations(&graph, &root, "", None, 1, 1); + assert_ne!(page1.results[0].name, page2.results[0].name); + } + + // -- query_get_declaration -- + + #[test] + fn get_declaration_class_with_ancestors_and_members() { + let graph = graph_with_source( + " + class Animal; end + class Dog < Animal + def speak; end + def fetch; end + end + ", + ); + let root = test_root(); + let result = query_get_declaration(&graph, &root, "Dog").unwrap(); + + assert_eq!(result.name, "Dog"); + assert_eq!(result.kind, "Class"); + assert!(!result.definitions.is_empty()); + assert!(result.ancestors.iter().any(|a| a.name == "Animal")); + assert!(result.members.iter().any(|m| m.name == "Dog#speak()")); + assert!(result.members.iter().any(|m| m.name == "Dog#fetch()")); + + let speak = result.members.iter().find(|m| m.name == "Dog#speak()").unwrap(); + assert_eq!(speak.kind, "Method"); + assert!(speak.location.as_ref().unwrap().path.ends_with("test.rb")); + assert_eq!(speak.location.as_ref().unwrap().line, 3); + } + + #[test] + fn get_declaration_module() { + let graph = graph_with_source("module Greetable; end"); + let root = test_root(); + let result = query_get_declaration(&graph, &root, "Greetable").unwrap(); + assert_eq!(result.kind, "Module"); + } + + #[test] + fn get_declaration_doc_comments() { + let graph = graph_with_source( + " + # The Animal class represents all animals. + class Animal; end + ", + ); + let root = test_root(); + let result = query_get_declaration(&graph, &root, "Animal").unwrap(); + assert!(result.definitions[0].comments.iter().any(|c| c.contains("Animal"))); + } + + #[test] + fn get_declaration_mixin_ancestors() { + let graph = graph_with_source( + " + module Greetable; end + class Person + include Greetable + end + ", + ); + let root = test_root(); + let result = query_get_declaration(&graph, &root, "Person").unwrap(); + assert!(result.ancestors.iter().any(|a| a.name == "Greetable")); + } + + #[test] + fn get_declaration_constant() { + let graph = graph_with_source( + " + class Animal + KINGDOM = 'Animalia' + end + ", + ); + let root = test_root(); + let result = query_get_declaration(&graph, &root, "Animal::KINGDOM").unwrap(); + assert_eq!(result.kind, "Constant"); + assert!(result.ancestors.is_empty()); + assert!(result.members.is_empty()); + } + + #[test] + fn get_declaration_not_found() { + let graph = graph_with_source("class Dog; end"); + let root = test_root(); + let err = query_get_declaration(&graph, &root, "DoesNotExist").unwrap_err(); + assert!(matches!(err, QueryError::NotFound { .. })); + } + + // -- query_get_descendants -- + + #[test] + fn get_descendants_with_subclasses() { + let graph = graph_with_source( + " + class Animal; end + class Dog < Animal; end + class Cat < Animal; end + ", + ); + + let result = query_get_descendants(&graph, "Animal", 50, 0).unwrap(); + assert_eq!(result.name, "Animal"); + assert!(result.descendants.iter().any(|d| d.name == "Animal")); + assert!(result.descendants.iter().any(|d| d.name == "Dog")); + assert!(result.descendants.iter().any(|d| d.name == "Cat")); + assert_eq!(result.total, 3); + + let result = query_get_descendants(&graph, "Cat", 50, 0).unwrap(); + assert_eq!(result.total, 1); + } + + #[test] + fn get_descendants_module() { + let graph = graph_with_source( + " + module Greetable; end + class Person + include Greetable + end + ", + ); + let result = query_get_descendants(&graph, "Greetable", 50, 0).unwrap(); + assert!(result.descendants.iter().any(|d| d.name == "Person")); + } + + #[test] + fn get_descendants_inheritance_chain() { + let graph = graph_with_source( + " + class Foo; end + class Bar < Foo; end + class Baz < Bar; end + ", + ); + let result = query_get_descendants(&graph, "Foo", 50, 0).unwrap(); + assert!(result.descendants.iter().any(|d| d.name == "Bar")); + assert!(result.descendants.iter().any(|d| d.name == "Baz")); + } + + #[test] + fn get_descendants_pagination() { + let graph = graph_with_source( + " + class Animal; end + class Dog < Animal; end + class Cat < Animal; end + ", + ); + + let page1 = query_get_descendants(&graph, "Animal", 1, 0).unwrap(); + assert_eq!(page1.descendants.len(), 1); + assert_eq!(page1.total, 3); + + let page2 = query_get_descendants(&graph, "Animal", 1, 1).unwrap(); + assert_ne!(page1.descendants[0].name, page2.descendants[0].name); + } + + #[test] + fn get_descendants_not_found() { + let graph = graph_with_source("class Dog; end"); + let err = query_get_descendants(&graph, "DoesNotExist", 50, 0).unwrap_err(); + assert!(matches!(err, QueryError::NotFound { .. })); + } + + #[test] + fn get_descendants_invalid_kind() { + let graph = graph_with_source( + " + class Animal + KINGDOM = 'Animalia' + end + ", + ); + let err = query_get_descendants(&graph, "Animal::KINGDOM", 50, 0).unwrap_err(); + assert!(matches!(err, QueryError::InvalidKind { .. })); + } + + // -- query_find_constant_references -- + + #[test] + fn find_references_success() { + let graph = graph_with_source( + " + class Animal; end + class Dog < Animal; end + class Kennel + def build + Animal.new + end + end + ", + ); + let root = test_root(); + let result = query_find_constant_references(&graph, &root, "Animal", 50, 0).unwrap(); + + assert_eq!(result.name, "Animal"); + assert_eq!(result.references.len(), 2); + assert_eq!(result.total, 2); + assert!(result.references[0].path.ends_with("test.rb")); + assert_eq!(result.references[0].line, 2); + assert_eq!(result.references[0].column, 13); + } + + #[test] + fn find_references_cross_file() { + let models = test_uri("models.rb"); + let services = test_uri("services.rb"); + let graph = graph_with_sources(&[ + (&models, "class Dog; end"), + ( + &services, + " + class Kennel + def adopt + Dog.new + end + end + ", + ), + ]); + let root = test_root(); + let result = query_find_constant_references(&graph, &root, "Dog", 50, 0).unwrap(); + assert!(result.references.iter().any(|r| r.path.contains("services"))); + } + + #[test] + fn find_references_pagination() { + let graph = graph_with_source( + " + class Animal; end + class Dog < Animal; end + class Cat < Animal; end + class Kennel + def build + Animal.new + end + end + ", + ); + let root = test_root(); + + let full = query_find_constant_references(&graph, &root, "Animal", 50, 0).unwrap(); + let page = query_find_constant_references(&graph, &root, "Animal", 1, 0).unwrap(); + assert_eq!(page.references.len(), 1); + assert_eq!(page.total, full.total); + } + + #[test] + fn find_references_not_found() { + let graph = graph_with_source("class Dog; end"); + let root = test_root(); + let err = query_find_constant_references(&graph, &root, "DoesNotExist", 50, 0).unwrap_err(); + assert!(matches!(err, QueryError::NotFound { .. })); + } + + // -- query_get_file_declarations -- + + #[test] + fn file_declarations_success() { + let graph = graph_with_source( + " + class Animal; end + class Dog < Animal; end + module Greetable; end + ", + ); + let root = test_root(); + let result = query_get_file_declarations(&graph, &root, &root.join("test.rb")).unwrap(); + + assert!(result.declarations.iter().any(|d| d.name == "Animal")); + assert!(result.declarations.iter().any(|d| d.name == "Dog")); + assert!(result.declarations.iter().any(|d| d.name == "Greetable")); + assert_eq!(result.declarations[0].name, "Animal"); + assert_eq!(result.declarations[0].kind, "Class"); + assert_eq!(result.declarations[0].line, 1); + } + + #[test] + fn file_declarations_multiple_files() { + let models = test_uri("models.rb"); + let services = test_uri("services.rb"); + let graph = graph_with_sources(&[ + (&models, "class Animal; end"), + (&services, "class Kennel; end"), + ]); + let root = test_root(); + let result = query_get_file_declarations(&graph, &root, &root.join("services.rb")).unwrap(); + assert!(result.declarations.iter().any(|d| d.name == "Kennel")); + } + + #[test] + fn file_declarations_not_found() { + let graph = graph_with_source("class Dog; end"); + let root = test_root(); + let err = query_get_file_declarations(&graph, &root, &root.join("nonexistent.rb")).unwrap_err(); + assert!(matches!(err, QueryError::NotFound { .. })); + } + + // -- query_codebase_stats -- + + #[test] + fn codebase_stats_returns_counts() { + let a = test_uri("a.rb"); + let b = test_uri("b.rb"); + let graph = graph_with_sources(&[(&a, "class Animal; end"), (&b, "module Greetable; end")]); + let result = query_codebase_stats(&graph); + + assert_eq!(result.files, 2); + assert_eq!(result.declarations, 5); + assert_eq!(result.definitions, 2); + assert_eq!(result.breakdown_by_kind["Class"], 4); + assert_eq!(result.breakdown_by_kind["Module"], 1); + } +} diff --git a/rust/rubydex-mcp/src/server.rs b/rust/rubydex-mcp/src/server.rs index 5a0194c8..66af80a7 100644 --- a/rust/rubydex-mcp/src/server.rs +++ b/rust/rubydex-mcp/src/server.rs @@ -1,4 +1,3 @@ -use std::collections::HashMap; use std::path::{Path, PathBuf}; use std::sync::{Arc, RwLock}; @@ -13,12 +12,7 @@ use rmcp::{ tool, tool_handler, tool_router, transport::io::stdio, }; -use rubydex::model::ids::{DeclarationId, UriId}; -use rubydex::model::{ - declaration::{Ancestor, Ancestors}, - graph::Graph, -}; -use url::Url; +use rubydex::model::graph::Graph; struct ServerState { graph: Option, @@ -128,132 +122,26 @@ macro_rules! ensure_graph_ready { }}; } -/// Looks up a declaration by name, returning an error JSON string from the caller if not found. -macro_rules! lookup_declaration { - ($graph:expr, $name:expr) => {{ - let declaration_id = DeclarationId::from($name); - match $graph.declarations().get(&declaration_id) { - Some(decl) => (declaration_id, decl), - None => { - return error_json( - "not_found", - &format!("Declaration '{}' not found", $name), - "Try search_declarations with a partial name to find the correct FQN", - ); - } - } - }}; -} - -/// Narrows a declaration to a namespace, returning an error JSON string if it's not a class or module. -macro_rules! require_namespace { - ($decl:expr, $name:expr, $tool_name:literal) => { - match $decl.as_namespace() { - Some(ns) => ns, - None => { - return error_json( - "invalid_kind", - &format!("'{}' is not a class or module (it is a {})", $name, $decl.kind()), - concat!( - $tool_name, - " only works on classes and modules, not methods or constants" - ), - ); - } - } - }; -} - -/// Parses a file URI into a platform-native absolute path. -fn uri_to_path(uri: &str) -> Option { - Url::parse(uri).ok()?.to_file_path().ok() -} - -/// Converts a file URI to a path relative to `root` when possible. -/// Falls back to an absolute display path if it cannot be relativized. -fn format_path(uri: &str, root: &Path) -> String { - let Some(path) = uri_to_path(uri) else { - return uri.to_string(); - }; - - path.strip_prefix(root) - .map_or_else(|_| path.display().to_string(), |rel| rel.display().to_string()) -} - -/// Formats an ancestor chain into a JSON array of `{"name": ..., "kind": ...}` objects. -fn format_ancestors(graph: &Graph, ancestors: &Ancestors) -> Vec { - ancestors - .iter() - .filter_map(|ancestor| match ancestor { - Ancestor::Complete(id) => { - let ancestor_decl = graph.declarations().get(id)?; - Some(serde_json::json!({ - "name": ancestor_decl.name(), - "kind": ancestor_decl.kind(), - })) - } - Ancestor::Partial(name_id) => { - let name_ref = graph.names().get(name_id)?; - Some(serde_json::json!({ - "name": format!("{name_ref:?}"), - "kind": "Unresolved", - })) - } - }) - .collect() -} #[tool_router] impl RubydexServer { #[tool( - description = "Search for Ruby classes, modules, methods, or constants by name. Use this INSTEAD OF Grep when you know part of a Ruby identifier name and want to find its definition. Returns fully qualified names, kinds, and file locations. Use the `kind` filter (\"Class\", \"Module\", \"Method\", \"Constant\") to narrow results." + description = "Search for Ruby classes, modules, methods, or constants by name. Use this INSTEAD OF Grep when you know part of a Ruby identifier name and want to find its definition. Returns fully qualified names, kinds, and file locations. Use the `kind` filter (\"Class\", \"Module\", \"Method\", \"Constant\") to narrow results. Results are paginated: the response includes `total` (the full count of matches). If `total` exceeds the number of returned results, use `offset` to fetch subsequent pages." )] fn search_declarations(&self, Parameters(params): Parameters) -> String { let state = ensure_graph_ready!(self); let graph = state.graph.as_ref().unwrap(); - let ids = rubydex::query::declaration_search(graph, ¶ms.query); - - let limit = params.limit.unwrap_or(25).min(100); - let kind_filter = params.kind.as_deref(); - - let mut results: Vec = Vec::new(); - for id in ids { - if results.len() >= limit { - break; - } - - let Some(decl) = graph.declarations().get(&id) else { - continue; - }; - - if let Some(kind) = kind_filter - && !decl.kind().eq_ignore_ascii_case(kind) - { - continue; - } - - let locations: Vec = decl - .definitions() - .iter() - .filter_map(|def_id| { - let def = graph.definitions().get(def_id)?; - let doc = graph.documents().get(def.uri_id())?; - let loc = def.offset().to_location(doc).to_presentation(); - Some(serde_json::json!({ - "path": format_path(doc.uri(), &self.root_path), - "line": loc.start_line(), - })) - }) - .collect(); - - results.push(serde_json::json!({ - "name": decl.name(), - "kind": decl.kind(), - "locations": locations, - })); - } - - serde_json::to_string(&results).unwrap_or_else(|_| "[]".to_string()) + let limit = params.limit.filter(|&l| l > 0).unwrap_or(50).min(100); + let offset = params.offset.unwrap_or(0); + let result = crate::queries::query_search_declarations( + graph, + &self.root_path, + ¶ms.query, + params.kind.as_deref(), + limit, + offset, + ); + serde_json::to_string(&result).unwrap_or_else(|_| "{}".to_string()) } #[tool( @@ -262,155 +150,38 @@ impl RubydexServer { fn get_declaration(&self, Parameters(params): Parameters) -> String { let state = ensure_graph_ready!(self); let graph = state.graph.as_ref().unwrap(); - let (_, decl) = lookup_declaration!(graph, ¶ms.name); - - let definitions: Vec = decl - .definitions() - .iter() - .filter_map(|def_id| { - let def = graph.definitions().get(def_id)?; - let doc = graph.documents().get(def.uri_id())?; - let loc = def.offset().to_location(doc).to_presentation(); - let path = format_path(doc.uri(), &self.root_path); - let comments: Vec = def - .comments() - .iter() - .map(|c| { - c.string() - .as_str() - .strip_prefix("# ") - .unwrap_or(c.string().as_str()) - .to_string() - }) - .collect(); - - Some(serde_json::json!({ - "path": path, - "line": loc.start_line(), - "comments": comments, - })) - }) - .collect(); - - let namespace = decl.as_namespace(); - let ancestors = namespace - .map(|ns| format_ancestors(graph, ns.ancestors())) - .unwrap_or_default(); - - let members: Vec = namespace - .map(|ns| { - ns.members() - .iter() - .filter_map(|(_, member_id)| { - let member_decl = graph.declarations().get(member_id)?; - let member_def = member_decl - .definitions() - .first() - .and_then(|def_id| graph.definitions().get(def_id)); - - let mut member = serde_json::json!({ - "name": member_decl.name(), - "kind": member_decl.kind(), - }); - - if let Some(def) = member_def - && let Some(doc) = graph.documents().get(def.uri_id()) - { - let loc = def.offset().to_location(doc).to_presentation(); - member["location"] = serde_json::json!({ - "path": format_path(doc.uri(), &self.root_path), - "line": loc.start_line(), - }); - } - - Some(member) - }) - .collect() - }) - .unwrap_or_default(); - - let result = serde_json::json!({ - "name": decl.name(), - "kind": decl.kind(), - "definitions": definitions, - "ancestors": ancestors, - "members": members, - }); - - serde_json::to_string(&result).unwrap_or_else(|_| "{}".to_string()) + match crate::queries::query_get_declaration(graph, &self.root_path, ¶ms.name) { + Ok(result) => serde_json::to_string(&result).unwrap_or_else(|_| "{}".to_string()), + Err(e) => e.to_json_string(), + } } #[tool( - description = "Find all classes and modules that inherit from or include a given class/module. Returns known descendants, including transitive relationships, for impact analysis before modifying a base class or module." + description = "Returns all known descendants for the given namespace including itself and all transitive descendants. Can be used to understand how a module/class is used across the codebase. Results are paginated: the response includes `total`. If `total` exceeds the number of returned results, use `offset` to fetch subsequent pages." )] fn get_descendants(&self, Parameters(params): Parameters) -> String { let state = ensure_graph_ready!(self); let graph = state.graph.as_ref().unwrap(); - let (_, decl) = lookup_declaration!(graph, ¶ms.name); - let namespace = require_namespace!(decl, ¶ms.name, "get_descendants"); - - let descendants: Vec = namespace - .descendants() - .iter() - .filter_map(|desc_id| { - let desc_decl = graph.declarations().get(desc_id)?; - Some(serde_json::json!({ - "name": desc_decl.name(), - "kind": desc_decl.kind(), - })) - }) - .collect(); - - let result = serde_json::json!({ - "name": decl.name(), - "descendants": descendants, - }); - - serde_json::to_string(&result).unwrap_or_else(|_| "{}".to_string()) + let limit = params.limit.filter(|&l| l > 0).unwrap_or(50).min(500); + let offset = params.offset.unwrap_or(0); + match crate::queries::query_get_descendants(graph, ¶ms.name, limit, offset) { + Ok(result) => serde_json::to_string(&result).unwrap_or_else(|_| "{}".to_string()), + Err(e) => e.to_json_string(), + } } #[tool( - description = "Find all resolved references to a Ruby class, module, or constant across the codebase. Returns file paths, line numbers, and columns for each usage." + description = "Find all resolved references to a Ruby class, module, or constant across the codebase. Returns file paths, line numbers, and columns for each usage. Results are paginated: the response includes `total`. If `total` exceeds the number of returned results, use `offset` to fetch subsequent pages." )] fn find_constant_references(&self, Parameters(params): Parameters) -> String { let state = ensure_graph_ready!(self); let graph = state.graph.as_ref().unwrap(); - let (_, decl) = lookup_declaration!(graph, ¶ms.name); - - let limit = params.limit.unwrap_or(50).min(200); - let mut references: Vec = Vec::with_capacity(limit.min(decl.references().len())); - let mut truncated = false; - - for ref_id in decl.references() { - if references.len() >= limit { - truncated = true; - break; - } - - let Some(const_ref) = graph.constant_references().get(ref_id) else { - continue; - }; - - let Some(doc) = graph.documents().get(&const_ref.uri_id()) else { - continue; - }; - let loc = const_ref.offset().to_location(doc).to_presentation(); - references.push(serde_json::json!({ - "path": format_path(doc.uri(), &self.root_path), - "line": loc.start_line(), - "column": loc.start_col(), - })); - } - - let mut result = serde_json::json!({ - "name": params.name, - "references": references, - }); - if truncated { - result["truncated"] = serde_json::Value::Bool(true); + let limit = params.limit.filter(|&l| l > 0).unwrap_or(50).min(200); + let offset = params.offset.unwrap_or(0); + match crate::queries::query_find_constant_references(graph, &self.root_path, ¶ms.name, limit, offset) { + Ok(result) => serde_json::to_string(&result).unwrap_or_else(|_| "{}".to_string()), + Err(e) => e.to_json_string(), } - - serde_json::to_string(&result).unwrap_or_else(|_| "{}".to_string()) } #[tool( @@ -419,60 +190,16 @@ impl RubydexServer { fn get_file_declarations(&self, Parameters(params): Parameters) -> String { let state = ensure_graph_ready!(self); let graph = state.graph.as_ref().unwrap(); - let absolute_target = if Path::new(¶ms.file_path).is_absolute() { PathBuf::from(¶ms.file_path) } else { self.root_path.join(¶ms.file_path) }; let canonical_target = std::fs::canonicalize(&absolute_target).unwrap_or(absolute_target); - - let Ok(uri) = Url::from_file_path(&canonical_target) else { - return error_json( - "invalid_path", - &format!("Cannot convert '{}' to a file URI", params.file_path), - "Use a relative path like 'app/models/user.rb' or an absolute path", - ); - }; - - let uri_id = UriId::from(uri.as_str()); - let Some(doc) = graph.documents().get(&uri_id) else { - return error_json( - "not_found", - &format!("File '{}' not found in the index", params.file_path), - "Use a relative path like 'app/models/user.rb' or an absolute path matching the indexed project", - ); - }; - - let mut declarations: Vec = Vec::new(); - - for def_id in doc.definitions() { - let Some(def) = graph.definitions().get(def_id) else { - continue; - }; - - let loc = def.offset().to_location(doc).to_presentation(); - - let decl_name = graph - .definition_id_to_declaration_id(*def_id) - .and_then(|decl_id| graph.declarations().get(decl_id)) - .map(|decl| (decl.name().to_string(), decl.kind())); - - if let Some((name, kind)) = decl_name { - declarations.push(serde_json::json!({ - "name": name, - "kind": kind, - "line": loc.start_line(), - })); - } + match crate::queries::query_get_file_declarations(graph, &self.root_path, &canonical_target) { + Ok(result) => serde_json::to_string(&result).unwrap_or_else(|_| "{}".to_string()), + Err(e) => e.to_json_string(), } - - let result = serde_json::json!({ - "file": format_path(doc.uri(), &self.root_path), - "declarations": declarations, - }); - - serde_json::to_string(&result).unwrap_or_else(|_| "{}".to_string()) } #[tool( @@ -481,31 +208,16 @@ impl RubydexServer { fn codebase_stats(&self) -> String { let state = ensure_graph_ready!(self); let graph = state.graph.as_ref().unwrap(); - - let mut breakdown: HashMap<&str, usize> = HashMap::new(); - for decl in graph.declarations().values() { - *breakdown.entry(decl.kind()).or_default() += 1; - } - - let breakdown_json: serde_json::Value = breakdown - .iter() - .map(|(k, v)| (k.to_string(), serde_json::json!(v))) - .collect(); - - let result = serde_json::json!({ - "files": graph.documents().len(), - "declarations": graph.declarations().len(), - "definitions": graph.definitions().len(), - "constant_references": graph.constant_references().len(), - "method_references": graph.method_references().len(), - "breakdown_by_kind": breakdown_json, - }); - + let result = crate::queries::query_codebase_stats(graph); serde_json::to_string(&result).unwrap_or_else(|_| "{}".to_string()) } } -const SERVER_INSTRUCTIONS: &str = r#"Rubydex provides semantic Ruby code intelligence. Use these tools INSTEAD OF Grep when working with Ruby code structure. +const SERVER_INSTRUCTIONS: &str = r#"Rubydex provides semantic Ruby code intelligence. + +ONLY use these tools for Ruby files (.rb, .rbi, .rbs) — never for Rust, JavaScript, or other languages. + +Use these tools INSTEAD OF Grep when working with Ruby code structure. Decision guide: - Know a name? -> search_declarations (fuzzy search by name) @@ -519,6 +231,8 @@ Typical workflow: search_declarations -> get_declaration -> find_constant_refere Fully qualified name format: "Foo::Bar" for classes/modules/constants, "Foo::Bar#method_name" for instance methods. +Pagination: tools that may return a high number of results include `total` for pagination. When `total` exceeds the number of returned items, use `offset` to fetch the next page. + Use Grep instead for: literal string search, log messages, comments, non-Ruby files, or content search rather than structural queries."#; #[tool_handler] @@ -531,3 +245,318 @@ impl ServerHandler for RubydexServer { } } } + +#[cfg(test)] +mod tests { + use super::*; + use rubydex::test_utils::GraphTest; + use serde_json::Value; + + fn parse(json_str: &str) -> Value { + serde_json::from_str(json_str).unwrap() + } + + /// Assert a JSON array field contains an entry with the given "name". + macro_rules! assert_includes { + ($json:expr, $field:literal, $name:expr) => {{ + let json = &$json; + let entries = json[$field] + .as_array() + .expect(concat!("expected '", $field, "' to be an array")); + assert!( + entries.iter().any(|e| e["name"].as_str() == Some($name)), + "Expected '{}' in '{}', got: {:?}", + $name, + $field, + entries.iter().filter_map(|e| e["name"].as_str()).collect::>() + ); + }}; + } + + /// Extract a JSON field as an array, panicking if not an array. + macro_rules! array { + ($json:expr, $field:literal) => { + $json[$field] + .as_array() + .expect(concat!("expected '", $field, "' to be an array")) + }; + } + + /// Assert a JSON field equals the expected u64 value. + macro_rules! assert_json_int { + ($json:expr, $field:literal, $val:expr) => { + assert_eq!( + $json[$field] + .as_u64() + .expect(concat!("expected '", $field, "' to be a number")), + $val + ); + }; + } + + fn assert_error(json_str: &str, expected_type: &str) { + let res = parse(json_str); + assert_eq!( + res["error"].as_str(), + Some(expected_type), + "Expected error '{expected_type}', got: {res}" + ); + assert!(res["message"].as_str().is_some()); + assert!(res["suggestion"].as_str().is_some()); + } + + /// Returns a platform-appropriate test root path and its file URI prefix. + fn test_root() -> (&'static str, &'static str) { + if cfg!(windows) { + ("C:\\test", "file:///C:/test") + } else { + ("/test", "file:///test") + } + } + + fn test_uri(filename: &str) -> String { + let (_, uri_prefix) = test_root(); + format!("{uri_prefix}/{filename}") + } + + /// Build a server from a single Ruby source. + fn server_with_source(source: &str) -> RubydexServer { + server_with_sources(&[(&test_uri("test.rb"), source)]) + } + + /// Build a server from multiple `(uri, source)` pairs. + fn server_with_sources(sources: &[(&str, &str)]) -> RubydexServer { + let mut gt = GraphTest::new(); + for (uri, source) in sources { + gt.index_uri(uri, source); + } + gt.resolve(); + + let (root, _) = test_root(); + let server = RubydexServer::new(root.to_string()); + { + let mut state = server.state.write().unwrap(); + state.graph = Some(gt.into_graph()); + } + server + } + + macro_rules! search_declarations { + ($server:expr, $($field:ident: $val:expr),* $(,)?) => { + parse(&$server.search_declarations(Parameters(SearchDeclarationsParams { + $($field: $val,)* + }))) + }; + } + + macro_rules! get_descendants { + ($server:expr, $($field:ident: $val:expr),* $(,)?) => { + parse(&$server.get_descendants(Parameters(GetDescendantsParams { + $($field: $val,)* + }))) + }; + } + + macro_rules! find_constant_references { + ($server:expr, $($field:ident: $val:expr),* $(,)?) => { + parse(&$server.find_constant_references(Parameters(FindConstantReferencesParams { + $($field: $val,)* + }))) + }; + } + + fn get_declaration(server: &RubydexServer, name: &str) -> Value { + parse(&server.get_declaration(Parameters(GetDeclarationParams { name: name.to_string() }))) + } + + fn get_file_declarations(server: &RubydexServer, file_path: &str) -> Value { + parse(&server.get_file_declarations(Parameters(GetFileDeclarationsParams { + file_path: file_path.to_string(), + }))) + } + + // -- JSON shape tests (one per handler) -- + + #[test] + fn search_declarations_json_shape() { + let s = server_with_source("class Dog; end"); + let res = search_declarations!(s, query: "Dog".into(), kind: None, limit: None, offset: None); + + assert_includes!(res, "results", "Dog"); + assert_json_int!(res, "total", 1); + + let first = &array!(res, "results")[0]; + assert_eq!(first["name"], "Dog"); + assert_eq!(first["kind"], "Class"); + assert!(first["locations"][0]["path"].as_str().unwrap().ends_with("test.rb")); + assert_json_int!(first["locations"][0], "line", 1); + } + + #[test] + fn get_declaration_json_shape() { + let s = server_with_source( + " + class Animal; end + class Dog < Animal + def speak; end + end + ", + ); + let res = get_declaration(&s, "Dog"); + + assert_eq!(res["name"], "Dog"); + assert_eq!(res["kind"], "Class"); + assert!(!array!(res, "definitions").is_empty()); + assert_includes!(res, "ancestors", "Animal"); + assert_includes!(res, "members", "Dog#speak()"); + + let member = array!(res, "members") + .iter() + .find(|m| m["name"].as_str() == Some("Dog#speak()")) + .unwrap(); + assert_eq!(member["kind"], "Method"); + assert!(member["location"]["path"].as_str().unwrap().ends_with("test.rb")); + } + + #[test] + fn get_descendants_json_shape() { + let s = server_with_source( + " + class Animal; end + class Dog < Animal; end + ", + ); + let res = get_descendants!(s, name: "Animal".into(), limit: None, offset: None); + assert_eq!(res["name"], "Animal"); + assert_includes!(res, "descendants", "Dog"); + assert!(res["total"].as_u64().unwrap() > 0); + } + + #[test] + fn find_constant_references_json_shape() { + let s = server_with_source( + " + class Animal; end + class Dog < Animal; end + ", + ); + let res = find_constant_references!(s, name: "Animal".into(), limit: None, offset: None); + + assert_eq!(res["name"], "Animal"); + assert!(res["total"].as_u64().unwrap() > 0); + let first_ref = &array!(res, "references")[0]; + assert!(first_ref["path"].as_str().is_some()); + assert!(first_ref["line"].as_u64().is_some()); + assert!(first_ref["column"].as_u64().is_some()); + } + + #[test] + fn get_file_declarations_json_shape() { + let s = server_with_source("class Animal; end"); + let res = get_file_declarations(&s, "test.rb"); + + assert_includes!(res, "declarations", "Animal"); + assert_eq!(array!(res, "declarations")[0]["kind"], "Class"); + assert!(array!(res, "declarations")[0]["line"].as_u64().is_some()); + } + + #[test] + fn codebase_stats_json_shape() { + let a = test_uri("a.rb"); + let b = test_uri("b.rb"); + let s = server_with_sources(&[(&a, "class Animal; end"), (&b, "module Greetable; end")]); + let res = parse(&s.codebase_stats()); + + assert_eq!(res["files"], 2); + assert!(res["declarations"].as_u64().is_some()); + assert!(res["definitions"].as_u64().is_some()); + assert!(res["breakdown_by_kind"]["Class"].as_u64().is_some()); + } + + // -- error tests -- + + #[test] + fn get_declaration_not_found() { + let s = server_with_source("class Dog; end"); + assert_error( + &s.get_declaration(Parameters(GetDeclarationParams { + name: "DoesNotExist".into(), + })), + "not_found", + ); + } + + #[test] + fn get_descendants_not_found() { + let s = server_with_source("class Dog; end"); + assert_error( + &s.get_descendants(Parameters(GetDescendantsParams { + name: "DoesNotExist".into(), + limit: None, + offset: None, + })), + "not_found", + ); + } + + #[test] + fn get_descendants_invalid_kind() { + let s = server_with_source( + " + class Animal + KINGDOM = 'Animalia' + end + ", + ); + assert_error( + &s.get_descendants(Parameters(GetDescendantsParams { + name: "Animal::KINGDOM".into(), + limit: None, + offset: None, + })), + "invalid_kind", + ); + } + + #[test] + fn find_constant_references_not_found() { + let s = server_with_source("class Dog; end"); + assert_error( + &s.find_constant_references(Parameters(FindConstantReferencesParams { + name: "DoesNotExist".into(), + limit: None, + offset: None, + })), + "not_found", + ); + } + + #[test] + fn get_file_declarations_not_found() { + let s = server_with_source("class Dog; end"); + assert_error( + &s.get_file_declarations(Parameters(GetFileDeclarationsParams { + file_path: "nonexistent.rb".into(), + })), + "not_found", + ); + } + + // -- server state errors -- + + #[test] + fn returns_indexing_error_when_graph_not_ready() { + let server = RubydexServer::new("/test".to_string()); + assert_error(&server.codebase_stats(), "indexing"); + } + + #[test] + fn returns_indexing_failed_error() { + let server = RubydexServer::new("/test".to_string()); + { + let mut state = server.state.write().unwrap(); + state.error = Some("something went wrong".into()); + } + assert_error(&server.codebase_stats(), "indexing_failed"); + } +} diff --git a/rust/rubydex-mcp/src/tools.rs b/rust/rubydex-mcp/src/tools.rs index a8d9fc32..97136f6d 100644 --- a/rust/rubydex-mcp/src/tools.rs +++ b/rust/rubydex-mcp/src/tools.rs @@ -6,8 +6,10 @@ pub struct SearchDeclarationsParams { pub query: String, #[schemars(description = "Filter by declaration kind: Class, Module, Method, Constant, etc.")] pub kind: Option, - #[schemars(description = "Maximum number of results to return (default 25, max 100)")] + #[schemars(description = "Maximum number of results to return (default 50, max 100)")] pub limit: Option, + #[schemars(description = "Number of results to skip for pagination (default 0)")] + pub offset: Option, } #[derive(Debug, serde::Deserialize, JsonSchema)] @@ -20,6 +22,10 @@ pub struct GetDeclarationParams { pub struct GetDescendantsParams { #[schemars(description = "Fully qualified name of the class or module")] pub name: String, + #[schemars(description = "Maximum number of descendants to return (default 50, max 500)")] + pub limit: Option, + #[schemars(description = "Number of descendants to skip for pagination (default 0)")] + pub offset: Option, } #[derive(Debug, serde::Deserialize, JsonSchema)] @@ -28,6 +34,8 @@ pub struct FindConstantReferencesParams { pub name: String, #[schemars(description = "Maximum number of references to return (default 50, max 200)")] pub limit: Option, + #[schemars(description = "Number of references to skip for pagination (default 0)")] + pub offset: Option, } #[derive(Debug, serde::Deserialize, JsonSchema)] diff --git a/rust/rubydex-mcp/tests/mcp.rs b/rust/rubydex-mcp/tests/mcp.rs index 895ccdc4..2a2cbb48 100644 --- a/rust/rubydex-mcp/tests/mcp.rs +++ b/rust/rubydex-mcp/tests/mcp.rs @@ -214,16 +214,17 @@ fn mcp_server_e2e() { assert!(stats["declarations"].as_u64().unwrap() > 0); // Semantic query: search declarations. - let results: Vec = serde_json::from_value(call_next_tool( + let search_response = call_next_tool( &mut stdin, &mut reader, &mut request_id, "search_declarations", &json!({ "query": "Dog" }), - )) - .unwrap(); - let result_names = names_from_entries(&results); + ); + let results = search_response["results"].as_array().unwrap(); + let result_names = names_from_entries(results); assert_has_name(&result_names, "Dog", "search results"); + assert!(search_response["total"].as_u64().unwrap() > 0); // Semantic query: inspect declaration details. let decl = call_next_tool( @@ -253,6 +254,7 @@ fn mcp_server_e2e() { let descendant_entries = descendants["descendants"].as_array().unwrap(); let descendant_names = names_from_entries(descendant_entries); assert_has_name(&descendant_names, "Dog", "Animal descendants"); + assert!(descendants["total"].as_u64().unwrap() > 0); // Semantic query: resolved constant references. let references = call_next_tool( @@ -271,6 +273,7 @@ fn mcp_server_e2e() { refs.iter().all(|entry| entry["path"].as_str().is_some()), "Expected references to include file paths, got: {references}" ); + assert!(references["total"].as_u64().unwrap() > 0); // Semantic query: file declarations. let file_declarations = call_next_tool( diff --git a/rust/rubydex/src/test_utils/graph_test.rs b/rust/rubydex/src/test_utils/graph_test.rs index a3e731bb..000ae9da 100644 --- a/rust/rubydex/src/test_utils/graph_test.rs +++ b/rust/rubydex/src/test_utils/graph_test.rs @@ -22,6 +22,11 @@ impl GraphTest { &self.graph } + #[must_use] + pub fn into_graph(self) -> Graph { + self.graph + } + /// Indexes a Ruby source pub fn index_uri(&mut self, uri: &str, source: &str) { let source = normalize_indentation(source);