magi-code 0.96.1

use super::{
    ToolResult, ToolResultDisplay, ToolRuntime,
    args::{AST_GREP_DEFAULT_LIMIT, AST_GREP_STDOUT_MAX_BYTES, AstGrepArgs, AstGrepOperation},
    contract::{metadata_key as meta, tool_name},
    fs::ExistingPathPolicy,
    workspace::WorkspaceWalkOptions,
};
use crate::cancellation::AgentCancellation;
use serde_json::json;
use std::{
    path::{Path, PathBuf},
    time::{Duration, Instant},
};

const AST_GREP_TIMEOUT: Duration = Duration::from_secs(30);

mod embedded;
mod outline;
mod search;

#[derive(Default)]
struct AstGrepRunState {
    stdout: String,
    exit_code: i32,
    timed_out: bool,
    stdout_truncated: bool,
    workset_truncated: bool,
}

impl ToolRuntime {
    pub(super) fn ast_grep(
        &self,
        args: AstGrepArgs,
        cancellation: &AgentCancellation,
    ) -> anyhow::Result<ToolResult> {
        self.ast_grep_with_timeout(args, cancellation, AST_GREP_TIMEOUT)
    }

    fn ast_grep_with_timeout(
        &self,
        args: AstGrepArgs,
        cancellation: &AgentCancellation,
        timeout: Duration,
    ) -> anyhow::Result<ToolResult> {
        let args = args.validate()?;
        let path_arg = args.path.clone().unwrap_or_else(|| ".".to_string());
        let resolved = self.resolve_existing_path(
            &path_arg,
            ExistingPathPolicy::ast_grep(self.ast_grep_absolute_paths),
        )?;
        cancellation.check()?;
        if args.operation == AstGrepOperation::Search {
            return search::run(self, &args, &resolved, cancellation, timeout);
        }
        outline::run(self, &args, &resolved, cancellation, timeout)
    }
}

fn append_batch_output(aggregate: &mut String, batch: &str) {
    if aggregate.is_empty() || batch.is_empty() {
        aggregate.push_str(batch);
    } else {
        if !aggregate.ends_with('\n') {
            aggregate.push('\n');
        }
        aggregate.push_str(batch);
    }
}

struct AstGrepResultInput<'a> {
    stdout: &'a str,
    exit_code: i32,
    timed_out: bool,
    stdout_truncated: bool,
    workset_truncated: bool,
}

fn ast_grep_result(input: AstGrepResultInput<'_>, args: &AstGrepArgs, limit: usize) -> ToolResult {
    let AstGrepResultInput {
        stdout,
        exit_code,
        timed_out,
        stdout_truncated,
        workset_truncated,
    } = input;
    let no_matches = exit_code == 1 && stdout.trim().is_empty();
    let lines: Vec<&str> = stdout.lines().take(limit).collect();
    let truncated = stdout.lines().count() > lines.len() || stdout_truncated || workset_truncated;
    let success = !timed_out && (exit_code == 0 || no_matches);
    let mut metadata = json!({meta::ENGINE:"ast-grep", meta::EXIT_CODE:exit_code, meta::TIMED_OUT:timed_out, meta::MATCHES_RETURNED:lines.len(), meta::TRUNCATED:truncated, meta::STDOUT_TRUNCATED:stdout_truncated});
    if let Some(lang) = &args.language {
        metadata[meta::LANGUAGE] = json!(lang);
    }
    if args.rewrite.is_some() {
        metadata[meta::REWRITE] = json!(true);
    }
    ToolResult {
        tool_name: tool_name::AST_GREP.to_string(),
        success,
        content: lines.join("\n"),
        metadata,
        display: ToolResultDisplay::default(),
    }
}