diff --git a/Cargo.lock b/Cargo.lock index 25b1b33cb2af53708c9cdc513e7520cdb730f6c8..82fc8b1bfe7c0cba939c0a716c6740b22f512bce 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -572,6 +572,7 @@ dependencies = [ "collections", "futures 0.3.31", "gpui", + "log", "mlua", "parking_lot", "project", diff --git a/crates/assistant_scripting/Cargo.toml b/crates/assistant_scripting/Cargo.toml index d1e06b0506c06f01dda64369ca87fef6fc79c8ab..2993b646f42a6481c1b227a4fca5eddd9f902033 100644 --- a/crates/assistant_scripting/Cargo.toml +++ b/crates/assistant_scripting/Cargo.toml @@ -17,6 +17,7 @@ anyhow.workspace = true collections.workspace = true futures.workspace = true gpui.workspace = true +log.workspace = true mlua.workspace = true parking_lot.workspace = true project.workspace = true diff --git a/crates/assistant_scripting/src/sandbox_preamble.lua b/crates/assistant_scripting/src/sandbox_preamble.lua index 03b0929b383591e950632bfc9e4127c23c984ea5..22cde87badf6fd29902fd87d4a591441c47d0072 100644 --- a/crates/assistant_scripting/src/sandbox_preamble.lua +++ b/crates/assistant_scripting/src/sandbox_preamble.lua @@ -13,7 +13,10 @@ sandbox.tostring = tostring sandbox.tonumber = tonumber sandbox.pairs = pairs sandbox.ipairs = ipairs + +-- Access to custom functions sandbox.search = search +sandbox.outline = outline -- Create a sandboxed version of LuaFileIO local io = {} diff --git a/crates/assistant_scripting/src/session.rs b/crates/assistant_scripting/src/session.rs index 59769b5cfb2918b76392728307e6fd4cce996d0a..08ed3f805a5165a6b1564624331b2c2230621d79 100644 --- a/crates/assistant_scripting/src/session.rs +++ b/crates/assistant_scripting/src/session.rs @@ -1,10 +1,11 @@ +use anyhow::anyhow; use collections::{HashMap, HashSet}; use futures::{ channel::{mpsc, oneshot}, pin_mut, SinkExt, StreamExt, }; use gpui::{AppContext, AsyncApp, Context, Entity, EventEmitter, SharedString, Task, WeakEntity}; -use mlua::{Lua, MultiValue, Table, UserData, UserDataMethods}; +use mlua::{ExternalResult, Lua, MultiValue, Table, UserData, UserDataMethods}; use parking_lot::Mutex; use project::{search::SearchQuery, Fs, Project}; use regex::Regex; @@ -129,9 +130,29 @@ impl ScriptSession { "search", lua.create_async_function({ let foreground_fns_tx = foreground_fns_tx.clone(); - let fs = fs.clone(); move |lua, regex| { - Self::search(lua, foreground_fns_tx.clone(), fs.clone(), regex) + let mut foreground_fns_tx = foreground_fns_tx.clone(); + let fs = fs.clone(); + async move { + Self::search(&lua, &mut foreground_fns_tx, fs, regex) + .await + .into_lua_err() + } + } + })?, + )?; + globals.set( + "outline", + lua.create_async_function({ + let root_dir = root_dir.clone(); + move |_lua, path| { + let mut foreground_fns_tx = foreground_fns_tx.clone(); + let root_dir = root_dir.clone(); + async move { + Self::outline(root_dir, &mut foreground_fns_tx, path) + .await + .into_lua_err() + } } })?, )?; @@ -211,27 +232,9 @@ impl ScriptSession { 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) - } + let path = match Self::parse_abs_path_in_root_dir(&root_dir, &path_str) { + Ok(path) => path, + Err(err) => return Ok((None, format!("{err}"))), }; // close method @@ -567,11 +570,11 @@ impl ScriptSession { } async fn search( - lua: Lua, - mut foreground_tx: mpsc::Sender, + lua: &Lua, + foreground_tx: &mut mpsc::Sender, fs: Arc, regex: String, - ) -> mlua::Result { + ) -> anyhow::Result
{ // TODO: Allow specification of these options. let search_query = SearchQuery::regex( ®ex, @@ -584,18 +587,17 @@ impl ScriptSession { ); let search_query = match search_query { Ok(query) => query, - Err(e) => return Err(mlua::Error::runtime(format!("Invalid search query: {}", e))), + Err(e) => return Err(anyhow!("Invalid search query: {}", e)), }; // TODO: Should use `search_query.regex`. The tool description should also be updated, // as it specifies standard regex. let search_regex = match Regex::new(®ex) { Ok(re) => re, - Err(e) => return Err(mlua::Error::runtime(format!("Invalid regex: {}", e))), + Err(e) => return Err(anyhow!("Invalid regex: {}", e)), }; - let mut abs_paths_rx = - Self::find_search_candidates(search_query, &mut foreground_tx).await?; + let mut abs_paths_rx = Self::find_search_candidates(search_query, foreground_tx).await?; let mut search_results: Vec
= Vec::new(); while let Some(path) = abs_paths_rx.next().await { @@ -643,7 +645,7 @@ impl ScriptSession { async fn find_search_candidates( search_query: SearchQuery, foreground_tx: &mut mpsc::Sender, - ) -> mlua::Result> { + ) -> anyhow::Result> { Self::run_foreground_fn( "finding search file candidates", foreground_tx, @@ -693,14 +695,62 @@ impl ScriptSession { }) }), ) - .await + .await? + } + + async fn outline( + root_dir: Option>, + foreground_tx: &mut mpsc::Sender, + path_str: String, + ) -> anyhow::Result { + let root_dir = root_dir + .ok_or_else(|| mlua::Error::runtime("cannot get outline without a root directory"))?; + let path = Self::parse_abs_path_in_root_dir(&root_dir, &path_str)?; + let outline = Self::run_foreground_fn( + "getting code outline", + foreground_tx, + Box::new(move |session, cx| { + cx.spawn(move |mut cx| async move { + // TODO: This will not use file content from `fs_changes`. It will also reflect + // user changes that have not been saved. + let buffer = session + .update(&mut cx, |session, cx| { + session + .project + .update(cx, |project, cx| project.open_local_buffer(&path, cx)) + })? + .await?; + buffer.update(&mut cx, |buffer, _cx| { + if let Some(outline) = buffer.snapshot().outline(None) { + Ok(outline) + } else { + Err(anyhow!("No outline for file {path_str}")) + } + }) + }) + }), + ) + .await? + .await??; + + Ok(outline + .items + .into_iter() + .map(|item| { + if item.text.contains('\n') { + log::error!("Outline item unexpectedly contains newline"); + } + format!("{}{}", " ".repeat(item.depth), item.text) + }) + .collect::>() + .join("\n")) } async fn run_foreground_fn( description: &str, foreground_tx: &mut mpsc::Sender, - function: Box, AsyncApp) -> anyhow::Result + Send>, - ) -> mlua::Result { + function: Box, AsyncApp) -> R + Send>, + ) -> anyhow::Result { let (response_tx, response_rx) = oneshot::channel(); let send_result = foreground_tx .send(ForegroundFn(Box::new(move |this, cx| { @@ -710,19 +760,34 @@ impl ScriptSession { match send_result { Ok(()) => (), Err(err) => { - return Err(mlua::Error::runtime(format!( - "Internal error while enqueuing work for {description}: {err}" - ))) + return Err(anyhow::Error::new(err).context(format!( + "Internal error while enqueuing work for {description}" + ))); } } match response_rx.await { - Ok(Ok(result)) => Ok(result), - Ok(Err(err)) => Err(mlua::Error::runtime(format!( - "Error while {description}: {err}" - ))), - Err(oneshot::Canceled) => Err(mlua::Error::runtime(format!( + Ok(result) => Ok(result), + Err(oneshot::Canceled) => Err(anyhow!( "Internal error: response oneshot was canceled while {description}." - ))), + )), + } + } + + fn parse_abs_path_in_root_dir(root_dir: &Path, path_str: &str) -> anyhow::Result { + let path = Path::new(&path_str); + if path.is_absolute() { + // Check if path starts with root_dir prefix without resolving symlinks + if path.starts_with(&root_dir) { + Ok(path.to_path_buf()) + } else { + Err(anyhow!( + "Error: Absolute path {} is outside the current working directory", + path_str + )) + } + } else { + // TODO: Does use of `../` break sandbox - is path canonicalization needed? + Ok(root_dir.join(path)) } } } diff --git a/crates/assistant_scripting/src/system_prompt.txt b/crates/assistant_scripting/src/system_prompt.txt index f02c870aa10e9e24d750f2ddd6a07fa72d995c53..359085b9dd4854776c90f89eff303adf8f5d65a0 100644 --- a/crates/assistant_scripting/src/system_prompt.txt +++ b/crates/assistant_scripting/src/system_prompt.txt @@ -16,13 +16,17 @@ tell you what the output was. Note that `io` only has `open`, and then the file it returns only has the methods read, write, and close - it doesn't have popen or anything else. -There will be a global called `search` which accepts a regex (it's implemented +There is a function 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 each string being a match that was found within the file). +There is a function called `outline` which accepts the path to a source file, +and returns a string where each line is a declaration. These lines are indented +with 2 spaces to indicate when a declaration is inside another. + When I send you the script output, do not thank me for running it, act as if you ran it yourself.