Richard Feldmanwip1a15c631 year agoHistory
893 lines · 35.1 KB · Rust
RawDownload
use anyhow::anyhow;use assistant_tool::{Tool, ToolFileChanges, ToolRegistry};use collections::HashMap;use gpui::{    AnyElement, App, AppContext as _, HighlightStyle, StyledText, Task, WeakEntity, Window,};use language::Language;use mlua::{Function, HookTriggers, Lua, MultiValue, Result, UserData, UserDataMethods, VmState};use parking_lot::Mutex;use rich_text::{self, Highlight};use schemars::JsonSchema;use serde::Deserialize;use settings::Settings;use std::ops::Range;use std::{    cell::RefCell,    path::{Path, PathBuf},    rc::Rc,    sync::Arc,};use theme::ThemeSettings;use ui::prelude::*;use workspace::Workspace; pub fn init(cx: &App) {    let registry = ToolRegistry::global(cx);    registry.register_tool(ScriptingTool);} #[derive(Debug, Deserialize, JsonSchema)]struct ScriptingToolInput {    lua_script: String,} pub struct ScriptingTool; impl Tool for ScriptingTool {    fn name(&self) -> String {        "lua-interpreter".to_string()    }     fn description(&self) -> String {        r#"You can write a Lua script and I'll run it on my code base and tell you what its output was,including both stdout as well as the git diff of changes it made to the filesystem. That way,you can get more information about the code base, or make changes to the code base directly.The lua script will have access to `io` and it will run with the current working directory being inthe root of the code base, so you can use it to explore, search, make changes, etc. You can also havethe script print things, and I'll tell you what the output was. Note that `io` only has `open`, andthen the file it returns only has the methods read, write, and close - it doesn't have popen oranything else. `os` is not available, so don't try to use it.  There will be a global called`search` which accepts a regex (it's implemented using Rust's regex crate,so use that regex syntax) and runs that regex on the contents of every file in the code base(aside from gitignored files), then returns an array of tables with two fields: "path"(the path to the file that had the matches) and "matches" (an array of strings, with eachstring being a match that was found within the file). If I reference something in the code base that you aren't familiar with, I suggest writingmultiple `search` calls to find what you're looking for in one script invocation, as opposedto trying multiple times. When you're done, if you have any proposed changes to files,you should write them directly to disk in the script. I'll review them later."#.into()    }     fn input_schema(&self) -> serde_json::Value {        let schema = schemars::schema_for!(ScriptingToolInput);        serde_json::to_value(&schema).unwrap()    }     fn run(        self: Arc<Self>,        input: serde_json::Value,        thread_id: Arc<str>,        workspace: WeakEntity<Workspace>,        _window: &mut Window,        cx: &mut App,    ) -> Task<anyhow::Result<String>> {        let root_dir = workspace.update(cx, |workspace, cx| {            let first_worktree = workspace                .visible_worktrees(cx)                .next()                .ok_or_else(|| anyhow!("no worktrees"))?;            workspace                .absolute_path_of_worktree(first_worktree.read(cx).id(), cx)                .ok_or_else(|| anyhow!("no worktree root"))        });        let root_dir = match root_dir {            Ok(root_dir) => root_dir,            Err(err) => return Task::ready(Err(err)),        };        let root_dir = match root_dir {            Ok(root_dir) => root_dir,            Err(err) => return Task::ready(Err(err)),        };        let input = match serde_json::from_value::<ScriptingToolInput>(input) {            Err(err) => return Task::ready(Err(err.into())),            Ok(input) => input,        };        let lua_script = input.lua_script;        let fs_changes = ToolFileChanges::get(thread_id, cx);        cx.background_spawn(async move {            let output = run_sandboxed_lua(&lua_script, fs_changes, root_dir)                .map_err(|err| anyhow!(format!("{err}")))?                .printed_lines                .join("\n");             if output.is_empty() {                Ok("(The script had no output.)".to_string())            } else {                Ok(output.to_string())            }        })    }     fn render_input(        self: Arc<Self>,        input: serde_json::Value,        lua_language: Option<Arc<Language>>,        cx: &mut App,    ) -> AnyElement {        let theme_settings = ThemeSettings::get_global(cx);        let theme = cx.theme();         if let Ok(input) = serde_json::from_value::<ScriptingToolInput>(input.clone()) {            let label;             // Use Lua syntax highlighting, if available            if let Some(lua_language) = &lua_language.clone() {                let mut highlights = Vec::new();                let mut buf = String::new();                 rich_text::render_code(&mut buf, &mut highlights, &input.lua_script, lua_language);                 let gpui_highlights: Vec<(Range<usize>, HighlightStyle)> = highlights                    .iter()                    .map(|(range, highlight)| {                        let style = match highlight {                            Highlight::Code => Default::default(),                            Highlight::Id(id) => id.style(theme.syntax()).unwrap_or_default(),                            Highlight::InlineCode(_link) => Default::default(), // Links won't come up                            Highlight::Highlight(highlight) => *highlight,                             _ => HighlightStyle::default(),                        };                        (range.clone(), style)                    })                    .collect();                label = StyledText::new(buf)                    .with_highlights(gpui_highlights)                    .into_any_element();            } else {                label = Label::new(&input.lua_script).into_any_element();            };             div()                .when(lua_language.is_some(), |this| {                    this.bg(theme.colors().editor_background)                })                .font_family(theme_settings.buffer_font.family.clone())                .child(label)                .into_any_element()        } else {            // Fallback to JSON if lua_language is unavailable            Label::new(serde_json::to_string_pretty(&input).unwrap_or_default()).into_any_element()        }    }     fn render_output(self: Arc<Self>, output: SharedString, cx: &mut App) -> AnyElement {        let theme_settings = ThemeSettings::get_global(cx);         div()            .font_family(theme_settings.buffer_font.family.clone())            .child(output)            .into_any_element()    }} const SANDBOX_PREAMBLE: &str = include_str!("sandbox_preamble.lua"); struct FileContent(RefCell<Vec<u8>>); impl UserData for FileContent {    fn add_methods<M: UserDataMethods<Self>>(_methods: &mut M) {        // FileContent doesn't have any methods so far.    }} /// Sandboxed print() function in Lua.fn print(lua: &Lua, printed_lines: Rc<RefCell<Vec<String>>>) -> Result<Function> {    lua.create_function(move |_, args: MultiValue| {        let mut string = String::new();         for arg in args.into_iter() {            // Lua's `print()` prints tab characters between each argument.            if !string.is_empty() {                string.push('\t');            }             // If the argument's to_string() fails, have the whole function call fail.            string.push_str(arg.to_string()?.as_str())        }         printed_lines.borrow_mut().push(string);         Ok(())    })} fn search(    lua: &Lua,    _fs_changes: Arc<Mutex<HashMap<PathBuf, Vec<u8>>>>,    root_dir: PathBuf,) -> Result<Function> {    lua.create_function(move |lua, regex: String| {        use mlua::Table;        use regex::Regex;        use std::fs;         // Function to recursively search directory        let search_regex = match Regex::new(&regex) {            Ok(re) => re,            Err(e) => return Err(mlua::Error::runtime(format!("Invalid regex: {}", e))),        };         let mut search_results: Vec<Result<Table>> = Vec::new();         // Create an explicit stack for directories to process        let mut dir_stack = vec![root_dir.clone()];         while let Some(current_dir) = dir_stack.pop() {            // Process each entry in the current directory            let entries = match fs::read_dir(&current_dir) {                Ok(entries) => entries,                Err(e) => return Err(e.into()),            };             for entry_result in entries {                let entry = match entry_result {                    Ok(e) => e,                    Err(e) => return Err(e.into()),                };                 let path = entry.path();                 if path.is_dir() {                    // Skip .git directory and other common directories to ignore                    let dir_name = path.file_name().unwrap_or_default().to_string_lossy();                    if !dir_name.starts_with('.')                        && dir_name != "node_modules"                        && dir_name != "target"                    {                        // Instead of recursive call, add to stack                        dir_stack.push(path);                    }                } else if path.is_file() {                    // Skip binary files and very large files                    if let Ok(metadata) = fs::metadata(&path) {                        if metadata.len() > 1_000_000 {                            // Skip files larger than 1MB                            continue;                        }                    }                     // Attempt to read the file as text                    if let Ok(content) = fs::read_to_string(&path) {                        let mut matches = Vec::new();                         // Find all regex matches in the content                        for capture in search_regex.find_iter(&content) {                            matches.push(capture.as_str().to_string());                        }                         // If we found matches, create a result entry                        if !matches.is_empty() {                            let result_entry = lua.create_table()?;                            result_entry.set("path", path.to_string_lossy().to_string())?;                             let matches_table = lua.create_table()?;                            for (i, m) in matches.iter().enumerate() {                                matches_table.set(i + 1, m.clone())?;                            }                            result_entry.set("matches", matches_table)?;                             search_results.push(Ok(result_entry));                        }                    }                }            }        }         // Create a table to hold our results        let results_table = lua.create_table()?;        for (i, result) in search_results.into_iter().enumerate() {            match result {                Ok(entry) => results_table.set(i + 1, entry)?,                Err(e) => return Err(e),            }        }         Ok(results_table)    })} /// Sandboxed io.open() function in Lua.fn io_open(    lua: &Lua,    fs_changes: Arc<Mutex<HashMap<PathBuf, Vec<u8>>>>,    root_dir: PathBuf,) -> Result<Function> {    lua.create_function(move |lua, (path_str, mode): (String, Option<String>)| {        let mode = mode.unwrap_or_else(|| "r".to_string());         // Parse the mode string to determine read/write permissions        let read_perm = mode.contains('r');        let write_perm = mode.contains('w') || mode.contains('a') || mode.contains('+');        let append = mode.contains('a');        let truncate = mode.contains('w');         // This will be the Lua value returned from the `open` function.        let file = lua.create_table()?;         // Store file metadata in the file        file.set("__path", path_str.clone())?;        file.set("__mode", mode.clone())?;        file.set("__read_perm", read_perm)?;        file.set("__write_perm", write_perm)?;         // Sandbox the path; it must be within root_dir        let path: PathBuf = {            let rust_path = Path::new(&path_str);             // Get absolute path            if rust_path.is_absolute() {                // Check if path starts with root_dir prefix without resolving symlinks                if !rust_path.starts_with(&root_dir) {                    return Ok((                        None,                        format!(                            "Error: Absolute path {} is outside the current working directory",                            path_str                        ),                    ));                }                rust_path.to_path_buf()            } else {                // Make relative path absolute relative to cwd                root_dir.join(rust_path)            }        };         // close method        let close_fn = {            let fs_changes = fs_changes.clone();            lua.create_function(move |_lua, file_userdata: mlua::Table| {                let write_perm = file_userdata.get::<bool>("__write_perm")?;                let path = file_userdata.get::<String>("__path")?;                 if write_perm {                    // When closing a writable file, record the content                    let content = file_userdata.get::<mlua::AnyUserData>("__content")?;                    let content_ref = content.borrow::<FileContent>()?;                    let content_vec = content_ref.0.borrow();                     // Don't actually write to disk; instead, just update fs_changes.                    let path_buf = PathBuf::from(&path);                    fs_changes                        .lock()                        .insert(path_buf.clone(), content_vec.clone());                }                 Ok(true)            })?        };        file.set("close", close_fn)?;         // If it's a directory, give it a custom read() and return early.        if path.is_dir() {            // TODO handle the case where we changed it in the in-memory fs             // Create a special directory handle            file.set("__is_directory", true)?;             // Store directory entries            let entries = match std::fs::read_dir(&path) {                Ok(entries) => {                    let mut entry_names = Vec::new();                    for entry in entries.flatten() {                        entry_names.push(entry.file_name().to_string_lossy().into_owned());                    }                    entry_names                }                Err(e) => return Ok((None, format!("Error reading directory: {}", e))),            };             // Save the list of entries            file.set("__dir_entries", entries)?;            file.set("__dir_position", 0usize)?;             // Create a directory-specific read function            let read_fn = lua.create_function(|_lua, file_userdata: mlua::Table| {                let position = file_userdata.get::<usize>("__dir_position")?;                let entries = file_userdata.get::<Vec<String>>("__dir_entries")?;                 if position >= entries.len() {                    return Ok(None); // No more entries                }                 let entry = entries[position].clone();                file_userdata.set("__dir_position", position + 1)?;                 Ok(Some(entry))            })?;            file.set("read", read_fn)?;             // If we got this far, the directory was opened successfully            return Ok((Some(file), String::new()));        }         let is_in_changes;        let file_exists;        let file_content = {            let fs_changes = fs_changes.lock();            is_in_changes = fs_changes.contains_key(&path);            file_exists = is_in_changes || path.exists();             if file_exists && !truncate {                if is_in_changes {                    fs_changes.get(&path).unwrap().clone()                } else {                    drop(fs_changes); // Release the lock before starting the read.                     match std::fs::read(&path) {                        Ok(content) => content,                        Err(e) => return Ok((None, format!("Error reading file: {}", e))),                    }                }            } else {                Vec::new()            }        };         // If in append mode, position should be at the end        let position = if append && file_exists {            file_content.len()        } else {            0        };         file.set("__position", position)?;        file.set(            "__content",            lua.create_userdata(FileContent(RefCell::new(file_content)))?,        )?;         // Create file methods         // read method        let read_fn = {            lua.create_function(                |_lua, (file_userdata, format): (mlua::Table, Option<mlua::Value>)| {                    let read_perm = file_userdata.get::<bool>("__read_perm")?;                    if !read_perm {                        return Err(mlua::Error::runtime("File not open for reading"));                    }                     let content = file_userdata.get::<mlua::AnyUserData>("__content")?;                    let mut position = file_userdata.get::<usize>("__position")?;                    let content_ref = content.borrow::<FileContent>()?;                    let content_vec = content_ref.0.borrow();                     if position >= content_vec.len() {                        return Ok(None); // EOF                    }                     match format {                        Some(mlua::Value::String(s)) => {                            let lossy_string = s.to_string_lossy();                            let format_str: &str = lossy_string.as_ref();                             // Only consider the first 2 bytes, since it's common to pass e.g. "*all"  instead of "*a"                            match &format_str[0..2] {                                "*a" => {                                    // Read entire file from current position                                    let result = String::from_utf8_lossy(&content_vec[position..])                                        .to_string();                                    position = content_vec.len();                                    file_userdata.set("__position", position)?;                                    Ok(Some(result))                                }                                "*l" => {                                    // Read next line                                    let mut line = Vec::new();                                    let mut found_newline = false;                                     while position < content_vec.len() {                                        let byte = content_vec[position];                                        position += 1;                                         if byte == b'\n' {                                            found_newline = true;                                            break;

Lines 1–500 of 893.