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

Release Notes:

- N/A

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<ForegroundFn>,
+        lua: &Lua,
+        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 =
-            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<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>,
-    ) -> mlua::Result<R> {
+        function: Box<dyn FnOnce(WeakEntity<Self>, AsyncApp) -> R + Send>,
+    ) -> 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!(
-                    "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<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))
         }
     }
 }