From 9fbdbd6bdb5840f8551089980042e2cd475f92d8 Mon Sep 17 00:00:00 2001 From: Stan Lo Date: Thu, 5 Mar 2026 17:31:45 +0000 Subject: [PATCH 1/3] Add pagination support and unit tests to rubydex MCP tools Add offset/total pagination to search_declarations, get_descendants, and find_constant_references. Uses a two-pass paginate! macro: filter pass counts total, map pass builds JSON only for the requested page. - search_declarations: response wrapped in {results, total} envelope - get_descendants: gains limit (default 100, max 500) and offset - find_constant_references: gains offset; total always present - limit=0 treated as default to guard against accidental empty pages - Tool descriptions and server instructions mention pagination - Add GraphTest::into_graph() for test graph injection - Add 26 unit tests calling tool handlers directly (no process spawn) - Windows-compatible test URIs via test_root()/test_uri() helpers --- rust/rubydex-mcp/src/server.rs | 749 +++++++++++++++++++--- rust/rubydex-mcp/src/tools.rs | 10 +- rust/rubydex-mcp/tests/mcp.rs | 11 +- rust/rubydex/src/test_utils/graph_test.rs | 5 + 4 files changed, 696 insertions(+), 79 deletions(-) diff --git a/rust/rubydex-mcp/src/server.rs b/rust/rubydex-mcp/src/server.rs index 5a0194c8..f75f4844 100644 --- a/rust/rubydex-mcp/src/server.rs +++ b/rust/rubydex-mcp/src/server.rs @@ -203,57 +203,80 @@ fn format_ancestors(graph: &Graph, ancestors: &Ancestors) -> Vec {{ + let filtered: Vec<_> = $items.filter($filter).collect(); + let total = filtered.len(); + let results: Vec = filtered + .into_iter() + .skip($offset) + .take($limit) + .filter_map($map) + .collect(); + (results, total) + }}; +} + #[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 limit = params.limit.filter(|&l| l > 0).unwrap_or(50).min(100); // default 50, max 100 + let offset = params.offset.unwrap_or(0); 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; - }; + let (results, total) = paginate!( + ids.iter(), + offset, + limit, + |id| { + let Some(decl) = graph.declarations().get(id) else { + return false; + }; + if let Some(kind) = kind_filter { + decl.kind().eq_ignore_ascii_case(kind) + } else { + true + } + }, + |id| { + let decl = graph.declarations().get(id)?; + 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(); - if let Some(kind) = kind_filter - && !decl.kind().eq_ignore_ascii_case(kind) - { - continue; - } + Some(serde_json::json!({ + "name": decl.name(), + "kind": decl.kind(), + "locations": locations, + })) + }, + ); - 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, - })); - } + let result = serde_json::json!({ + "results": results, + "total": total, + }); - serde_json::to_string(&results).unwrap_or_else(|_| "[]".to_string()) + serde_json::to_string(&result).unwrap_or_else(|_| "{}".to_string()) } #[tool( @@ -341,7 +364,7 @@ impl RubydexServer { } #[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); @@ -349,66 +372,71 @@ impl RubydexServer { 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)?; + let limit = params.limit.filter(|&l| l > 0).unwrap_or(50).min(500); // default 50, max 500 + let offset = params.offset.unwrap_or(0); + + let (descendants, total) = paginate!( + namespace.descendants().iter(), + offset, + limit, + |id| graph.declarations().get(id).is_some(), + |id| { + let desc_decl = graph.declarations().get(id)?; Some(serde_json::json!({ "name": desc_decl.name(), "kind": desc_decl.kind(), })) - }) - .collect(); + }, + ); let result = serde_json::json!({ "name": decl.name(), "descendants": descendants, + "total": total, }); serde_json::to_string(&result).unwrap_or_else(|_| "{}".to_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 limit = params.limit.filter(|&l| l > 0).unwrap_or(50).min(200); // default 50, max 200 + let offset = params.offset.unwrap_or(0); - 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 (references, total) = paginate!( + decl.references().iter(), + offset, + limit, + |ref_id| { + graph + .constant_references() + .get(ref_id) + .and_then(|r| graph.documents().get(&r.uri_id())) + .is_some() + }, + |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(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!({ + let result = serde_json::json!({ "name": params.name, "references": references, + "total": total, }); - if truncated { - result["truncated"] = serde_json::Value::Bool(true); - } serde_json::to_string(&result).unwrap_or_else(|_| "{}".to_string()) } @@ -505,7 +533,11 @@ impl RubydexServer { } } -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 +551,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 +565,570 @@ 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(), + }))) + } + + // -- search_declarations -- + + #[test] + fn search_declarations_returns_matching_results() { + 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 search_declarations_kind_filter() { + let s = server_with_source( + " + class Dog; end + module Walkable; end + ", + ); + + let res = search_declarations!(s, query: "Dog".into(), kind: Some("Class".into()), limit: None, offset: None); + assert_includes!(res, "results", "Dog"); + + let res = search_declarations!(s, query: "Dog".into(), kind: Some("Module".into()), limit: None, offset: None); + assert!(array!(res, "results").is_empty()); + + // Case-insensitive + let res = search_declarations!(s, query: "Dog".into(), kind: Some("class".into()), limit: None, offset: None); + assert_includes!(res, "results", "Dog"); + + let res = search_declarations!(s, query: "dog".into(), kind: None, limit: None, offset: None); + assert_includes!(res, "results", "Dog"); + } + + #[test] + fn search_declarations_no_match() { + let s = server_with_source("class Dog; end"); + let res = search_declarations!(s, query: "Zzzzzzzzz".into(), kind: None, limit: None, offset: None); + assert!(array!(res, "results").is_empty()); + assert_json_int!(res, "total", 0); + } + + #[test] + fn search_declarations_pagination() { + let s = server_with_source( + " + class A; end + class B; end + class C; end + ", + ); + + let res = search_declarations!(s, query: String::new(), kind: None, limit: Some(2), offset: Some(0)); + assert_eq!(array!(res, "results").len(), 2); + let total = res["total"].as_u64().unwrap(); + + let res = search_declarations!(s, query: String::new(), kind: None, limit: Some(2), offset: Some(9999)); + assert!(array!(res, "results").is_empty()); + assert_json_int!(res, "total", total); + + // Verify consecutive pages return different items + let page1 = search_declarations!(s, query: String::new(), kind: None, limit: Some(1), offset: Some(0)); + let page2 = search_declarations!(s, query: String::new(), kind: None, limit: Some(1), offset: Some(1)); + let name1 = array!(page1, "results")[0]["name"].as_str().unwrap(); + let name2 = array!(page2, "results")[0]["name"].as_str().unwrap(); + assert_ne!(name1, name2, "Page 1 and page 2 should return different items"); + } + + // -- get_declaration -- + + #[test] + fn get_declaration_class_with_ancestors_and_members() { + let s = server_with_source( + " + class Animal; end + class Dog < Animal + def speak; end + def fetch; 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()"); + assert_includes!(res, "members", "Dog#fetch()"); + + 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")); + assert_json_int!(member["location"], "line", 3); + } + + #[test] + fn get_declaration_module() { + let s = server_with_source("module Greetable; end"); + assert_eq!(get_declaration(&s, "Greetable")["kind"], "Module"); + } + + #[test] + fn get_declaration_doc_comments() { + let s = server_with_source( + " + # The Animal class represents all animals. + class Animal; end + ", + ); + let res = get_declaration(&s, "Animal"); + let comments = array!(res["definitions"][0], "comments"); + assert!( + comments.iter().any(|c| c.as_str().unwrap().contains("Animal")), + "Expected doc comment on Animal, got: {comments:?}" + ); + } + + #[test] + fn get_declaration_mixin_ancestors() { + let s = server_with_source( + " + module Greetable; end + class Person + include Greetable + end + ", + ); + assert_includes!(get_declaration(&s, "Person"), "ancestors", "Greetable"); + } + + #[test] + fn get_declaration_constant() { + let s = server_with_source( + " + class Animal + KINGDOM = 'Animalia' + end + ", + ); + let res = get_declaration(&s, "Animal::KINGDOM"); + assert_eq!(res["kind"], "Constant"); + assert!(array!(res, "ancestors").is_empty()); + assert!(array!(res, "members").is_empty()); + } + + #[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", + ); + } + + // -- get_descendants -- + + #[test] + fn get_descendants_with_subclasses() { + let s = server_with_source( + " + class Animal; end + class Dog < Animal; end + class Cat < Animal; end + ", + ); + + let res = get_descendants!(s, name: "Animal".into(), limit: None, offset: None); + assert_eq!(res["name"], "Animal"); + assert_includes!(res, "descendants", "Animal"); + assert_includes!(res, "descendants", "Dog"); + assert_includes!(res, "descendants", "Cat"); + assert_json_int!(res, "total", 3); + + // Cat: 1 descendant (itself only, no subclasses) + let res = get_descendants!(s, name: "Cat".into(), limit: None, offset: None); + assert_json_int!(res, "total", 1); + } + + #[test] + fn get_descendants_module() { + let s = server_with_source( + " + module Greetable; end + + class Person + include Greetable + end + ", + ); + let res = get_descendants!(s, name: "Greetable".into(), limit: None, offset: None); + assert_includes!(res, "descendants", "Person"); + } + + #[test] + fn get_descendants_inheritance_chain() { + let s = server_with_source( + " + class Foo; end + class Bar < Foo; end + class Baz < Bar; end + ", + ); + let res = get_descendants!(s, name: "Foo".into(), limit: None, offset: None); + assert_includes!(res, "descendants", "Bar"); + assert_includes!(res, "descendants", "Baz"); + } + + #[test] + fn get_descendants_pagination() { + let s = server_with_source( + " + class Animal; end + class Dog < Animal; end + class Cat < Animal; end + ", + ); + let page1 = get_descendants!(s, name: "Animal".into(), limit: Some(1), offset: Some(0)); + assert_eq!(array!(page1, "descendants").len(), 1); + assert_json_int!(page1, "total", 3); + + let page2 = get_descendants!(s, name: "Animal".into(), limit: Some(1), offset: Some(1)); + let name1 = array!(page1, "descendants")[0]["name"].as_str().unwrap(); + let name2 = array!(page2, "descendants")[0]["name"].as_str().unwrap(); + assert_ne!(name1, name2, "Page 1 and page 2 should return different descendants"); + } + + #[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", + ); + } + + // -- find_constant_references -- + + #[test] + fn find_constant_references_success() { + let s = server_with_source( + " + class Animal; end + class Dog < Animal; end + class Kennel + def build + Animal.new + end + end + ", + ); + let res = find_constant_references!(s, name: "Animal".into(), limit: None, offset: None); + + assert_eq!(res["name"], "Animal"); + assert_eq!(array!(res, "references").len(), 2); + assert_json_int!(res, "total", 2); + let first_ref = &array!(res, "references")[0]; + assert!(first_ref["path"].as_str().unwrap().ends_with("test.rb")); + assert_json_int!(first_ref, "line", 2); + assert_json_int!(first_ref, "column", 13); + } + + #[test] + fn find_constant_references_cross_file() { + let models = test_uri("models.rb"); + let services = test_uri("services.rb"); + let s = server_with_sources(&[ + (&models, "class Dog; end"), + ( + &services, + " + class Kennel + def adopt + Dog.new + end + end + ", + ), + ]); + let res = find_constant_references!(s, name: "Dog".into(), limit: None, offset: None); + let paths: Vec<&str> = array!(res, "references") + .iter() + .filter_map(|r| r["path"].as_str()) + .collect(); + assert!( + paths.iter().any(|p| p.contains("services")), + "Expected cross-file ref from services, got: {paths:?}" + ); + } + + #[test] + fn find_constant_references_pagination() { + let s = server_with_source( + " + class Animal; end + class Dog < Animal; end + class Cat < Animal; end + class Kennel + def build + Animal.new + end + end + ", + ); + let full = find_constant_references!(s, name: "Animal".into(), limit: None, offset: None); + let full_total = full["total"].as_u64().unwrap(); + + let page = find_constant_references!(s, name: "Animal".into(), limit: Some(1), offset: Some(0)); + assert_eq!(array!(page, "references").len(), 1); + assert_json_int!(page, "total", full_total); + } + + #[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", + ); + } + + // -- get_file_declarations -- + + #[test] + fn get_file_declarations_success() { + let s = server_with_source( + " + class Animal; end + class Dog < Animal; end + module Greetable; end + ", + ); + let res = get_file_declarations(&s, "test.rb"); + + assert_includes!(res, "declarations", "Animal"); + assert_includes!(res, "declarations", "Dog"); + assert_includes!(res, "declarations", "Greetable"); + assert_eq!(array!(res, "declarations")[0]["name"], "Animal"); + assert_eq!(array!(res, "declarations")[0]["kind"], "Class"); + assert_json_int!(array!(res, "declarations")[0], "line", 1); + } + + #[test] + fn get_file_declarations_multiple_files() { + let models = test_uri("models.rb"); + let services = test_uri("services.rb"); + let s = server_with_sources(&[(&models, "class Animal; end"), (&services, "class Kennel; end")]); + let res = get_file_declarations(&s, "services.rb"); + assert_includes!(res, "declarations", "Kennel"); + } + + #[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", + ); + } + + // -- codebase_stats -- + + #[test] + fn codebase_stats_returns_counts() { + 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_json_int!(res, "declarations", 5); + assert_json_int!(res, "definitions", 2); + + let breakdown = &res["breakdown_by_kind"]; + assert_json_int!(breakdown, "Class", 4); + assert_json_int!(breakdown, "Module", 1); + } + + // -- error states -- + + #[test] + fn returns_indexing_error_when_graph_not_ready() { + let server = RubydexServer::new("/test".to_string()); + // graph is None (still indexing) + 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); From ecd1dc03922b0095ab7a4e7c7c23058504d10d3f Mon Sep 17 00:00:00 2001 From: Stan Lo Date: Thu, 5 Mar 2026 17:31:45 +0000 Subject: [PATCH 2/3] Document MCP server pagination and result ordering in architecture.md Add MCP Server section covering the two-pass pagination approach, result ordering guarantees (stable within a session, may vary across restarts due to parallel indexing), and key file references. --- docs/architecture.md | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) 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. From 5a9b4cf3982535fd814d1f9c60162a89e3a9eef9 Mon Sep 17 00:00:00 2001 From: Stan Lo Date: Fri, 6 Mar 2026 22:30:20 +0000 Subject: [PATCH 3/3] Separate MCP query logic from tool handler serialization Extract core query logic from server.rs tool handlers into a new queries.rs module with typed Rust result structs. This enables testing business logic directly without JSON serialization. - Add queries.rs with QueryError enum, 12 typed result structs, 6 pure query functions, and 22 unit tests - Refactor server.rs handlers to thin adapters (~5 lines each): acquire graph, normalize params, call query, serialize - Split tests: query logic tests in queries.rs (typed), JSON shape and error tests in server.rs (integration) - Remove intermediate Vec allocations in pagination (use iterator count) - Use &'static str for kind fields and HashMap keys instead of String --- rust/rubydex-mcp/src/main.rs | 1 + rust/rubydex-mcp/src/queries.rs | 923 ++++++++++++++++++++++++++++++++ rust/rubydex-mcp/src/server.rs | 708 +++--------------------- 3 files changed, 992 insertions(+), 640 deletions(-) create mode 100644 rust/rubydex-mcp/src/queries.rs 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 f75f4844..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,96 +122,6 @@ 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() -} - -/// Filters, paginates, and maps items. Returns `(results, total)` where `total` is the -/// count of all items passing the filter, and `results` contains only the requested page. -macro_rules! paginate { - ($items:expr, $offset:expr, $limit:expr, $filter:expr, $map:expr $(,)?) => {{ - let filtered: Vec<_> = $items.filter($filter).collect(); - let total = filtered.len(); - let results: Vec = filtered - .into_iter() - .skip($offset) - .take($limit) - .filter_map($map) - .collect(); - (results, total) - }}; -} #[tool_router] impl RubydexServer { @@ -227,55 +131,16 @@ impl RubydexServer { 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.filter(|&l| l > 0).unwrap_or(50).min(100); // default 50, max 100 + let limit = params.limit.filter(|&l| l > 0).unwrap_or(50).min(100); let offset = params.offset.unwrap_or(0); - let kind_filter = params.kind.as_deref(); - - let (results, total) = paginate!( - ids.iter(), - offset, + let result = crate::queries::query_search_declarations( + graph, + &self.root_path, + ¶ms.query, + params.kind.as_deref(), limit, - |id| { - let Some(decl) = graph.declarations().get(id) else { - return false; - }; - if let Some(kind) = kind_filter { - decl.kind().eq_ignore_ascii_case(kind) - } else { - true - } - }, - |id| { - let decl = graph.declarations().get(id)?; - 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(); - - Some(serde_json::json!({ - "name": decl.name(), - "kind": decl.kind(), - "locations": locations, - })) - }, + offset, ); - - let result = serde_json::json!({ - "results": results, - "total": total, - }); - serde_json::to_string(&result).unwrap_or_else(|_| "{}".to_string()) } @@ -285,82 +150,10 @@ 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( @@ -369,33 +162,12 @@ impl RubydexServer { 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 limit = params.limit.filter(|&l| l > 0).unwrap_or(50).min(500); // default 50, max 500 + let limit = params.limit.filter(|&l| l > 0).unwrap_or(50).min(500); let offset = params.offset.unwrap_or(0); - - let (descendants, total) = paginate!( - namespace.descendants().iter(), - offset, - limit, - |id| graph.declarations().get(id).is_some(), - |id| { - let desc_decl = graph.declarations().get(id)?; - Some(serde_json::json!({ - "name": desc_decl.name(), - "kind": desc_decl.kind(), - })) - }, - ); - - let result = serde_json::json!({ - "name": decl.name(), - "descendants": descendants, - "total": total, - }); - - serde_json::to_string(&result).unwrap_or_else(|_| "{}".to_string()) + 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( @@ -404,41 +176,12 @@ impl RubydexServer { 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.filter(|&l| l > 0).unwrap_or(50).min(200); // default 50, max 200 + let limit = params.limit.filter(|&l| l > 0).unwrap_or(50).min(200); let offset = params.offset.unwrap_or(0); - - let (references, total) = paginate!( - decl.references().iter(), - offset, - limit, - |ref_id| { - graph - .constant_references() - .get(ref_id) - .and_then(|r| graph.documents().get(&r.uri_id())) - .is_some() - }, - |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(serde_json::json!({ - "path": format_path(doc.uri(), &self.root_path), - "line": loc.start_line(), - "column": loc.start_col(), - })) - }, - ); - - let result = serde_json::json!({ - "name": params.name, - "references": references, - "total": total, - }); - - serde_json::to_string(&result).unwrap_or_else(|_| "{}".to_string()) + 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(), + } } #[tool( @@ -447,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( @@ -509,26 +208,7 @@ 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()) } } @@ -695,10 +375,10 @@ mod tests { }))) } - // -- search_declarations -- + // -- JSON shape tests (one per handler) -- #[test] - fn search_declarations_returns_matching_results() { + 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); @@ -713,72 +393,12 @@ mod tests { } #[test] - fn search_declarations_kind_filter() { - let s = server_with_source( - " - class Dog; end - module Walkable; end - ", - ); - - let res = search_declarations!(s, query: "Dog".into(), kind: Some("Class".into()), limit: None, offset: None); - assert_includes!(res, "results", "Dog"); - - let res = search_declarations!(s, query: "Dog".into(), kind: Some("Module".into()), limit: None, offset: None); - assert!(array!(res, "results").is_empty()); - - // Case-insensitive - let res = search_declarations!(s, query: "Dog".into(), kind: Some("class".into()), limit: None, offset: None); - assert_includes!(res, "results", "Dog"); - - let res = search_declarations!(s, query: "dog".into(), kind: None, limit: None, offset: None); - assert_includes!(res, "results", "Dog"); - } - - #[test] - fn search_declarations_no_match() { - let s = server_with_source("class Dog; end"); - let res = search_declarations!(s, query: "Zzzzzzzzz".into(), kind: None, limit: None, offset: None); - assert!(array!(res, "results").is_empty()); - assert_json_int!(res, "total", 0); - } - - #[test] - fn search_declarations_pagination() { - let s = server_with_source( - " - class A; end - class B; end - class C; end - ", - ); - - let res = search_declarations!(s, query: String::new(), kind: None, limit: Some(2), offset: Some(0)); - assert_eq!(array!(res, "results").len(), 2); - let total = res["total"].as_u64().unwrap(); - - let res = search_declarations!(s, query: String::new(), kind: None, limit: Some(2), offset: Some(9999)); - assert!(array!(res, "results").is_empty()); - assert_json_int!(res, "total", total); - - // Verify consecutive pages return different items - let page1 = search_declarations!(s, query: String::new(), kind: None, limit: Some(1), offset: Some(0)); - let page2 = search_declarations!(s, query: String::new(), kind: None, limit: Some(1), offset: Some(1)); - let name1 = array!(page1, "results")[0]["name"].as_str().unwrap(); - let name2 = array!(page2, "results")[0]["name"].as_str().unwrap(); - assert_ne!(name1, name2, "Page 1 and page 2 should return different items"); - } - - // -- get_declaration -- - - #[test] - fn get_declaration_class_with_ancestors_and_members() { + fn get_declaration_json_shape() { let s = server_with_source( " class Animal; end class Dog < Animal def speak; end - def fetch; end end ", ); @@ -789,7 +409,6 @@ mod tests { assert!(!array!(res, "definitions").is_empty()); assert_includes!(res, "ancestors", "Animal"); assert_includes!(res, "members", "Dog#speak()"); - assert_includes!(res, "members", "Dog#fetch()"); let member = array!(res, "members") .iter() @@ -797,140 +416,74 @@ mod tests { .unwrap(); assert_eq!(member["kind"], "Method"); assert!(member["location"]["path"].as_str().unwrap().ends_with("test.rb")); - assert_json_int!(member["location"], "line", 3); } #[test] - fn get_declaration_module() { - let s = server_with_source("module Greetable; end"); - assert_eq!(get_declaration(&s, "Greetable")["kind"], "Module"); - } - - #[test] - fn get_declaration_doc_comments() { + fn get_descendants_json_shape() { let s = server_with_source( " - # The Animal class represents all animals. class Animal; end + class Dog < Animal; end ", ); - let res = get_declaration(&s, "Animal"); - let comments = array!(res["definitions"][0], "comments"); - assert!( - comments.iter().any(|c| c.as_str().unwrap().contains("Animal")), - "Expected doc comment on Animal, got: {comments:?}" - ); - } - - #[test] - fn get_declaration_mixin_ancestors() { - let s = server_with_source( - " - module Greetable; end - class Person - include Greetable - end - ", - ); - assert_includes!(get_declaration(&s, "Person"), "ancestors", "Greetable"); - } - - #[test] - fn get_declaration_constant() { - let s = server_with_source( - " - class Animal - KINGDOM = 'Animalia' - end - ", - ); - let res = get_declaration(&s, "Animal::KINGDOM"); - assert_eq!(res["kind"], "Constant"); - assert!(array!(res, "ancestors").is_empty()); - assert!(array!(res, "members").is_empty()); - } - - #[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", - ); + 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); } - // -- get_descendants -- - #[test] - fn get_descendants_with_subclasses() { + fn find_constant_references_json_shape() { let s = server_with_source( " class Animal; end class Dog < Animal; end - class Cat < Animal; end ", ); + let res = find_constant_references!(s, name: "Animal".into(), limit: None, offset: None); - let res = get_descendants!(s, name: "Animal".into(), limit: None, offset: None); assert_eq!(res["name"], "Animal"); - assert_includes!(res, "descendants", "Animal"); - assert_includes!(res, "descendants", "Dog"); - assert_includes!(res, "descendants", "Cat"); - assert_json_int!(res, "total", 3); - - // Cat: 1 descendant (itself only, no subclasses) - let res = get_descendants!(s, name: "Cat".into(), limit: None, offset: None); - assert_json_int!(res, "total", 1); + 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_descendants_module() { - let s = server_with_source( - " - module Greetable; end + fn get_file_declarations_json_shape() { + let s = server_with_source("class Animal; end"); + let res = get_file_declarations(&s, "test.rb"); - class Person - include Greetable - end - ", - ); - let res = get_descendants!(s, name: "Greetable".into(), limit: None, offset: None); - assert_includes!(res, "descendants", "Person"); + 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 get_descendants_inheritance_chain() { - let s = server_with_source( - " - class Foo; end - class Bar < Foo; end - class Baz < Bar; end - ", - ); - let res = get_descendants!(s, name: "Foo".into(), limit: None, offset: None); - assert_includes!(res, "descendants", "Bar"); - assert_includes!(res, "descendants", "Baz"); + 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_descendants_pagination() { - let s = server_with_source( - " - class Animal; end - class Dog < Animal; end - class Cat < Animal; end - ", + 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", ); - let page1 = get_descendants!(s, name: "Animal".into(), limit: Some(1), offset: Some(0)); - assert_eq!(array!(page1, "descendants").len(), 1); - assert_json_int!(page1, "total", 3); - - let page2 = get_descendants!(s, name: "Animal".into(), limit: Some(1), offset: Some(1)); - let name1 = array!(page1, "descendants")[0]["name"].as_str().unwrap(); - let name2 = array!(page2, "descendants")[0]["name"].as_str().unwrap(); - assert_ne!(name1, name2, "Page 1 and page 2 should return different descendants"); } #[test] @@ -965,82 +518,6 @@ mod tests { ); } - // -- find_constant_references -- - - #[test] - fn find_constant_references_success() { - let s = server_with_source( - " - class Animal; end - class Dog < Animal; end - class Kennel - def build - Animal.new - end - end - ", - ); - let res = find_constant_references!(s, name: "Animal".into(), limit: None, offset: None); - - assert_eq!(res["name"], "Animal"); - assert_eq!(array!(res, "references").len(), 2); - assert_json_int!(res, "total", 2); - let first_ref = &array!(res, "references")[0]; - assert!(first_ref["path"].as_str().unwrap().ends_with("test.rb")); - assert_json_int!(first_ref, "line", 2); - assert_json_int!(first_ref, "column", 13); - } - - #[test] - fn find_constant_references_cross_file() { - let models = test_uri("models.rb"); - let services = test_uri("services.rb"); - let s = server_with_sources(&[ - (&models, "class Dog; end"), - ( - &services, - " - class Kennel - def adopt - Dog.new - end - end - ", - ), - ]); - let res = find_constant_references!(s, name: "Dog".into(), limit: None, offset: None); - let paths: Vec<&str> = array!(res, "references") - .iter() - .filter_map(|r| r["path"].as_str()) - .collect(); - assert!( - paths.iter().any(|p| p.contains("services")), - "Expected cross-file ref from services, got: {paths:?}" - ); - } - - #[test] - fn find_constant_references_pagination() { - let s = server_with_source( - " - class Animal; end - class Dog < Animal; end - class Cat < Animal; end - class Kennel - def build - Animal.new - end - end - ", - ); - let full = find_constant_references!(s, name: "Animal".into(), limit: None, offset: None); - let full_total = full["total"].as_u64().unwrap(); - - let page = find_constant_references!(s, name: "Animal".into(), limit: Some(1), offset: Some(0)); - assert_eq!(array!(page, "references").len(), 1); - assert_json_int!(page, "total", full_total); - } - #[test] fn find_constant_references_not_found() { let s = server_with_source("class Dog; end"); @@ -1054,36 +531,6 @@ mod tests { ); } - // -- get_file_declarations -- - - #[test] - fn get_file_declarations_success() { - let s = server_with_source( - " - class Animal; end - class Dog < Animal; end - module Greetable; end - ", - ); - let res = get_file_declarations(&s, "test.rb"); - - assert_includes!(res, "declarations", "Animal"); - assert_includes!(res, "declarations", "Dog"); - assert_includes!(res, "declarations", "Greetable"); - assert_eq!(array!(res, "declarations")[0]["name"], "Animal"); - assert_eq!(array!(res, "declarations")[0]["kind"], "Class"); - assert_json_int!(array!(res, "declarations")[0], "line", 1); - } - - #[test] - fn get_file_declarations_multiple_files() { - let models = test_uri("models.rb"); - let services = test_uri("services.rb"); - let s = server_with_sources(&[(&models, "class Animal; end"), (&services, "class Kennel; end")]); - let res = get_file_declarations(&s, "services.rb"); - assert_includes!(res, "declarations", "Kennel"); - } - #[test] fn get_file_declarations_not_found() { let s = server_with_source("class Dog; end"); @@ -1095,30 +542,11 @@ mod tests { ); } - // -- codebase_stats -- - - #[test] - fn codebase_stats_returns_counts() { - 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_json_int!(res, "declarations", 5); - assert_json_int!(res, "definitions", 2); - - let breakdown = &res["breakdown_by_kind"]; - assert_json_int!(breakdown, "Class", 4); - assert_json_int!(breakdown, "Module", 1); - } - - // -- error states -- + // -- server state errors -- #[test] fn returns_indexing_error_when_graph_not_ready() { let server = RubydexServer::new("/test".to_string()); - // graph is None (still indexing) assert_error(&server.codebase_stats(), "indexing"); }