diff --git a/crates/goose-mcp/src/developer/analyze/formatter.rs b/crates/goose-mcp/src/developer/analyze/formatter.rs index 504d44c1bfeb..c05011647820 100644 --- a/crates/goose-mcp/src/developer/analyze/formatter.rs +++ b/crates/goose-mcp/src/developer/analyze/formatter.rs @@ -179,9 +179,106 @@ impl Formatter { output.push('\n'); } + // References (type tracking) - only show if present + if !result.references.is_empty() { + Self::append_references(&mut output, result); + } + output } + /// Append reference tracking information (method-to-type associations, type usage) + fn append_references(output: &mut String, result: &AnalysisResult) { + use crate::developer::analyze::types::ReferenceType; + + // Group references by type + let mut method_defs = Vec::new(); + let mut type_inst = Vec::new(); + let mut field_types = Vec::new(); + let mut var_types = Vec::new(); + let mut param_types = Vec::new(); + + for ref_info in &result.references { + match ref_info.ref_type { + ReferenceType::MethodDefinition => method_defs.push(ref_info), + ReferenceType::TypeInstantiation => type_inst.push(ref_info), + ReferenceType::FieldType => field_types.push(ref_info), + ReferenceType::VariableType => var_types.push(ref_info), + ReferenceType::ParameterType => param_types.push(ref_info), + ReferenceType::Call | ReferenceType::Definition | ReferenceType::Import => {} + } + } + + // Only show section if we have non-call references + if method_defs.is_empty() + && type_inst.is_empty() + && field_types.is_empty() + && var_types.is_empty() + && param_types.is_empty() + { + return; + } + + output.push_str("\nR: "); + + let mut sections = Vec::new(); + + // Method definitions (methods associated with types) + if !method_defs.is_empty() { + let mut method_strs: Vec = method_defs + .iter() + .map(|r| { + if let Some(type_name) = &r.associated_type { + format!("{}({})", r.symbol, type_name) + } else { + r.symbol.clone() + } + }) + .collect(); + method_strs.sort(); + method_strs.dedup(); + sections.push(format!("methods[{}]", method_strs.join(" "))); + } + + // Type instantiations (struct literals) + if !type_inst.is_empty() { + let mut type_names: Vec = type_inst.iter().map(|r| r.symbol.clone()).collect(); + type_names.sort(); + type_names.dedup(); + sections.push(format!("types[{}]", type_names.join(" "))); + } + + // Field types (only show unique types, not all occurrences) + if !field_types.is_empty() { + let mut field_type_names: Vec = + field_types.iter().map(|r| r.symbol.clone()).collect(); + field_type_names.sort(); + field_type_names.dedup(); + sections.push(format!("fields[{}]", field_type_names.join(" "))); + } + + // Variable types (only show unique types) + if !var_types.is_empty() { + let mut var_type_names: Vec = + var_types.iter().map(|r| r.symbol.clone()).collect(); + var_type_names.sort(); + var_type_names.dedup(); + sections.push(format!("vars[{}]", var_type_names.join(" "))); + } + + // Parameter types (only show unique types) + if !param_types.is_empty() { + let mut param_type_names: Vec = + param_types.iter().map(|r| r.symbol.clone()).collect(); + param_type_names.sort(); + param_type_names.dedup(); + sections.push(format!("params[{}]", param_type_names.join(" "))); + } + + output.push_str(§ions.join("; ")); + output.push('\n'); + } + /// Format directory structure with summary pub fn format_directory_structure( base_path: &Path, diff --git a/crates/goose-mcp/src/developer/analyze/languages/mod.rs b/crates/goose-mcp/src/developer/analyze/languages/mod.rs index db9372ee2423..c9bae7bc5a78 100644 --- a/crates/goose-mcp/src/developer/analyze/languages/mod.rs +++ b/crates/goose-mcp/src/developer/analyze/languages/mod.rs @@ -41,6 +41,10 @@ type ExtractFunctionNameHandler = fn(&tree_sitter::Node, &str, &str) -> Option) -> Option; +/// Handler for finding the receiver type from a receiver node +/// Takes: (receiver_node, source) +type FindReceiverTypeHandler = fn(&tree_sitter::Node, &str) -> Option; + /// Language configuration containing all language-specific information /// /// This struct serves as a single source of truth for language support. @@ -61,6 +65,8 @@ pub struct LanguageInfo { pub extract_function_name_handler: Option, /// Optional handler for finding method names from receiver nodes pub find_method_for_receiver_handler: Option, + /// Optional handler for finding receiver type from receiver nodes + pub find_receiver_type_handler: Option, } /// Get language configuration for a given language @@ -76,15 +82,17 @@ pub fn get_language_info(language: &str) -> Option { function_name_kinds: &["identifier", "field_identifier", "property_identifier"], extract_function_name_handler: None, find_method_for_receiver_handler: None, + find_receiver_type_handler: None, }), "rust" => Some(LanguageInfo { element_query: rust::ELEMENT_QUERY, call_query: rust::CALL_QUERY, - reference_query: "", + reference_query: rust::REFERENCE_QUERY, function_node_kinds: &["function_item", "impl_item"], function_name_kinds: &["identifier", "field_identifier", "property_identifier"], extract_function_name_handler: Some(rust::extract_function_name_for_kind), - find_method_for_receiver_handler: None, + find_method_for_receiver_handler: Some(rust::find_method_for_receiver), + find_receiver_type_handler: Some(rust::find_receiver_type), }), "javascript" | "typescript" => Some(LanguageInfo { element_query: javascript::ELEMENT_QUERY, @@ -98,6 +106,7 @@ pub fn get_language_info(language: &str) -> Option { function_name_kinds: &["identifier", "field_identifier", "property_identifier"], extract_function_name_handler: None, find_method_for_receiver_handler: None, + find_receiver_type_handler: None, }), "go" => Some(LanguageInfo { element_query: go::ELEMENT_QUERY, @@ -107,6 +116,7 @@ pub fn get_language_info(language: &str) -> Option { function_name_kinds: &["identifier", "field_identifier", "property_identifier"], extract_function_name_handler: None, find_method_for_receiver_handler: Some(go::find_method_for_receiver), + find_receiver_type_handler: None, }), "java" => Some(LanguageInfo { element_query: java::ELEMENT_QUERY, @@ -116,6 +126,7 @@ pub fn get_language_info(language: &str) -> Option { function_name_kinds: &["identifier", "field_identifier", "property_identifier"], extract_function_name_handler: None, find_method_for_receiver_handler: None, + find_receiver_type_handler: None, }), "kotlin" => Some(LanguageInfo { element_query: kotlin::ELEMENT_QUERY, @@ -125,6 +136,7 @@ pub fn get_language_info(language: &str) -> Option { function_name_kinds: &["identifier", "field_identifier", "property_identifier"], extract_function_name_handler: None, find_method_for_receiver_handler: None, + find_receiver_type_handler: None, }), "swift" => Some(LanguageInfo { element_query: swift::ELEMENT_QUERY, @@ -139,6 +151,7 @@ pub fn get_language_info(language: &str) -> Option { function_name_kinds: &["simple_identifier"], extract_function_name_handler: Some(swift::extract_function_name_for_kind), find_method_for_receiver_handler: None, + find_receiver_type_handler: None, }), "ruby" => Some(LanguageInfo { element_query: ruby::ELEMENT_QUERY, @@ -148,6 +161,7 @@ pub fn get_language_info(language: &str) -> Option { function_name_kinds: &["identifier", "field_identifier", "property_identifier"], extract_function_name_handler: None, find_method_for_receiver_handler: Some(ruby::find_method_for_receiver), + find_receiver_type_handler: None, }), _ => None, } diff --git a/crates/goose-mcp/src/developer/analyze/languages/rust.rs b/crates/goose-mcp/src/developer/analyze/languages/rust.rs index 31c46cfd8f6d..bc1a980db1cf 100644 --- a/crates/goose-mcp/src/developer/analyze/languages/rust.rs +++ b/crates/goose-mcp/src/developer/analyze/languages/rust.rs @@ -27,6 +27,48 @@ pub const CALL_QUERY: &str = r#" macro: (identifier) @macro.call) "#; +/// Tree-sitter query for extracting Rust type references and usage patterns +pub const REFERENCE_QUERY: &str = r#" + ; Method receivers - capture self parameters to associate methods with impl types + (self_parameter) @method.receiver + + ; Struct instantiation - struct literals + (struct_expression + name: (type_identifier) @struct.literal) + + ; Field type declarations in structs + (field_declaration + type: (type_identifier) @field.type) + + ; Field with reference type + (field_declaration + type: (reference_type + (type_identifier) @field.type)) + + ; Field with generic type + (field_declaration + type: (generic_type + type: (type_identifier) @field.type)) + + ; Variable type annotations + (let_declaration + type: (type_identifier) @var.type) + + ; Variable with reference type + (let_declaration + type: (reference_type + (type_identifier) @var.type)) + + ; Function parameter types + (parameter + type: (type_identifier) @param.type) + + ; Parameter with reference type + (parameter + type: (reference_type + (type_identifier) @param.type)) +"#; + /// Extract function name for Rust-specific node kinds /// /// Rust has special cases like impl_item blocks that should be @@ -48,3 +90,55 @@ pub fn extract_function_name_for_kind( } None } + +/// Find the method name for a method receiver node in Rust +/// +/// The receiver_node is a self_parameter. This walks up to find the +/// containing function_item and returns the method name. +pub fn find_method_for_receiver( + receiver_node: &tree_sitter::Node, + source: &str, + _ast_recursion_limit: Option, +) -> Option { + // Walk up to find the function_item that contains this self_parameter + let mut current = *receiver_node; + + while let Some(parent) = current.parent() { + if parent.kind() == "function_item" { + // Found the function, get its name + for i in 0..parent.child_count() { + if let Some(child) = parent.child(i) { + if child.kind() == "identifier" { + return Some(source[child.byte_range()].to_string()); + } + } + } + } + current = parent; + } + None +} + +/// Find the receiver type for a self parameter in Rust +/// +/// In Rust, self parameters are special - they don't explicitly state their type. +/// This function walks up from a self_parameter node to find the impl block +/// and extracts the type being implemented. +pub fn find_receiver_type(node: &tree_sitter::Node, source: &str) -> Option { + // Walk up from self_parameter to find the impl_item + let mut current = *node; + while let Some(parent) = current.parent() { + if parent.kind() == "impl_item" { + // Find the type_identifier in the impl block + for i in 0..parent.child_count() { + if let Some(child) = parent.child(i) { + if child.kind() == "type_identifier" { + return Some(source[child.byte_range()].to_string()); + } + } + } + } + current = parent; + } + None +} diff --git a/crates/goose-mcp/src/developer/analyze/parser.rs b/crates/goose-mcp/src/developer/analyze/parser.rs index 9edb360ab9e9..d5cdd87d2d18 100644 --- a/crates/goose-mcp/src/developer/analyze/parser.rs +++ b/crates/goose-mcp/src/developer/analyze/parser.rs @@ -380,11 +380,19 @@ impl ElementExtractor { ast_recursion_limit, ); if let Some(method_name) = method_name { - ( - ReferenceType::MethodDefinition, - method_name, - Some(text.to_string()), - ) + // Use language-specific handler to find receiver type, or fall back to text + let type_name = Self::find_receiver_type(&node, source, language) + .or_else(|| Some(text.to_string())); + + if let Some(type_name) = type_name { + ( + ReferenceType::MethodDefinition, + method_name, + Some(type_name), + ) + } else { + continue; + } } else { continue; } @@ -428,6 +436,18 @@ impl ElementExtractor { .and_then(|handler| handler(receiver_node, source, ast_recursion_limit)) } + fn find_receiver_type( + receiver_node: &tree_sitter::Node, + source: &str, + language: &str, + ) -> Option { + use crate::developer::analyze::languages; + + languages::get_language_info(language) + .and_then(|info| info.find_receiver_type_handler) + .and_then(|handler| handler(receiver_node, source)) + } + fn find_containing_function( node: &tree_sitter::Node, source: &str, diff --git a/crates/goose-mcp/src/developer/analyze/tests/mod.rs b/crates/goose-mcp/src/developer/analyze/tests/mod.rs index f74ec3ba37ec..6da0e66d1c26 100644 --- a/crates/goose-mcp/src/developer/analyze/tests/mod.rs +++ b/crates/goose-mcp/src/developer/analyze/tests/mod.rs @@ -9,4 +9,5 @@ pub mod integration_tests; pub mod large_output_tests; pub mod parser_tests; pub mod ruby_test; +pub mod rust_test; pub mod traversal_tests; diff --git a/crates/goose-mcp/src/developer/analyze/tests/rust_test.rs b/crates/goose-mcp/src/developer/analyze/tests/rust_test.rs new file mode 100644 index 000000000000..d0bae3c91d4c --- /dev/null +++ b/crates/goose-mcp/src/developer/analyze/tests/rust_test.rs @@ -0,0 +1,179 @@ +use crate::developer::analyze::graph::CallGraph; +use crate::developer::analyze::parser::{ElementExtractor, ParserManager}; +use crate::developer::analyze::types::{AnalysisResult, ReferenceType}; +use std::collections::HashSet; +use std::path::PathBuf; + +fn parse_and_extract(code: &str) -> AnalysisResult { + let manager = ParserManager::new(); + let tree = manager.parse(code, "rust").unwrap(); + ElementExtractor::extract_with_depth(&tree, code, "rust", "semantic", None).unwrap() +} + +fn build_test_graph(files: Vec<(&str, &str)>) -> CallGraph { + let manager = ParserManager::new(); + let results: Vec<_> = files + .iter() + .map(|(path, code)| { + let tree = manager.parse(code, "rust").unwrap(); + let result = + ElementExtractor::extract_with_depth(&tree, code, "rust", "semantic", None) + .unwrap(); + (PathBuf::from(*path), result) + }) + .collect(); + CallGraph::build_from_results(&results) +} + +#[test] +fn test_rust_self_parameter_type_resolution() { + // Test that self parameters correctly resolve to their impl type + let code = r#" +struct MyStruct { + value: i32, +} + +impl MyStruct { + fn method_with_self(&self) -> i32 { + self.value + } + + fn method_with_mut_self(&mut self) { + self.value += 1; + } + + fn associated_function() -> Self { + MyStruct { value: 0 } + } +} +"#; + + let result = parse_and_extract(code); + + // Find method references with self parameters + let self_methods: Vec<_> = result + .references + .iter() + .filter(|r| r.ref_type == ReferenceType::MethodDefinition) + .collect(); + + // Should find both methods with self parameters + assert_eq!( + self_methods.len(), + 2, + "Expected 2 methods with self parameters" + ); + + // Both should be associated with MyStruct + for method_ref in &self_methods { + assert_eq!( + method_ref.associated_type.as_deref(), + Some("MyStruct"), + "Method {} should be associated with MyStruct", + method_ref.symbol + ); + } + + // Verify the specific methods + let method_names: HashSet<_> = self_methods.iter().map(|r| r.symbol.as_str()).collect(); + assert!(method_names.contains("method_with_self")); + assert!(method_names.contains("method_with_mut_self")); +} + +#[test] +fn test_rust_struct_and_impl_tracking() { + let code = r#" +struct Config { + host: String, + port: u16, +} + +struct Handler { + cfg: Config, +} + +impl Handler { + fn new(cfg: Config) -> Self { + Handler { cfg } + } + + fn start(&self) -> Result<(), String> { + Ok(()) + } +} + +fn main() { + let cfg = Config { host: "localhost".to_string(), port: 8080 }; + let handler = Handler::new(cfg); + let _ = handler.start(); +} +"#; + + let result = parse_and_extract(code); + let graph = build_test_graph(vec![("test.rs", code)]); + + // Test struct extraction (includes impl blocks) + assert_eq!(result.class_count, 3); // Config, Handler, impl Handler + let struct_names: HashSet<_> = result.classes.iter().map(|c| c.name.as_str()).collect(); + assert!(struct_names.contains("Config")); + assert!(struct_names.contains("Handler")); + + // Test method extraction + let method_names: HashSet<_> = result.functions.iter().map(|f| f.name.as_str()).collect(); + assert!(method_names.contains("new")); + assert!(method_names.contains("start")); + assert!(method_names.contains("main")); + + // Test method-to-type associations (only methods with self parameter) + let handler_methods: Vec<_> = result + .references + .iter() + .filter(|r| { + r.ref_type == ReferenceType::MethodDefinition + && r.associated_type.as_deref() == Some("Handler") + }) + .collect(); + assert!( + !handler_methods.is_empty(), + "Expected at least 1 method on Handler (start), found {}", + handler_methods.len() + ); + + // Verify the method is 'start' (new doesn't have self, so it's not tracked) + assert!( + handler_methods.iter().any(|r| r.symbol == "start"), + "Expected to find 'start' method on Handler" + ); + + // Test field type tracking + let field_type_refs: Vec<_> = result + .references + .iter() + .filter(|r| r.ref_type == ReferenceType::FieldType) + .collect(); + assert!( + !field_type_refs.is_empty(), + "Expected to find field type references" + ); + + // Test struct instantiation + let config_literals: Vec<_> = result + .references + .iter() + .filter(|r| r.symbol == "Config" && r.ref_type == ReferenceType::TypeInstantiation) + .collect(); + assert!( + !config_literals.is_empty(), + "Expected to find Config struct literals" + ); + + // Test call graph integration + let incoming = graph.find_incoming_chains("Handler", 1); + assert!( + !incoming.is_empty(), + "Expected to find incoming references to Handler" + ); + + let outgoing = graph.find_outgoing_chains("Handler", 1); + assert!(!outgoing.is_empty(), "Expected to find methods on Handler"); +}