tree-sitter replaces ast-grep

This commit is contained in:
Dhanji R. Prasanna
2025-11-05 13:56:23 +11:00
parent 71e9e46f74
commit f25a3d5e06
12 changed files with 1003 additions and 876 deletions

69
Cargo.lock generated
View File

@@ -318,14 +318,13 @@ dependencies = [
[[package]] [[package]]
name = "cc" name = "cc"
version = "1.2.43" version = "1.0.106"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "739eb0f94557554b3ca9a86d2d37bebd49c5e6d0c1d2bda35ba5bdac830befc2" checksum = "066fce287b1d4eafef758e89e09d724a24808a9196fe9756b8ca90e86d0719a2"
dependencies = [ dependencies = [
"find-msvc-tools",
"jobserver", "jobserver",
"libc", "libc",
"shlex", "once_cell",
] ]
[[package]] [[package]]
@@ -1136,12 +1135,6 @@ dependencies = [
"simd-adler32", "simd-adler32",
] ]
[[package]]
name = "find-msvc-tools"
version = "0.1.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "52051878f80a721bb68ebfbc930e07b65ba72f2da88968ea5c06fd6ca3d3a127"
[[package]] [[package]]
name = "flate2" name = "flate2"
version = "1.1.5" version = "1.1.5"
@@ -1398,7 +1391,13 @@ dependencies = [
"tokio-stream", "tokio-stream",
"tokio-util", "tokio-util",
"tracing", "tracing",
"tree-sitter",
"tree-sitter-javascript",
"tree-sitter-python",
"tree-sitter-rust",
"tree-sitter-typescript",
"uuid", "uuid",
"walkdir",
] ]
[[package]] [[package]]
@@ -3628,6 +3627,56 @@ dependencies = [
"tracing-log", "tracing-log",
] ]
[[package]]
name = "tree-sitter"
version = "0.22.6"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "df7cc499ceadd4dcdf7ec6d4cbc34ece92c3fa07821e287aedecd4416c516dca"
dependencies = [
"cc",
"regex",
]
[[package]]
name = "tree-sitter-javascript"
version = "0.21.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "8710a71bc6779e33811a8067bdda3ed08bed1733296ff915e44faf60f8c533d7"
dependencies = [
"cc",
"tree-sitter",
]
[[package]]
name = "tree-sitter-python"
version = "0.21.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "b4066c6cf678f962f8c2c4561f205945c84834cce73d981e71392624fdc390a9"
dependencies = [
"cc",
"tree-sitter",
]
[[package]]
name = "tree-sitter-rust"
version = "0.21.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "277690f420bf90741dea984f3da038ace46c4fe6047cba57a66822226cde1c93"
dependencies = [
"cc",
"tree-sitter",
]
[[package]]
name = "tree-sitter-typescript"
version = "0.21.2"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "ecb35d98a688378e56c18c9c159824fd16f730ccbea19aacf4f206e5d5438ed9"
dependencies = [
"cc",
"tree-sitter",
]
[[package]] [[package]]
name = "try-lock" name = "try-lock"
version = "0.2.5" version = "0.2.5"

View File

@@ -94,6 +94,7 @@ These commands give you fine-grained control over context management, allowing y
- Screenshot capture and window management - Screenshot capture and window management
- OCR text extraction from images and screen regions - OCR text extraction from images and screen regions
- Window listing and identification - Window listing and identification
- **Code Search**: Embedded tree-sitter for syntax-aware code search (Rust, Python, JavaScript, TypeScript) - see [Code Search Guide](docs/CODE_SEARCH.md)
- **Final Output**: Formatted result presentation - **Final Output**: Formatted result presentation
### Provider Flexibility ### Provider Flexibility

19
TODO
View File

@@ -1,19 +0,0 @@
next tasks
x get something working with autonomous mode
- g3d
- bug where it prints everything in a conversation turn all over again before final_output
x ui abstraction from core
- context token counting bug
- embedded model
- prompt rewriting
- generates status messages "ruffling feathers..."
- project description?
- treesitter + friends
x error where it just gives up turn
- "project" behaviors (read readme first)
- advance project mgmt
- git for reverting
- swarm
- ui tests / computer controller

View File

@@ -26,3 +26,11 @@ rand = "0.8"
regex = "1.0" regex = "1.0"
shellexpand = "3.1" shellexpand = "3.1"
serde_yaml = "0.9" serde_yaml = "0.9"
# tree-sitter for embedded code search
tree-sitter = "0.22"
tree-sitter-rust = "0.21"
tree-sitter-python = "0.21"
tree-sitter-javascript = "0.21"
tree-sitter-typescript = "0.21"
walkdir = "2.4"

View File

@@ -0,0 +1,58 @@
//! Inspect tree-sitter AST structure for Rust code
use tree_sitter::{Parser, Language};
fn print_tree(node: tree_sitter::Node, source: &str, indent: usize) {
let indent_str = " ".repeat(indent);
let node_text = &source[node.byte_range()];
let preview = if node_text.len() > 50 {
format!("{}...", &node_text[..50])
} else {
node_text.to_string()
};
println!(
"{}{} [{}:{}] '{}'",
indent_str,
node.kind(),
node.start_position().row + 1,
node.start_position().column + 1,
preview.replace('\n', "\\n")
);
let mut cursor = node.walk();
for child in node.children(&mut cursor) {
print_tree(child, source, indent + 1);
}
}
fn main() -> anyhow::Result<()> {
let source_code = r#"
pub async fn example_async() {
println!("Hello");
}
fn regular_function() {
println!("Regular");
}
pub async fn another_async(x: i32) -> Result<(), ()> {
Ok(())
}
"#;
println!("Source code:");
println!("{}", source_code);
println!("\n{}", "=".repeat(80));
println!("AST Structure:");
println!("{}\n", "=".repeat(80));
let mut parser = Parser::new();
let language: Language = tree_sitter_rust::language().into();
parser.set_language(&language)?;
let tree = parser.parse(source_code, None).unwrap();
print_tree(tree.root_node(), source_code, 0);
Ok(())
}

View File

@@ -0,0 +1,56 @@
//! Inspect tree-sitter AST structure for Python code
use tree_sitter::{Parser, Language};
fn print_tree(node: tree_sitter::Node, source: &str, indent: usize) {
let indent_str = " ".repeat(indent);
let node_text = &source[node.byte_range()];
let preview = if node_text.len() > 50 {
format!("{}...", &node_text[..50])
} else {
node_text.to_string()
};
println!(
"{}{} [{}:{}] '{}'",
indent_str,
node.kind(),
node.start_position().row + 1,
node.start_position().column + 1,
preview.replace('\n', "\\n")
);
let mut cursor = node.walk();
for child in node.children(&mut cursor) {
print_tree(child, source, indent + 1);
}
}
fn main() -> anyhow::Result<()> {
let source_code = r#"
def regular_function():
pass
async def async_function():
pass
class MyClass:
def method(self):
pass
"#;
println!("Source code:");
println!("{}", source_code);
println!("\n{}", "=".repeat(80));
println!("AST Structure:");
println!("{}\n", "=".repeat(80));
let mut parser = Parser::new();
let language: Language = tree_sitter_python::language().into();
parser.set_language(&language)?;
let tree = parser.parse(source_code, None).unwrap();
print_tree(tree.root_node(), source_code, 0);
Ok(())
}

View File

@@ -0,0 +1,43 @@
//! Test Python async query
use tree_sitter::{Parser, Query, QueryCursor, Language};
fn main() -> anyhow::Result<()> {
let source_code = r#"
def regular_function():
pass
async def async_function():
pass
"#;
let mut parser = Parser::new();
let language: Language = tree_sitter_python::language().into();
parser.set_language(&language)?;
let tree = parser.parse(source_code, None).unwrap();
// Try different queries
let queries = vec![
"(function_definition (async) name: (identifier) @name)",
"(function_definition (async))",
"(async)",
];
for query_str in queries {
println!("\nTrying query: {}", query_str);
match Query::new(&language, query_str) {
Ok(query) => {
let mut cursor = QueryCursor::new();
let matches = cursor.matches(&query, tree.root_node(), source_code.as_bytes());
let count = matches.count();
println!(" ✓ Valid query, found {} matches", count);
}
Err(e) => {
println!(" ✗ Invalid query: {}", e);
}
}
}
Ok(())
}

View File

@@ -1,787 +0,0 @@
//! Code search functionality using ast-grep for syntax-aware semantic searches
use anyhow::{anyhow, Result};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::collections::HashMap;
use std::process::Stdio;
use std::time::{Duration, Instant};
use tokio::io::{AsyncBufReadExt, BufReader};
use tokio::process::Command;
use tokio::sync::Semaphore;
use tracing::{debug, error, info, warn};
/// Maximum number of searches allowed per request
const MAX_SEARCHES: usize = 20;
/// Default timeout for individual searches in seconds
const DEFAULT_TIMEOUT_SECS: u64 = 60;
/// Default maximum concurrency
const DEFAULT_MAX_CONCURRENCY: usize = 4;
/// Default maximum matches per search
const DEFAULT_MAX_MATCHES: usize = 500;
/// Search specification for a single ast-grep search
#[derive(Debug, Clone, Deserialize)]
pub struct SearchSpec {
pub name: String,
pub mode: SearchMode,
// Pattern mode fields
pub pattern: Option<String>,
pub language: Option<String>,
// YAML mode fields
pub rule_yaml: Option<String>,
// Common fields
pub paths: Option<Vec<String>>,
pub globs: Option<Vec<String>>,
pub json_style: Option<JsonStyle>,
pub context: Option<u32>,
pub threads: Option<u32>,
pub include_metadata: Option<bool>,
pub no_ignore: Option<Vec<NoIgnoreType>>,
pub severity: Option<HashMap<String, SeverityLevel>>,
pub timeout_secs: Option<u64>,
}
/// Search mode: pattern or yaml
#[derive(Debug, Clone, Deserialize, PartialEq)]
#[serde(rename_all = "lowercase")]
pub enum SearchMode {
Pattern,
Yaml,
}
/// JSON output style
#[derive(Debug, Clone, Deserialize, PartialEq)]
#[serde(rename_all = "lowercase")]
pub enum JsonStyle {
Pretty,
Stream,
Compact,
}
impl Default for JsonStyle {
fn default() -> Self {
JsonStyle::Stream
}
}
/// No-ignore types
#[derive(Debug, Clone, Deserialize, PartialEq)]
#[serde(rename_all = "lowercase")]
pub enum NoIgnoreType {
Hidden,
Dot,
Exclude,
Global,
Parent,
Vcs,
}
/// Severity levels for YAML rules
#[derive(Debug, Clone, Deserialize, PartialEq)]
#[serde(rename_all = "lowercase")]
pub enum SeverityLevel {
Error,
Warning,
Info,
Hint,
Off,
}
/// Request structure for code search
#[derive(Debug, Deserialize)]
pub struct CodeSearchRequest {
pub searches: Vec<SearchSpec>,
pub max_concurrency: Option<usize>,
pub max_matches_per_search: Option<usize>,
}
/// Result of a single search
#[derive(Debug, Serialize)]
pub struct SearchResult {
pub name: String,
pub mode: String,
pub status: String,
pub cmd: Vec<String>,
pub match_count: Option<usize>,
pub truncated: Option<bool>,
pub matches: Option<Vec<Value>>,
pub stderr: Option<String>,
pub exit_code: Option<i32>,
pub duration_ms: u64,
}
/// Summary of all searches
#[derive(Debug, Serialize)]
pub struct SearchSummary {
pub completed: usize,
pub total: usize,
pub total_matches: usize,
pub duration_ms: u64,
}
/// Complete response structure
#[derive(Debug, Serialize)]
pub struct CodeSearchResponse {
pub summary: SearchSummary,
pub searches: Vec<SearchResult>,
}
/// YAML rule structure for validation
#[derive(Debug, Deserialize)]
struct YamlRule {
pub id: String,
pub language: String,
pub rule: Value,
}
/// Execute a batch of code searches using ast-grep
pub async fn execute_code_search(request: CodeSearchRequest) -> Result<CodeSearchResponse> {
let start_time = Instant::now();
// Validate request
if request.searches.is_empty() {
return Err(anyhow!("No searches specified"));
}
if request.searches.len() > MAX_SEARCHES {
return Err(anyhow!(
"Too many searches: {} (max: {})",
request.searches.len(),
MAX_SEARCHES
));
}
// Check if ast-grep is available
check_ast_grep_available().await?;
let max_concurrency = request.max_concurrency.unwrap_or(DEFAULT_MAX_CONCURRENCY);
let max_matches = request.max_matches_per_search.unwrap_or(DEFAULT_MAX_MATCHES);
// Create semaphore for concurrency control
let semaphore = std::sync::Arc::new(Semaphore::new(max_concurrency));
// Execute searches concurrently
let mut tasks = Vec::new();
for search in request.searches {
let sem = semaphore.clone();
let task = tokio::spawn(async move {
let _permit = sem.acquire().await.unwrap();
execute_single_search(search, max_matches).await
});
tasks.push(task);
}
// Wait for all searches to complete
let mut results = Vec::new();
let mut total_matches = 0;
let mut completed = 0;
for task in tasks {
match task.await {
Ok(result) => {
if result.status == "ok" {
completed += 1;
if let Some(count) = result.match_count {
total_matches += count;
}
}
results.push(result);
}
Err(e) => {
error!("Task join error: {}", e);
// Create an error result
results.push(SearchResult {
name: "unknown".to_string(),
mode: "unknown".to_string(),
status: "error".to_string(),
cmd: vec![],
match_count: None,
truncated: None,
matches: None,
stderr: Some(format!("Task execution error: {}", e)),
exit_code: None,
duration_ms: 0,
});
}
}
}
let total_duration = start_time.elapsed();
Ok(CodeSearchResponse {
summary: SearchSummary {
completed,
total: results.len(),
total_matches,
duration_ms: total_duration.as_millis() as u64,
},
searches: results,
})
}
/// Execute a single search
async fn execute_single_search(search: SearchSpec, max_matches: usize) -> SearchResult {
let start_time = Instant::now();
let timeout_secs = search.timeout_secs.unwrap_or(DEFAULT_TIMEOUT_SECS);
// Validate the search specification
if let Err(e) = validate_search_spec(&search) {
return SearchResult {
name: search.name,
mode: format!("{:?}", search.mode).to_lowercase(),
status: "error".to_string(),
cmd: vec![],
match_count: None,
truncated: None,
matches: None,
stderr: Some(format!("Validation error: {}", e)),
exit_code: None,
duration_ms: start_time.elapsed().as_millis() as u64,
};
}
// Build command
let cmd_args = match build_ast_grep_command(&search) {
Ok(args) => args,
Err(e) => {
return SearchResult {
name: search.name,
mode: format!("{:?}", search.mode).to_lowercase(),
status: "error".to_string(),
cmd: vec![],
match_count: None,
truncated: None,
matches: None,
stderr: Some(format!("Command build error: {}", e)),
exit_code: None,
duration_ms: start_time.elapsed().as_millis() as u64,
};
}
};
debug!("Executing ast-grep command: {:?}", cmd_args);
// Execute with timeout
let timeout_duration = Duration::from_secs(timeout_secs);
match tokio::time::timeout(timeout_duration, run_ast_grep_command(&cmd_args)).await {
Ok(Ok((stdout, stderr, exit_code))) => {
let duration_ms = start_time.elapsed().as_millis() as u64;
if exit_code == 0 {
// Parse JSON output
match parse_ast_grep_output(&stdout, max_matches) {
Ok((matches, truncated)) => {
SearchResult {
name: search.name,
mode: format!("{:?}", search.mode).to_lowercase(),
status: "ok".to_string(),
cmd: cmd_args,
match_count: Some(matches.len()),
truncated: Some(truncated),
matches: Some(matches),
stderr: if stderr.is_empty() { None } else { Some(stderr) },
exit_code: None,
duration_ms,
}
}
Err(e) => {
SearchResult {
name: search.name,
mode: format!("{:?}", search.mode).to_lowercase(),
status: "error".to_string(),
cmd: cmd_args,
match_count: None,
truncated: None,
matches: None,
stderr: Some(format!("JSON parse error: {}\nRaw output: {}", e, stdout)),
exit_code: Some(exit_code),
duration_ms,
}
}
}
} else {
SearchResult {
name: search.name,
mode: format!("{:?}", search.mode).to_lowercase(),
status: "error".to_string(),
cmd: cmd_args,
match_count: None,
truncated: None,
matches: None,
stderr: Some(stderr),
exit_code: Some(exit_code),
duration_ms,
}
}
}
Ok(Err(e)) => {
SearchResult {
name: search.name,
mode: format!("{:?}", search.mode).to_lowercase(),
status: "error".to_string(),
cmd: cmd_args,
match_count: None,
truncated: None,
matches: None,
stderr: Some(format!("Execution error: {}", e)),
exit_code: None,
duration_ms: start_time.elapsed().as_millis() as u64,
}
}
Err(_) => {
SearchResult {
name: search.name,
mode: format!("{:?}", search.mode).to_lowercase(),
status: "timeout".to_string(),
cmd: cmd_args,
match_count: None,
truncated: None,
matches: None,
stderr: Some(format!("Search timed out after {} seconds", timeout_secs)),
exit_code: None,
duration_ms: start_time.elapsed().as_millis() as u64,
}
}
}
}
/// Validate a search specification
fn validate_search_spec(search: &SearchSpec) -> Result<()> {
match search.mode {
SearchMode::Pattern => {
if search.pattern.is_none() || search.pattern.as_ref().unwrap().is_empty() {
return Err(anyhow!("Pattern mode requires non-empty 'pattern' field"));
}
}
SearchMode::Yaml => {
let rule_yaml = search.rule_yaml.as_ref()
.ok_or_else(|| anyhow!("YAML mode requires 'rule_yaml' field"))?;
if rule_yaml.is_empty() {
return Err(anyhow!("YAML mode requires non-empty 'rule_yaml' field"));
}
// Parse and validate YAML structure
let parsed: YamlRule = serde_yaml::from_str(rule_yaml)
.map_err(|e| anyhow!("Invalid YAML rule: {}", e))?;
if parsed.id.is_empty() {
return Err(anyhow!("YAML rule must have non-empty 'id' field"));
}
if parsed.language.is_empty() {
return Err(anyhow!("YAML rule must have non-empty 'language' field"));
}
// Validate language is supported (basic check)
validate_language(&parsed.language)?;
}
}
// Validate context range
if let Some(context) = search.context {
if context > 20 {
return Err(anyhow!("Context lines cannot exceed 20"));
}
}
Ok(())
}
/// Validate that a language is supported by ast-grep
fn validate_language(language: &str) -> Result<()> {
let supported_languages = [
"rust", "javascript", "typescript", "python", "java", "c", "cpp", "csharp",
"go", "html", "css", "json", "yaml", "xml", "bash", "kotlin", "swift",
"php", "ruby", "scala", "dart", "lua", "r", "sql", "dockerfile",
"Rust", "JavaScript", "TypeScript", "Python", "Java", "C", "Cpp", "CSharp",
"Go", "Html", "Css", "Json", "Yaml", "Xml", "Bash", "Kotlin", "Swift",
"Php", "Ruby", "Scala", "Dart", "Lua", "R", "Sql", "Dockerfile"
];
if !supported_languages.contains(&language) {
warn!("Language '{}' may not be supported by ast-grep", language);
}
Ok(())
}
/// Build ast-grep command arguments
fn build_ast_grep_command(search: &SearchSpec) -> Result<Vec<String>> {
let mut args = vec!["ast-grep".to_string()];
match search.mode {
SearchMode::Pattern => {
args.push("run".to_string());
// Add pattern
args.push("-p".to_string());
args.push(search.pattern.as_ref().unwrap().clone());
// Add language if specified
if let Some(ref lang) = search.language {
args.push("-l".to_string());
args.push(lang.clone());
}
}
SearchMode::Yaml => {
args.push("scan".to_string());
// Add inline rules
args.push("--inline-rules".to_string());
args.push(search.rule_yaml.as_ref().unwrap().clone());
// Add include-metadata if requested
if search.include_metadata.unwrap_or(false) {
args.push("--include-metadata".to_string());
}
// Add severity overrides
if let Some(ref severity_map) = search.severity {
for (rule_id, severity) in severity_map {
match severity {
SeverityLevel::Error => {
args.push("--error".to_string());
args.push(rule_id.clone());
}
SeverityLevel::Warning => {
args.push("--warning".to_string());
args.push(rule_id.clone());
}
SeverityLevel::Info => {
args.push("--info".to_string());
args.push(rule_id.clone());
}
SeverityLevel::Hint => {
args.push("--hint".to_string());
args.push(rule_id.clone());
}
SeverityLevel::Off => {
args.push("--off".to_string());
args.push(rule_id.clone());
}
}
}
}
}
}
// Add common arguments
// Add globs if specified
if let Some(ref globs) = search.globs {
if !globs.is_empty() {
args.push("--globs".to_string());
args.push(globs.join(","));
}
}
// Add context
if let Some(context) = search.context {
args.push("-C".to_string());
args.push(context.to_string());
}
// Add threads
if let Some(threads) = search.threads {
args.push("-j".to_string());
args.push(threads.to_string());
}
// Add JSON output style
let json_style = search.json_style.as_ref().unwrap_or(&JsonStyle::Stream);
let json_arg = match json_style {
JsonStyle::Pretty => "--json=pretty",
JsonStyle::Stream => "--json=stream",
JsonStyle::Compact => "--json=compact",
};
args.push(json_arg.to_string());
// Add no-ignore options
if let Some(ref no_ignore_list) = search.no_ignore {
for no_ignore_type in no_ignore_list {
let flag = match no_ignore_type {
NoIgnoreType::Hidden => "--no-ignore=hidden",
NoIgnoreType::Dot => "--no-ignore=dot",
NoIgnoreType::Exclude => "--no-ignore=exclude",
NoIgnoreType::Global => "--no-ignore=global",
NoIgnoreType::Parent => "--no-ignore=parent",
NoIgnoreType::Vcs => "--no-ignore=vcs",
};
args.push(flag.to_string());
}
}
// Add paths (default to current directory if none specified)
if let Some(ref paths) = search.paths {
if !paths.is_empty() {
args.extend(paths.clone());
} else {
args.push(".".to_string());
}
} else {
args.push(".".to_string());
}
Ok(args)
}
/// Run ast-grep command and capture output
async fn run_ast_grep_command(args: &[String]) -> Result<(String, String, i32)> {
let mut cmd = Command::new(&args[0]);
cmd.args(&args[1..]);
cmd.stdout(Stdio::piped());
cmd.stderr(Stdio::piped());
debug!("Running command: {:?}", args);
let mut child = cmd.spawn()
.map_err(|e| anyhow!("Failed to spawn ast-grep process: {}", e))?;
let stdout = child.stdout.take().unwrap();
let stderr = child.stderr.take().unwrap();
let stdout_reader = BufReader::new(stdout);
let stderr_reader = BufReader::new(stderr);
let stdout_task = tokio::spawn(async move {
let mut lines = stdout_reader.lines();
let mut output = String::new();
while let Ok(Some(line)) = lines.next_line().await {
if !output.is_empty() {
output.push('\n');
}
output.push_str(&line);
}
output
});
let stderr_task = tokio::spawn(async move {
let mut lines = stderr_reader.lines();
let mut output = String::new();
while let Ok(Some(line)) = lines.next_line().await {
if !output.is_empty() {
output.push('\n');
}
output.push_str(&line);
}
output
});
let status = child.wait().await
.map_err(|e| anyhow!("Failed to wait for ast-grep process: {}", e))?;
let stdout_output = stdout_task.await
.map_err(|e| anyhow!("Failed to read stdout: {}", e))?;
let stderr_output = stderr_task.await
.map_err(|e| anyhow!("Failed to read stderr: {}", e))?;
let exit_code = status.code().unwrap_or(-1);
Ok((stdout_output, stderr_output, exit_code))
}
/// Parse ast-grep JSON output
fn parse_ast_grep_output(output: &str, max_matches: usize) -> Result<(Vec<Value>, bool)> {
if output.trim().is_empty() {
return Ok((vec![], false));
}
let mut matches = Vec::new();
let mut truncated = false;
// Handle stream format (line-delimited JSON)
for line in output.lines() {
let line = line.trim();
if line.is_empty() {
continue;
}
match serde_json::from_str::<Value>(line) {
Ok(match_obj) => {
if matches.len() >= max_matches {
truncated = true;
break;
}
matches.push(match_obj);
}
Err(e) => {
debug!("Failed to parse JSON line '{}': {}", line, e);
// Try to parse the entire output as a single JSON array
match serde_json::from_str::<Vec<Value>>(output) {
Ok(array_matches) => {
let take_count = array_matches.len().min(max_matches);
let total_count = array_matches.len();
matches = array_matches.into_iter().take(take_count).collect();
truncated = take_count < total_count;
break;
}
Err(e2) => {
return Err(anyhow!(
"Failed to parse ast-grep output as line-delimited JSON or JSON array. Line error: {}, Array error: {}",
e, e2
));
}
}
}
}
}
Ok((matches, truncated))
}
/// Check if ast-grep is available and provide installation hints if not
async fn check_ast_grep_available() -> Result<()> {
match Command::new("ast-grep")
.arg("--version")
.output()
.await
{
Ok(output) => {
if output.status.success() {
let version = String::from_utf8_lossy(&output.stdout);
info!("Found ast-grep: {}", version.trim());
Ok(())
} else {
Err(anyhow!("ast-grep command failed: {}", String::from_utf8_lossy(&output.stderr)))
}
}
Err(_) => {
Err(anyhow!(
"ast-grep not found. Please install it using one of these methods:\n\n\
• Homebrew (macOS): brew install ast-grep\n\
• MacPorts (macOS): sudo port install ast-grep\n\
• Nix: nix-env -iA nixpkgs.ast-grep\n\
• Cargo: cargo install ast-grep\n\
• npm: npm install -g @ast-grep/cli\n\
• pip: pip install ast-grep\n\n\
For more installation options, visit: https://ast-grep.github.io/guide/quick-start.html"
))
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_validate_pattern_search() {
let search = SearchSpec {
name: "test".to_string(),
mode: SearchMode::Pattern,
pattern: Some("fn $NAME() {}".to_string()),
language: Some("rust".to_string()),
rule_yaml: None,
paths: None,
globs: None,
json_style: None,
context: None,
threads: None,
include_metadata: None,
no_ignore: None,
severity: None,
timeout_secs: None,
};
assert!(validate_search_spec(&search).is_ok());
}
#[test]
fn test_validate_yaml_search() {
let yaml_rule = r#"
id: test-rule
language: Rust
rule:
pattern: "fn $NAME() {}"
"#;
let search = SearchSpec {
name: "test".to_string(),
mode: SearchMode::Yaml,
pattern: None,
language: None,
rule_yaml: Some(yaml_rule.to_string()),
paths: None,
globs: None,
json_style: None,
context: None,
threads: None,
include_metadata: None,
no_ignore: None,
severity: None,
timeout_secs: None,
};
assert!(validate_search_spec(&search).is_ok());
}
#[test]
fn test_build_pattern_command() {
let search = SearchSpec {
name: "test".to_string(),
mode: SearchMode::Pattern,
pattern: Some("fn $NAME() {}".to_string()),
language: Some("rust".to_string()),
rule_yaml: None,
paths: Some(vec!["src/".to_string()]),
globs: None,
json_style: Some(JsonStyle::Stream),
context: Some(2),
threads: Some(4),
include_metadata: None,
no_ignore: None,
severity: None,
timeout_secs: None,
};
let cmd = build_ast_grep_command(&search).unwrap();
assert_eq!(cmd[0], "ast-grep");
assert_eq!(cmd[1], "run");
assert!(cmd.contains(&"-p".to_string()));
assert!(cmd.contains(&"fn $NAME() {}".to_string()));
assert!(cmd.contains(&"-l".to_string()));
assert!(cmd.contains(&"rust".to_string()));
assert!(cmd.contains(&"--json=stream".to_string()));
assert!(cmd.contains(&"-C".to_string()));
assert!(cmd.contains(&"2".to_string()));
assert!(cmd.contains(&"-j".to_string()));
assert!(cmd.contains(&"4".to_string()));
assert!(cmd.contains(&"src/".to_string()));
}
#[test]
fn test_parse_stream_json() {
let output = r#"{"file":"test.rs","text":"fn hello() {}"}
{"file":"test2.rs","text":"fn world() {}"}"#;
let (matches, truncated) = parse_ast_grep_output(output, 10).unwrap();
assert_eq!(matches.len(), 2);
assert!(!truncated);
assert_eq!(matches[0]["file"], "test.rs");
assert_eq!(matches[1]["file"], "test2.rs");
}
#[test]
fn test_parse_truncated_output() {
let output = r#"{"file":"test1.rs","text":"fn a() {}"}
{"file":"test2.rs","text":"fn b() {}"}
{"file":"test3.rs","text":"fn c() {}"}"#;
let (matches, truncated) = parse_ast_grep_output(output, 2).unwrap();
assert_eq!(matches.len(), 2);
assert!(truncated);
}
}

View File

@@ -0,0 +1,81 @@
//! Code search functionality using tree-sitter for syntax-aware searches
use anyhow::Result;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
mod searcher;
pub use searcher::TreeSitterSearcher;
/// Request for batch code searches
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CodeSearchRequest {
pub searches: Vec<SearchSpec>,
#[serde(default = "default_concurrency")]
pub max_concurrency: usize,
#[serde(default = "default_max_matches")]
pub max_matches_per_search: usize,
}
fn default_concurrency() -> usize {
4
}
fn default_max_matches() -> usize {
500
}
/// Individual search specification
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SearchSpec {
/// Name/label for this search
pub name: String,
/// tree-sitter query (S-expression format)
pub query: String,
/// Language: "rust", "python", "javascript", "typescript"
pub language: String,
/// Paths to search (default: current directory)
#[serde(default)]
pub paths: Vec<String>,
/// Lines of context around each match
#[serde(default)]
pub context_lines: usize,
}
/// Response containing all search results
#[derive(Debug, Serialize, Deserialize)]
pub struct CodeSearchResponse {
pub searches: Vec<SearchResult>,
pub total_matches: usize,
pub total_files_searched: usize,
}
/// Result for a single search
#[derive(Debug, Serialize, Deserialize)]
pub struct SearchResult {
pub name: String,
pub matches: Vec<Match>,
pub match_count: usize,
pub files_searched: usize,
#[serde(skip_serializing_if = "Option::is_none")]
pub error: Option<String>,
}
/// A single match
#[derive(Debug, Serialize, Deserialize)]
pub struct Match {
pub file: String,
pub line: usize,
pub column: usize,
pub text: String,
#[serde(skip_serializing_if = "HashMap::is_empty")]
pub captures: HashMap<String, String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub context: Option<String>,
}
/// Main entry point for code search
pub async fn execute_code_search(request: CodeSearchRequest) -> Result<CodeSearchResponse> {
let mut searcher = TreeSitterSearcher::new()?;
searcher.execute_search(request).await
}

View File

@@ -0,0 +1,267 @@
use super::{CodeSearchRequest, CodeSearchResponse, Match, SearchResult, SearchSpec};
use anyhow::{anyhow, Result};
use std::collections::HashMap;
use std::fs;
use std::path::Path;
use tree_sitter::{Language, Parser, Query, QueryCursor};
use walkdir::WalkDir;
pub struct TreeSitterSearcher {
parsers: HashMap<String, Parser>,
languages: HashMap<String, Language>,
}
impl TreeSitterSearcher {
pub fn new() -> Result<Self> {
let mut parsers = HashMap::new();
let mut languages = HashMap::new();
// Initialize Rust
{
let mut parser = Parser::new();
let language: Language = tree_sitter_rust::language().into();
parser
.set_language(&language)
.map_err(|e| anyhow!("Failed to set Rust language: {}", e))?;
parsers.insert("rust".to_string(), parser);
languages.insert("rust".to_string(), language);
}
// Initialize Python
{
let mut parser = Parser::new();
let language: Language = tree_sitter_python::language().into();
parser
.set_language(&language)
.map_err(|e| anyhow!("Failed to set Python language: {}", e))?;
parsers.insert("python".to_string(), parser);
languages.insert("python".to_string(), language);
}
// Initialize JavaScript
{
let mut parser = Parser::new();
let language: Language = tree_sitter_javascript::language().into();
parser
.set_language(&language)
.map_err(|e| anyhow!("Failed to set JavaScript language: {}", e))?;
parsers.insert("javascript".to_string(), parser);
// Create separate parser for "js" alias
let mut parser_js = Parser::new();
parser_js.set_language(&language)
.map_err(|e| anyhow!("Failed to set JavaScript language: {}", e))?;
parsers.insert("js".to_string(), parser_js);
languages.insert("javascript".to_string(), language.clone());
languages.insert("js".to_string(), language.clone());
}
// Initialize TypeScript
{
let mut parser = Parser::new();
let language: Language = tree_sitter_typescript::language_typescript().into();
parser
.set_language(&language)
.map_err(|e| anyhow!("Failed to set TypeScript language: {}", e))?;
parsers.insert("typescript".to_string(), parser);
// Create separate parser for "ts" alias
let mut parser_ts = Parser::new();
parser_ts.set_language(&language)
.map_err(|e| anyhow!("Failed to set TypeScript language: {}", e))?;
parsers.insert("ts".to_string(), parser_ts);
languages.insert("typescript".to_string(), language.clone());
languages.insert("ts".to_string(), language.clone());
}
if parsers.is_empty() {
return Err(anyhow!(
"No language parsers available. Enable at least one language feature."
));
}
Ok(Self { parsers, languages })
}
pub async fn execute_search(
&mut self,
request: CodeSearchRequest,
) -> Result<CodeSearchResponse> {
let mut all_results = Vec::new();
let mut total_matches = 0;
let mut total_files = 0;
// Execute searches sequentially (could parallelize with tokio::spawn if needed)
for spec in request.searches {
let result = self
.search_single(&spec, request.max_matches_per_search)
.await;
match result {
Ok(search_result) => {
total_matches += search_result.match_count;
total_files += search_result.files_searched;
all_results.push(search_result);
}
Err(e) => {
all_results.push(SearchResult {
name: spec.name.clone(),
matches: vec![],
match_count: 0,
files_searched: 0,
error: Some(e.to_string()),
});
}
}
}
Ok(CodeSearchResponse {
searches: all_results,
total_matches,
total_files_searched: total_files,
})
}
async fn search_single(
&mut self,
spec: &SearchSpec,
max_matches: usize,
) -> Result<SearchResult> {
// Get parser and language
let parser = self
.parsers
.get_mut(&spec.language)
.ok_or_else(|| anyhow!("Unsupported language: {}", spec.language))?;
let language = self
.languages
.get(&spec.language)
.ok_or_else(|| anyhow!("Language not found: {}", spec.language))?;
// Parse query
let query = Query::new(language, &spec.query)
.map_err(|e| anyhow!("Invalid query: {}", e))?;
let mut matches = Vec::new();
let mut files_searched = 0;
// Determine search paths
let search_paths = if spec.paths.is_empty() {
vec![".".to_string()]
} else {
spec.paths.clone()
};
// Walk directories and search files
for search_path in search_paths {
for entry in WalkDir::new(&search_path)
.follow_links(true)
.into_iter()
.filter_map(|e| e.ok())
{
if matches.len() >= max_matches {
break;
}
let path = entry.path();
if !path.is_file() {
continue;
}
// Check file extension matches language
if !Self::is_language_file(path, &spec.language) {
continue;
}
files_searched += 1;
// Read and parse file
if let Ok(source_code) = fs::read_to_string(path) {
if let Some(tree) = parser.parse(&source_code, None) {
let mut cursor = QueryCursor::new();
let query_matches = cursor.matches(
&query,
tree.root_node(),
source_code.as_bytes(),
);
for query_match in query_matches {
if matches.len() >= max_matches {
break;
}
// Extract captures
let mut captures_map = HashMap::new();
let mut match_text = String::new();
let mut match_line = 0;
let mut match_column = 0;
for capture in query_match.captures {
let capture_name = query.capture_names()[capture.index as usize];
let node = capture.node;
let text = &source_code[node.byte_range()];
captures_map.insert(capture_name.to_string(), text.to_string());
// Use first capture for position
if match_text.is_empty() {
match_text = text.to_string();
let start = node.start_position();
match_line = start.row + 1;
match_column = start.column + 1;
}
}
// Get context if requested
let context = if spec.context_lines > 0 {
Some(Self::get_context(
&source_code,
match_line,
spec.context_lines,
))
} else {
None
};
matches.push(Match {
file: path.display().to_string(),
line: match_line,
column: match_column,
text: match_text,
captures: captures_map,
context,
});
}
}
}
}
}
Ok(SearchResult {
name: spec.name.clone(),
match_count: matches.len(),
files_searched,
matches,
error: None,
})
}
fn is_language_file(path: &Path, language: &str) -> bool {
let ext = path.extension().and_then(|e| e.to_str());
match (language, ext) {
("rust", Some("rs")) => true,
("python", Some("py")) => true,
("javascript" | "js", Some("js" | "jsx" | "mjs")) => true,
("typescript" | "ts", Some("ts" | "tsx")) => true,
_ => false,
}
}
fn get_context(source: &str, line: usize, context_lines: usize) -> String {
let lines: Vec<&str> = source.lines().collect();
// line is 1-indexed, convert to 0-indexed
let line_idx = line.saturating_sub(1);
// Get context_lines before and after
let start = line_idx.saturating_sub(context_lines);
let end = (line_idx + context_lines + 1).min(lines.len());
lines[start..end].join("\n")
}
}

View File

@@ -1096,11 +1096,11 @@ If you create temporary files for verification, place these in a subdir named 't
For reading files, prioritize use of code_search tool use with multiple search requests per call instead of read_file, if it makes sense. For reading files, prioritize use of code_search tool use with multiple search requests per call instead of read_file, if it makes sense.
Additional examples for the 'code_search' tool: Additional examples for the 'code_search' tool:
- Example for pattern mode: {\"tool\": \"code_search\", \"args\": {\"searches\": [{\"name\": \"find_functions\", \"mode\": \"pattern\", \"pattern\": \"fn $NAME($$$ARGS) { $$$ }\", \"language\": \"rust\", \"paths\": [\"src/\"]}]}} - Find functions: {\"tool\": \"code_search\", \"args\": {\"searches\": [{\"name\": \"find_functions\", \"query\": \"(function_item name: (identifier) @name)\", \"language\": \"rust\", \"paths\": [\"src/\"]}]}}
- Example for YAML mode: {\"tool\": \"code_search\", \"args\": {\"searches\": [{\"name\": \"find_async\", \"mode\": \"yaml\", \"rule_yaml\": \"id: async-fn\nlanguage: Rust\nrule:\n pattern: async fn $NAME($$$) { $$$ }\"}]}} - Find async functions: {\"tool\": \"code_search\", \"args\": {\"searches\": [{\"name\": \"find_async\", \"query\": \"(function_item (function_modifiers) name: (identifier) @name)\", \"language\": \"rust\"}]}}
- Example for multiple searches: {\"tool\": \"code_search\", \"args\": {\"searches\": [{\"name\": \"funcs\", \"mode\": \"pattern\", \"pattern\": \"fn $NAME\", \"language\": \"rust\"}, {\"name\": \"structs\", \"mode\": \"pattern\", \"pattern\": \"struct $NAME\", \"language\": \"rust\"}]}} - Find structs: {\"tool\": \"code_search\", \"args\": {\"searches\": [{\"name\": \"structs\", \"query\": \"(struct_item name: (type_identifier) @name)\", \"language\": \"rust\"}]}}
- Example for passing optional args like \"context\": {\"tool\": \"code_search\", \"args\": {\"searches\": [{\"name\": \"funcs\", \"mode\": \"pattern\", \"context\": 3, \"pattern\": \"fn $NAME\", \"language\": \"rust\"}]} - Multiple searches: {\"tool\": \"code_search\", \"args\": {\"searches\": [{\"name\": \"funcs\", \"query\": \"(function_item name: (identifier) @name)\", \"language\": \"rust\"}, {\"name\": \"structs\", \"query\": \"(struct_item name: (type_identifier) @name)\", \"language\": \"rust\"}]}}
- Common optional args for searches: - With context lines: {\"tool\": \"code_search\", \"args\": {\"searches\": [{\"name\": \"funcs\", \"query\": \"(function_item name: (identifier) @name)\", \"language\": \"rust\", \"context_lines\": 3}]}}
- \"context\": 3 (show surrounding lines), - \"context\": 3 (show surrounding lines),
- \"json_style\": \"stream\" (for large results) - \"json_style\": \"stream\" (for large results)
@@ -1184,13 +1184,13 @@ Short description for providers without native calling specs:
- Format: {\"tool\": \"todo_write\", \"args\": {\"content\": \"- [ ] Task 1\\n- [ ] Task 2\"}} - Format: {\"tool\": \"todo_write\", \"args\": {\"content\": \"- [ ] Task 1\\n- [ ] Task 2\"}}
- Example: {\"tool\": \"todo_write\", \"args\": {\"content\": \"- [ ] Implement feature\\n - [ ] Write tests\\n - [ ] Run tests\"}} - Example: {\"tool\": \"todo_write\", \"args\": {\"content\": \"- [ ] Implement feature\\n - [ ] Write tests\\n - [ ] Run tests\"}}
- **code_search**: Batch syntax-aware searches via ast-grep. Supports up to 20 pattern or YAML-rule searches in parallel. - **code_search**: Syntax-aware code search using tree-sitter. Supports Rust, Python, JavaScript, TypeScript.
- Format: {\"tool\": \"code_search\", \"args\": {\"searches\": [{\"name\": \"search_label\", \"mode\": \"pattern|yaml\", ...}], \"max_concurrency\": 4, \"max_matches_per_search\": 500}} - Format: {\"tool\": \"code_search\", \"args\": {\"searches\": [{\"name\": \"label\", \"query\": \"tree-sitter query\", \"language\": \"rust|python|javascript|typescript\", \"paths\": [\"src/\"], \"context_lines\": 0}]}}
- Example for pattern mode: {\"tool\": \"code_search\", \"args\": {\"searches\": [{\"name\": \"find_functions\", \"mode\": \"pattern\", \"pattern\": \"fn $NAME($$$ARGS) { $$$ }\", \"language\": \"rust\", \"paths\": [\"src/\"]}]}} - Find functions: {\"tool\": \"code_search\", \"args\": {\"searches\": [{\"name\": \"find_functions\", \"query\": \"(function_item name: (identifier) @name)\", \"language\": \"rust\", \"paths\": [\"src/\"]}]}}
- Example for YAML mode: {\"tool\": \"code_search\", \"args\": {\"searches\": [{\"name\": \"find_async\", \"mode\": \"yaml\", \"rule_yaml\": \"id: async-fn\nlanguage: Rust\nrule:\n pattern: async fn $NAME($$$) { $$$ }\"}]}} - Find async functions: {\"tool\": \"code_search\", \"args\": {\"searches\": [{\"name\": \"find_async\", \"query\": \"(function_item (function_modifiers) name: (identifier) @name)\", \"language\": \"rust\"}]}}
- Example for multiple searches: {\"tool\": \"code_search\", \"args\": {\"searches\": [{\"name\": \"funcs\", \"mode\": \"pattern\", \"pattern\": \"fn $NAME\", \"language\": \"rust\"}, {\"name\": \"structs\", \"mode\": \"pattern\", \"pattern\": \"struct $NAME\", \"language\": \"rust\"}]}} - Find structs: {\"tool\": \"code_search\", \"args\": {\"searches\": [{\"name\": \"structs\", \"query\": \"(struct_item name: (type_identifier) @name)\", \"language\": \"rust\"}]}}
- Example for passing optional args like \"context\": {\"tool\": \"code_search\", \"args\": {\"searches\": [{\"name\": \"funcs\", \"mode\": \"pattern\", \"context\": 3, \"pattern\": \"fn $NAME\", \"language\": \"rust\"}]} - Multiple searches: {\"tool\": \"code_search\", \"args\": {\"searches\": [{\"name\": \"funcs\", \"query\": \"(function_item name: (identifier) @name)\", \"language\": \"rust\"}, {\"name\": \"structs\", \"query\": \"(struct_item name: (type_identifier) @name)\", \"language\": \"rust\"}]}}
- Common optional args for searches: - With context lines: {\"tool\": \"code_search\", \"args\": {\"searches\": [{\"name\": \"funcs\", \"query\": \"(function_item name: (identifier) @name)\", \"language\": \"rust\", \"context_lines\": 3}]}}
- \"context\": 3 (show surrounding lines), - \"context\": 3 (show surrounding lines),
- \"json_style\": \"stream\" (for large results) - \"json_style\": \"stream\" (for large results)
@@ -1908,7 +1908,7 @@ Template:
// Add code_search tool // Add code_search tool
tools.push(Tool { tools.push(Tool {
name: "code_search".to_string(), name: "code_search".to_string(),
description: "Batch syntax-aware searches via ast-grep. Supports up to 20 pattern or YAML-rule searches in parallel; returns JSON matches (stream-collated).".to_string(), description: "Batch syntax-aware code searches using embedded tree-sitter. Supports up to 20 searches in parallel for Rust, Python, JavaScript, and TypeScript. Uses tree-sitter query syntax (S-expressions).".to_string(),
input_schema: json!({ input_schema: json!({
"type": "object", "type": "object",
"properties": { "properties": {
@@ -1919,43 +1919,14 @@ Template:
"type": "object", "type": "object",
"properties": { "properties": {
"name": { "type": "string", "description": "Label for this search." }, "name": { "type": "string", "description": "Label for this search." },
"mode": { "query": { "type": "string", "description": "tree-sitter query in S-expression format (e.g., \"(function_item name: (identifier) @name)\")"},
"type": "string", "language": { "type": "string", "enum": ["rust", "python", "javascript", "typescript"], "description": "Programming language to search." },
"enum": ["pattern", "yaml"],
"description": "`pattern` uses `ast-grep run`; `yaml` uses `ast-grep scan --inline-rules`."
},
// pattern mode (fast one-off)
"pattern": { "type": "string", "description": "ast-grep pattern code (e.g., \"async fn $NAME($$$ARGS) { $$$ }\")"},
"language": { "type": "string", "description": "Optional language for pattern mode; ast-grep may infer from file extensions if omitted." },
// yaml mode (full rule object)
"rule_yaml": { "type": "string", "description": "A full YAML rule object text. Must include `id`, `language`, and `rule`." },
// targeting
"paths": { "type": "array", "items": { "type": "string" }, "description": "Paths/dirs to search. Defaults to current dir if empty." }, "paths": { "type": "array", "items": { "type": "string" }, "description": "Paths/dirs to search. Defaults to current dir if empty." },
"globs": { "type": "array", "items": { "type": "string" }, "description": "Optional include/exclude globs for CLI --globs." }, "context_lines": { "type": "integer", "minimum": 0, "maximum": 20, "default": 0, "description": "Lines of context to include around each match." }
// result formatting & performance knobs
"json_style": { "type": "string", "enum": ["pretty","stream","compact"], "default": "stream", "description": "Use stream for large codebases." },
"context": { "type": "integer", "minimum": 0, "maximum": 20, "default": 0, "description": "CLI -C context lines in text output; also affects JSON `lines` field." },
"threads": { "type": "integer", "minimum": 1, "description": "Optional override for ast-grep -j (per process)." },
"include_metadata": { "type": "boolean", "default": false, "description": "If yaml mode and rule has metadata, add --include-metadata." },
// robustness
"no_ignore": {
"type": "array",
"items": { "type": "string", "enum": ["hidden","dot","exclude","global","parent","vcs"] },
"description": "Forwarded to --no-ignore to bypass ignore files/hidden."
}, },
// severity overrides for yaml mode "required": ["name", "query", "language"]
"severity": {
"type": "object",
"additionalProperties": { "type": "string", "enum": ["error","warning","info","hint","off"] },
"description": "Optional map<ruleId, severity> -> passed via --error/--warning/--info/--hint/--off."
},
// per-search timeout seconds (default 60)
"timeout_secs": { "type": "integer", "minimum": 1, "default": 60 }
},
"required": ["name","mode"]
} }
}, },
// global concurrency & truncation
"max_concurrency": { "type": "integer", "minimum": 1, "default": 4 }, "max_concurrency": { "type": "integer", "minimum": 1, "default": 4 },
"max_matches_per_search": { "type": "integer", "minimum": 1, "default": 500 } "max_matches_per_search": { "type": "integer", "minimum": 1, "default": 500 }
}, },
@@ -2395,13 +2366,6 @@ Template:
// Check if we need to summarize before starting // Check if we need to summarize before starting
if self.context_window.should_summarize() { if self.context_window.should_summarize() {
info!(
"Context window at {}% ({}/{} tokens), triggering auto-summarization",
self.context_window.percentage_used() as u32,
self.context_window.used_tokens,
self.context_window.total_tokens
);
// Notify user about summarization // Notify user about summarization
self.ui_writer.print_context_status(&format!( self.ui_writer.print_context_status(&format!(
"\n🗜️ Context window reaching capacity ({}%). Creating summary...", "\n🗜️ Context window reaching capacity ({}%). Creating summary...",
@@ -4590,13 +4554,7 @@ Template:
} }
} }
Err(e) => { Err(e) => {
// Check if it's an ast-grep not found error and provide helpful message Ok(format!("❌ Code search failed: {}", e))
let error_msg = e.to_string();
if error_msg.contains("ast-grep not found") {
Ok(format!("{}", error_msg))
} else {
Ok(format!("❌ Code search failed: {}", error_msg))
}
} }
} }
} }

View File

@@ -0,0 +1,412 @@
//! Integration tests for tree-sitter code search
use g3_core::code_search::{execute_code_search, CodeSearchRequest, SearchSpec};
use std::fs;
#[tokio::test]
async fn test_find_async_functions() {
// Create a temporary test file
let test_dir = std::env::temp_dir().join("g3_test_code_search");
fs::create_dir_all(&test_dir).unwrap();
let test_file = test_dir.join("test.rs");
fs::write(
&test_file,
r#"
pub async fn example_async() {
println!("Hello");
}
fn regular_function() {
println!("Regular");
}
pub async fn another_async(x: i32) -> Result<(), ()> {
Ok(())
}
"#,
)
.unwrap();
// Test 1: Find async functions
let request = CodeSearchRequest {
searches: vec![SearchSpec {
name: "find_async_functions".to_string(),
// In tree-sitter-rust, async is a token inside function_modifiers
query: "(function_item (function_modifiers) name: (identifier) @name)".to_string(),
language: "rust".to_string(),
paths: vec![test_dir.to_string_lossy().to_string()],
context_lines: 0,
}],
max_concurrency: 4,
max_matches_per_search: 100,
};
let response = execute_code_search(request).await.unwrap();
assert_eq!(response.searches.len(), 1);
let search_result = &response.searches[0];
assert_eq!(search_result.name, "find_async_functions");
assert_eq!(search_result.match_count, 2, "Should find 2 async functions");
assert!(search_result.error.is_none());
// Check that we found the right functions
let function_names: Vec<String> = search_result
.matches
.iter()
.filter_map(|m| m.captures.get("name").cloned())
.collect();
assert!(function_names.contains(&"example_async".to_string()));
assert!(function_names.contains(&"another_async".to_string()));
// Cleanup
fs::remove_dir_all(&test_dir).ok();
}
#[tokio::test]
async fn test_find_all_functions() {
// Create a temporary test file
let test_dir = std::env::temp_dir().join("g3_test_code_search_2");
fs::create_dir_all(&test_dir).unwrap();
let test_file = test_dir.join("test.rs");
fs::write(
&test_file,
r#"
pub async fn example_async() {
println!("Hello");
}
fn regular_function() {
println!("Regular");
}
pub async fn another_async(x: i32) -> Result<(), ()> {
Ok(())
}
"#,
)
.unwrap();
// Test 2: Find all functions (async and regular)
let request = CodeSearchRequest {
searches: vec![SearchSpec {
name: "find_all_functions".to_string(),
query: "(function_item name: (identifier) @name)".to_string(),
language: "rust".to_string(),
paths: vec![test_dir.to_string_lossy().to_string()],
context_lines: 0,
}],
max_concurrency: 4,
max_matches_per_search: 100,
};
let response = execute_code_search(request).await.unwrap();
assert_eq!(response.searches.len(), 1);
let search_result = &response.searches[0];
assert_eq!(search_result.name, "find_all_functions");
assert_eq!(search_result.match_count, 3, "Should find 3 functions total");
assert!(search_result.error.is_none());
// Check that we found all functions
let function_names: Vec<String> = search_result
.matches
.iter()
.filter_map(|m| m.captures.get("name").cloned())
.collect();
assert!(function_names.contains(&"example_async".to_string()));
assert!(function_names.contains(&"regular_function".to_string()));
assert!(function_names.contains(&"another_async".to_string()));
// Cleanup
fs::remove_dir_all(&test_dir).ok();
}
#[tokio::test]
async fn test_find_structs() {
// Create a temporary test file
let test_dir = std::env::temp_dir().join("g3_test_code_search_3");
fs::create_dir_all(&test_dir).unwrap();
let test_file = test_dir.join("test.rs");
fs::write(
&test_file,
r#"
pub struct MyStruct {
field: String,
}
struct AnotherStruct;
enum MyEnum {
Variant,
}
"#,
)
.unwrap();
// Test 3: Find structs
let request = CodeSearchRequest {
searches: vec![SearchSpec {
name: "find_structs".to_string(),
query: "(struct_item name: (type_identifier) @name)".to_string(),
language: "rust".to_string(),
paths: vec![test_dir.to_string_lossy().to_string()],
context_lines: 0,
}],
max_concurrency: 4,
max_matches_per_search: 100,
};
let response = execute_code_search(request).await.unwrap();
assert_eq!(response.searches.len(), 1);
let search_result = &response.searches[0];
assert_eq!(search_result.name, "find_structs");
assert_eq!(search_result.match_count, 2, "Should find 2 structs");
assert!(search_result.error.is_none());
// Check that we found the right structs
let struct_names: Vec<String> = search_result
.matches
.iter()
.filter_map(|m| m.captures.get("name").cloned())
.collect();
assert!(struct_names.contains(&"MyStruct".to_string()));
assert!(struct_names.contains(&"AnotherStruct".to_string()));
// Cleanup
fs::remove_dir_all(&test_dir).ok();
}
#[tokio::test]
async fn test_context_lines() {
// Create a temporary test file
let test_dir = std::env::temp_dir().join("g3_test_code_search_4");
fs::create_dir_all(&test_dir).unwrap();
let test_file = test_dir.join("test.rs");
fs::write(
&test_file,
r#"
// Line 1
// Line 2
pub fn target_function() {
// Line 4
println!("target");
}
// Line 7
// Line 8
"#,
)
.unwrap();
// Test 4: Context lines
let request = CodeSearchRequest {
searches: vec![SearchSpec {
name: "find_with_context".to_string(),
query: "(function_item name: (identifier) @name)".to_string(),
language: "rust".to_string(),
paths: vec![test_dir.to_string_lossy().to_string()],
context_lines: 2,
}],
max_concurrency: 4,
max_matches_per_search: 100,
};
let response = execute_code_search(request).await.unwrap();
assert_eq!(response.searches.len(), 1);
let search_result = &response.searches[0];
assert_eq!(search_result.match_count, 1);
let match_result = &search_result.matches[0];
assert!(match_result.context.is_some());
let context = match_result.context.as_ref().unwrap();
assert!(context.contains("Line 2"), "Should include 2 lines before");
assert!(context.contains("target_function"), "Should include the function");
// Note: context_lines=2 means 2 lines before and after the match line (line 4)
// So we get lines 2-6, which includes up to println but not the closing brace
assert!(context.contains("println"), "Should include 2 lines after the match");
// Cleanup
fs::remove_dir_all(&test_dir).ok();
}
#[tokio::test]
async fn test_multiple_searches() {
// Create a temporary test file
let test_dir = std::env::temp_dir().join("g3_test_code_search_5");
fs::create_dir_all(&test_dir).unwrap();
let test_file = test_dir.join("test.rs");
fs::write(
&test_file,
r#"
pub async fn async_func() {}
fn regular_func() {}
pub struct MyStruct;
"#,
)
.unwrap();
// Test 5: Multiple searches in one request
let request = CodeSearchRequest {
searches: vec![
SearchSpec {
name: "async_functions".to_string(),
query: "(function_item (function_modifiers) name: (identifier) @name)".to_string(),
language: "rust".to_string(),
paths: vec![test_dir.to_string_lossy().to_string()],
context_lines: 0,
},
SearchSpec {
name: "structs".to_string(),
query: "(struct_item name: (type_identifier) @name)".to_string(),
language: "rust".to_string(),
paths: vec![test_dir.to_string_lossy().to_string()],
context_lines: 0,
},
],
max_concurrency: 4,
max_matches_per_search: 100,
};
let response = execute_code_search(request).await.unwrap();
assert_eq!(response.searches.len(), 2);
assert_eq!(response.total_matches, 2); // 1 async function + 1 struct
// Check first search (async functions)
let async_search = &response.searches[0];
assert_eq!(async_search.name, "async_functions");
assert_eq!(async_search.match_count, 1);
// Check second search (structs)
let struct_search = &response.searches[1];
assert_eq!(struct_search.name, "structs");
assert_eq!(struct_search.match_count, 1);
// Cleanup
fs::remove_dir_all(&test_dir).ok();
}
#[tokio::test]
async fn test_python_search() {
// Create a temporary Python test file
let test_dir = std::env::temp_dir().join("g3_test_code_search_python");
fs::create_dir_all(&test_dir).unwrap();
let test_file = test_dir.join("test.py");
fs::write(
&test_file,
r#"
def regular_function():
pass
async def async_function():
pass
class MyClass:
def method(self):
pass
"#,
)
.unwrap();
// Test 6: Python async functions
let request = CodeSearchRequest {
searches: vec![SearchSpec {
name: "python_async".to_string(),
// Note: tree-sitter-python doesn't expose 'async' as a queryable node
// For now, we'll just find all functions (async detection would need text matching)
query: "(function_definition name: (identifier) @name)".to_string(),
language: "python".to_string(),
paths: vec![test_dir.to_string_lossy().to_string()],
context_lines: 0,
}],
max_concurrency: 4,
max_matches_per_search: 100,
};
let response = execute_code_search(request).await.unwrap();
assert_eq!(response.searches.len(), 1);
let search_result = &response.searches[0];
assert_eq!(search_result.match_count, 3, "Should find 3 functions in Python (2 regular + 1 async + 1 method)");
let function_names: Vec<String> = search_result
.matches
.iter()
.filter_map(|m| m.captures.get("name").cloned())
.collect();
assert!(function_names.contains(&"regular_function".to_string()));
assert!(function_names.contains(&"async_function".to_string()));
assert!(function_names.contains(&"method".to_string()));
// Cleanup
fs::remove_dir_all(&test_dir).ok();
}
#[tokio::test]
async fn test_javascript_search() {
// Create a temporary JavaScript test file
let test_dir = std::env::temp_dir().join("g3_test_code_search_js");
fs::create_dir_all(&test_dir).unwrap();
let test_file = test_dir.join("test.js");
fs::write(
&test_file,
r#"
function regularFunction() {
console.log("regular");
}
async function asyncFunction() {
console.log("async");
}
class MyClass {
constructor() {}
}
"#,
)
.unwrap();
// Test 7: JavaScript functions
let request = CodeSearchRequest {
searches: vec![SearchSpec {
name: "js_functions".to_string(),
query: "(function_declaration name: (identifier) @name)".to_string(),
language: "javascript".to_string(),
paths: vec![test_dir.to_string_lossy().to_string()],
context_lines: 0,
}],
max_concurrency: 4,
max_matches_per_search: 100,
};
let response = execute_code_search(request).await.unwrap();
assert_eq!(response.searches.len(), 1);
let search_result = &response.searches[0];
assert_eq!(search_result.match_count, 2, "Should find 2 functions in JavaScript");
let function_names: Vec<String> = search_result
.matches
.iter()
.filter_map(|m| m.captures.get("name").cloned())
.collect();
assert!(function_names.contains(&"regularFunction".to_string()));
assert!(function_names.contains(&"asyncFunction".to_string()));
// Cleanup
fs::remove_dir_all(&test_dir).ok();
}