Add support for querying file outline in assistant script (#26351)

Release Notes:

- N/A

Cargo.lock

@@ -572,6 +572,7 @@ dependencies = [
  "collections",
  "futures 0.3.31",
  "gpui",
+ "log",
  "mlua",
  "parking_lot",
  "project",

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

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 = {}

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 = match Self::parse_abs_path_in_root_dir(&root_dir, &path_str) {
-        let path: PathBuf = {
+            Ok(path) => path,
-            let rust_path = Path::new(&path_str);
+            Err(err) => return Ok((None, format!("{err}"))),
-
-            // 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
@@ -567,11 +570,11 @@ impl ScriptSession {
     }
 
     async fn search(
-        lua: Lua,
+        lua: &Lua,
-        mut foreground_tx: mpsc::Sender<ForegroundFn>,
+        foreground_tx: &mut mpsc::Sender<ForegroundFn>,
         fs: Arc<dyn Fs>,
         regex: String,
-    ) -> mlua::Result<Table> {
+    ) -> anyhow::Result<Table> {
         // TODO: Allow specification of these options.
         let search_query = SearchQuery::regex(
             &regex,
@@ -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(&regex) {
             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 =
+        let mut abs_paths_rx = Self::find_search_candidates(search_query, foreground_tx).await?;
-            Self::find_search_candidates(search_query, &mut foreground_tx).await?;
 
         let mut search_results: Vec<Table> = 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<ForegroundFn>,
-    ) -> mlua::Result<mpsc::UnboundedReceiver<PathBuf>> {
+    ) -> anyhow::Result<mpsc::UnboundedReceiver<PathBuf>> {
         Self::run_foreground_fn(
             "finding search file candidates",
             foreground_tx,
@@ -693,14 +695,62 @@ impl ScriptSession {
                 })
             }),
         )
-        .await
+        .await?
+    }
+
+    async fn outline(
+        root_dir: Option<Arc<Path>>,
+        foreground_tx: &mut mpsc::Sender<ForegroundFn>,
+        path_str: String,
+    ) -> anyhow::Result<String> {
+        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::<Vec<String>>()
+            .join("\n"))
     }
 
     async fn run_foreground_fn<R: Send + 'static>(
         description: &str,
         foreground_tx: &mut mpsc::Sender<ForegroundFn>,
-        function: Box<dyn FnOnce(WeakEntity<Self>, AsyncApp) -> anyhow::Result<R> + Send>,
+        function: Box<dyn FnOnce(WeakEntity<Self>, AsyncApp) -> R + Send>,
-    ) -> mlua::Result<R> {
+    ) -> anyhow::Result<R> {
         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!(
+                return Err(anyhow::Error::new(err).context(format!(
-                    "Internal error while enqueuing work for {description}: {err}"
+                    "Internal error while enqueuing work for {description}"
-                )))
+                )));
             }
         }
         match response_rx.await {
-            Ok(Ok(result)) => Ok(result),
+            Ok(result) => Ok(result),
-            Ok(Err(err)) => Err(mlua::Error::runtime(format!(
+            Err(oneshot::Canceled) => Err(anyhow!(
-                "Error while {description}: {err}"
-            ))),
-            Err(oneshot::Canceled) => Err(mlua::Error::runtime(format!(
                 "Internal error: response oneshot was canceled while {description}."
-            ))),
+            )),
+        }
+    }
+
+    fn parse_abs_path_in_root_dir(root_dir: &Path, path_str: &str) -> anyhow::Result<PathBuf> {
+        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))
         }
     }
 }

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.