From 7c3eecc9c737b23f048ca82f9cefa8160c8cd35c Mon Sep 17 00:00:00 2001 From: Michael Sloan Date: Sun, 9 Mar 2025 23:26:17 -0600 Subject: [PATCH] Add support for querying file outline in assistant script (#26351) Release Notes: - N/A --- Cargo.lock | 1 + crates/assistant_scripting/Cargo.toml | 1 + .../src/sandbox_preamble.lua | 3 + crates/assistant_scripting/src/session.rs | 153 +++++++++++++----- .../assistant_scripting/src/system_prompt.txt | 6 +- 5 files changed, 119 insertions(+), 45 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 25b1b33cb2af53..82fc8b1bfe7c0c 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 d1e06b0506c06f..2993b646f42a64 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 03b0929b383591..22cde87badf6fd 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 59769b5cfb2918..08ed3f805a5165 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 f02c870aa10e9e..359085b9dd4854 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.