assistant: Display edits from scripts in panel (#26441)

https://github.com/user-attachments/assets/a486ff2a-4aa1-4c0d-be6c-1dea2a8d60c8
 
- [x] Track buffer changes in `ScriptingSession`
- [x] Show edited files in thread

Reviewing diffs and displaying line counts will be part of an upcoming
PR.

Release Notes:

- N/A

---------

Co-authored-by: Antonio Scandurra <me@as-cii.com>

crates/assistant2/src/message_editor.rs

@@ -2,6 +2,7 @@ use std::sync::Arc;
 
 use editor::actions::MoveUp;
 use editor::{Editor, EditorElement, EditorEvent, EditorStyle};
+use file_icons::FileIcons;
 use fs::Fs;
 use gpui::{
     Animation, AnimationExt, App, DismissEvent, Entity, Focusable, Subscription, TextStyle,
@@ -15,8 +16,8 @@ use std::time::Duration;
 use text::Bias;
 use theme::ThemeSettings;
 use ui::{
-    prelude::*, ButtonLike, KeyBinding, PlatformStyle, PopoverMenu, PopoverMenuHandle, Switch,
+    prelude::*, ButtonLike, Disclosure, KeyBinding, PlatformStyle, PopoverMenu, PopoverMenuHandle,
-    Tooltip,
+    Switch, Tooltip,
 };
 use vim_mode_setting::VimModeSetting;
 use workspace::Workspace;
@@ -39,6 +40,7 @@ pub struct MessageEditor {
     inline_context_picker_menu_handle: PopoverMenuHandle<ContextPicker>,
     model_selector: Entity<AssistantModelSelector>,
     use_tools: bool,
+    edits_expanded: bool,
     _subscriptions: Vec<Subscription>,
 }
 
@@ -117,6 +119,7 @@ impl MessageEditor {
                 )
             }),
             use_tools: false,
+            edits_expanded: false,
             _subscriptions: subscriptions,
         }
     }
@@ -303,6 +306,9 @@ impl Render for MessageEditor {
             px(64.)
         };
 
+        let changed_buffers = self.thread.read(cx).scripting_changed_buffers(cx);
+        let changed_buffers_count = changed_buffers.len();
+
         v_flex()
             .size_full()
             .when(is_streaming_completion, |parent| {
@@ -363,6 +369,109 @@ impl Render for MessageEditor {
                     ),
                 )
             })
+            .when(changed_buffers_count > 0, |parent| {
+                parent.child(
+                    v_flex()
+                        .mx_2()
+                        .bg(cx.theme().colors().element_background)
+                        .border_1()
+                        .border_b_0()
+                        .border_color(cx.theme().colors().border)
+                        .rounded_t_md()
+                        .child(
+                            h_flex()
+                                .gap_2()
+                                .p_2()
+                                .child(
+                                    Disclosure::new("edits-disclosure", self.edits_expanded)
+                                        .on_click(cx.listener(|this, _ev, _window, cx| {
+                                            this.edits_expanded = !this.edits_expanded;
+                                            cx.notify();
+                                        })),
+                                )
+                                .child(
+                                    Label::new("Edits")
+                                        .size(LabelSize::XSmall)
+                                        .color(Color::Muted),
+                                )
+                                .child(Label::new("•").size(LabelSize::XSmall).color(Color::Muted))
+                                .child(
+                                    Label::new(format!(
+                                        "{} {}",
+                                        changed_buffers_count,
+                                        if changed_buffers_count == 1 {
+                                            "file"
+                                        } else {
+                                            "files"
+                                        }
+                                    ))
+                                    .size(LabelSize::XSmall)
+                                    .color(Color::Muted),
+                                ),
+                        )
+                        .when(self.edits_expanded, |parent| {
+                            parent.child(
+                                v_flex().bg(cx.theme().colors().editor_background).children(
+                                    changed_buffers.enumerate().flat_map(|(index, buffer)| {
+                                        let file = buffer.read(cx).file()?;
+                                        let path = file.path();
+
+                                        let parent_label = path.parent().and_then(|parent| {
+                                            let parent_str = parent.to_string_lossy();
+
+                                            if parent_str.is_empty() {
+                                                None
+                                            } else {
+                                                Some(
+                                                    Label::new(format!(
+                                                        "{}{}",
+                                                        parent_str,
+                                                        std::path::MAIN_SEPARATOR_STR
+                                                    ))
+                                                    .color(Color::Muted)
+                                                    .size(LabelSize::Small),
+                                                )
+                                            }
+                                        });
+
+                                        let name_label = path.file_name().map(|name| {
+                                            Label::new(name.to_string_lossy().to_string())
+                                                .size(LabelSize::Small)
+                                        });
+
+                                        let file_icon = FileIcons::get_icon(&path, cx)
+                                            .map(Icon::from_path)
+                                            .unwrap_or_else(|| Icon::new(IconName::File));
+
+                                        let element = div()
+                                            .p_2()
+                                            .when(index + 1 < changed_buffers_count, |parent| {
+                                                parent
+                                                    .border_color(cx.theme().colors().border)
+                                                    .border_b_1()
+                                            })
+                                            .child(
+                                                h_flex()
+                                                    .gap_2()
+                                                    .child(file_icon)
+                                                    .child(
+                                                        // TODO: handle overflow
+                                                        h_flex()
+                                                            .children(parent_label)
+                                                            .children(name_label),
+                                                    )
+                                                    // TODO: show lines changed
+                                                    .child(Label::new("+").color(Color::Created))
+                                                    .child(Label::new("-").color(Color::Deleted)),
+                                            );
+
+                                        Some(element)
+                                    }),
+                                ),
+                            )
+                        }),
+                )
+            })
             .child(
                 v_flex()
                     .key_context("MessageEditor")

crates/assistant2/src/thread.rs

@@ -5,7 +5,7 @@ use assistant_tool::ToolWorkingSet;
 use chrono::{DateTime, Utc};
 use collections::{BTreeMap, HashMap, HashSet};
 use futures::StreamExt as _;
-use gpui::{App, Context, Entity, EventEmitter, SharedString, Task};
+use gpui::{App, AppContext, Context, Entity, EventEmitter, SharedString, Task};
 use language_model::{
     LanguageModel, LanguageModelCompletionEvent, LanguageModelRegistry, LanguageModelRequest,
     LanguageModelRequestMessage, LanguageModelRequestTool, LanguageModelToolResult,
@@ -13,7 +13,7 @@ use language_model::{
     Role, StopReason,
 };
 use project::Project;
-use scripting_tool::ScriptingTool;
+use scripting_tool::{ScriptingSession, ScriptingTool};
 use serde::{Deserialize, Serialize};
 use util::{post_inc, TryFutureExt as _};
 use uuid::Uuid;
@@ -76,6 +76,7 @@ pub struct Thread {
     project: Entity<Project>,
     tools: Arc<ToolWorkingSet>,
     tool_use: ToolUseState,
+    scripting_session: Entity<ScriptingSession>,
     scripting_tool_use: ToolUseState,
 }
 
@@ -83,8 +84,10 @@ impl Thread {
     pub fn new(
         project: Entity<Project>,
         tools: Arc<ToolWorkingSet>,
-        _cx: &mut Context<Self>,
+        cx: &mut Context<Self>,
     ) -> Self {
+        let scripting_session = cx.new(|cx| ScriptingSession::new(project.clone(), cx));
+
         Self {
             id: ThreadId::new(),
             updated_at: Utc::now(),
@@ -99,6 +102,7 @@ impl Thread {
             project,
             tools,
             tool_use: ToolUseState::new(),
+            scripting_session,
             scripting_tool_use: ToolUseState::new(),
         }
     }
@@ -108,7 +112,7 @@ impl Thread {
         saved: SavedThread,
         project: Entity<Project>,
         tools: Arc<ToolWorkingSet>,
-        _cx: &mut Context<Self>,
+        cx: &mut Context<Self>,
     ) -> Self {
         let next_message_id = MessageId(
             saved
@@ -121,6 +125,7 @@ impl Thread {
             ToolUseState::from_saved_messages(&saved.messages, |name| name != ScriptingTool::NAME);
         let scripting_tool_use =
             ToolUseState::from_saved_messages(&saved.messages, |name| name == ScriptingTool::NAME);
+        let scripting_session = cx.new(|cx| ScriptingSession::new(project.clone(), cx));
 
         Self {
             id,
@@ -144,6 +149,7 @@ impl Thread {
             project,
             tools,
             tool_use,
+            scripting_session,
             scripting_tool_use,
         }
     }
@@ -237,6 +243,13 @@ impl Thread {
         self.scripting_tool_use.tool_results_for_message(id)
     }
 
+    pub fn scripting_changed_buffers<'a>(
+        &self,
+        cx: &'a App,
+    ) -> impl ExactSizeIterator<Item = &'a Entity<language::Buffer>> {
+        self.scripting_session.read(cx).changed_buffers()
+    }
+
     pub fn message_has_tool_results(&self, message_id: MessageId) -> bool {
         self.tool_use.message_has_tool_results(message_id)
     }
@@ -637,7 +650,32 @@ impl Thread {
             .collect::<Vec<_>>();
 
         for scripting_tool_use in pending_scripting_tool_uses {
-            let task = ScriptingTool.run(scripting_tool_use.input, self.project.clone(), cx);
+            let task = match ScriptingTool::deserialize_input(scripting_tool_use.input) {
+                Err(err) => Task::ready(Err(err.into())),
+                Ok(input) => {
+                    let (script_id, script_task) =
+                        self.scripting_session.update(cx, move |session, cx| {
+                            session.run_script(input.lua_script, cx)
+                        });
+
+                    let session = self.scripting_session.clone();
+                    cx.spawn(|_, cx| async move {
+                        script_task.await;
+
+                        let message = session.read_with(&cx, |session, _cx| {
+                            // Using a id to get the script output seems impractical.
+                            // Why not just include it in the Task result?
+                            // This is because we'll later report the script state as it runs,
+                            session
+                                .get(script_id)
+                                .output_message_for_llm()
+                                .expect("Script shouldn't still be running")
+                        })?;
+
+                        Ok(message)
+                    })
+                }
+            };
 
             self.insert_scripting_tool_output(scripting_tool_use.id.clone(), task, cx);
         }

crates/scripting_tool/src/scripting_session.rs

@@ -5,6 +5,7 @@ use futures::{
     pin_mut, SinkExt, StreamExt,
 };
 use gpui::{AppContext, AsyncApp, Context, Entity, Task, WeakEntity};
+use language::Buffer;
 use mlua::{ExternalResult, Lua, MultiValue, Table, UserData, UserDataMethods};
 use parking_lot::Mutex;
 use project::{search::SearchQuery, Fs, Project, ProjectPath, WorktreeId};
@@ -19,9 +20,10 @@ struct ForegroundFn(Box<dyn FnOnce(WeakEntity<ScriptingSession>, AsyncApp) + Sen
 
 pub struct ScriptingSession {
     project: Entity<Project>,
+    scripts: Vec<Script>,
+    changed_buffers: HashSet<Entity<Buffer>>,
     foreground_fns_tx: mpsc::Sender<ForegroundFn>,
     _invoke_foreground_fns: Task<()>,
-    scripts: Vec<Script>,
 }
 
 impl ScriptingSession {
@@ -29,16 +31,21 @@ impl ScriptingSession {
         let (foreground_fns_tx, mut foreground_fns_rx) = mpsc::channel(128);
         ScriptingSession {
             project,
+            scripts: Vec::new(),
+            changed_buffers: HashSet::default(),
             foreground_fns_tx,
             _invoke_foreground_fns: cx.spawn(|this, cx| async move {
                 while let Some(foreground_fn) = foreground_fns_rx.next().await {
                     foreground_fn.0(this.clone(), cx.clone());
                 }
             }),
-            scripts: Vec::new(),
         }
     }
 
+    pub fn changed_buffers(&self) -> impl ExactSizeIterator<Item = &Entity<Buffer>> {
+        self.changed_buffers.iter()
+    }
+
     pub fn run_script(
         &mut self,
         script_src: String,
@@ -340,7 +347,6 @@ impl ScriptingSession {
         let write_perm = file_userdata.get::<bool>("__write_perm")?;
 
         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 text = {
@@ -383,12 +389,21 @@ impl ScriptingSession {
                         })?;
 
                         session
-                            .update(&mut cx, |session, cx| {
+                            .update(&mut cx, {
-                                session
+                                let buffer = buffer.clone();
-                                    .project
+
-                                    .update(cx, |project, cx| project.save_buffer(buffer, cx))
+                                |session, cx| {
+                                    session
+                                        .project
+                                        .update(cx, |project, cx| project.save_buffer(buffer, cx))
+                                }
                             })?
-                            .await
+                            .await?;
+
+                        // If we saved successfully, mark buffer as changed
+                        session.update(&mut cx, |session, _cx| {
+                            session.changed_buffers.insert(buffer);
+                        })
                     })
                 })
             }),
@@ -880,7 +895,6 @@ impl Script {
         }
     }
 }
-
 #[cfg(test)]
 mod tests {
     use gpui::TestAppContext;
@@ -897,7 +911,8 @@ mod tests {
             print("Goodbye", "moon!")
         "#;
 
-        let output = test_script(script, cx).await.unwrap();
+        let test_session = TestSession::init(cx).await;
+        let output = test_session.test_success(script, cx).await;
         assert_eq!(output, "Hello\tworld!\nGoodbye\tmoon!\n");
     }
 
@@ -916,7 +931,8 @@ mod tests {
             end
         "#;
 
-        let output = test_script(script, cx).await.unwrap();
+        let test_session = TestSession::init(cx).await;
+        let output = test_session.test_success(script, cx).await;
         assert_eq!(output, "File: /file1.txt\nMatches:\n  world\n");
     }
 
@@ -925,60 +941,129 @@ mod tests {
     #[gpui::test]
     async fn test_open_and_read_file(cx: &mut TestAppContext) {
         let script = r#"
-                local file = io.open("file1.txt", "r")
+            local file = io.open("file1.txt", "r")
-                local content = file:read()
+            local content = file:read()
-                print("Content:", content)
+            print("Content:", content)
-                file:close()
+            file:close()
-            "#;
+        "#;
 
-        let output = test_script(script, cx).await.unwrap();
+        let test_session = TestSession::init(cx).await;
+        let output = test_session.test_success(script, cx).await;
         assert_eq!(output, "Content:\tHello world!\n");
+
+        // Only read, should not be marked as changed
+        assert!(!test_session.was_marked_changed("file1.txt", cx));
     }
 
     #[gpui::test]
     async fn test_read_write_roundtrip(cx: &mut TestAppContext) {
         let script = r#"
-                local file = io.open("new_file.txt", "w")
+            local file = io.open("file1.txt", "w")
-                file:write("This is new content")
+            file:write("This is new content")
-                file:close()
+            file:close()
 
-                -- Read back to verify
+            -- Read back to verify
-                local read_file = io.open("new_file.txt", "r")
+            local read_file = io.open("file1.txt", "r")
-                if read_file then
+            local content = read_file:read("*a")
-                    local content = read_file:read("*a")
+            print("Written content:", content)
-                    print("Written content:", content)
+            read_file:close()
-                    read_file:close()
+        "#;
-                end
-            "#;
 
-        let output = test_script(script, cx).await.unwrap();
+        let test_session = TestSession::init(cx).await;
+        let output = test_session.test_success(script, cx).await;
         assert_eq!(output, "Written content:\tThis is new content\n");
+        assert!(test_session.was_marked_changed("file1.txt", cx));
     }
 
     #[gpui::test]
     async fn test_multiple_writes(cx: &mut TestAppContext) {
         let script = r#"
-                -- Test writing to a file multiple times
+            -- Test writing to a file multiple times
-                local file = io.open("multiwrite.txt", "w")
+            local file = io.open("multiwrite.txt", "w")
-                file:write("First line\n")
+            file:write("First line\n")
-                file:write("Second line\n")
+            file:write("Second line\n")
-                file:write("Third line")
+            file:write("Third line")
-                file:close()
+            file:close()
 
-                -- Read back to verify
+            -- Read back to verify
-                local read_file = io.open("multiwrite.txt", "r")
+            local read_file = io.open("multiwrite.txt", "r")
-                if read_file then
+            if read_file then
-                    local content = read_file:read("*a")
+                local content = read_file:read("*a")
-                    print("Full content:", content)
+                print("Full content:", content)
-                    read_file:close()
+                read_file:close()
-                end
+            end
-            "#;
+        "#;
 
-        let output = test_script(script, cx).await.unwrap();
+        let test_session = TestSession::init(cx).await;
+        let output = test_session.test_success(script, cx).await;
         assert_eq!(
             output,
             "Full content:\tFirst line\nSecond line\nThird line\n"
         );
+        assert!(test_session.was_marked_changed("multiwrite.txt", cx));
+    }
+
+    #[gpui::test]
+    async fn test_multiple_writes_diff_handles(cx: &mut TestAppContext) {
+        let script = r#"
+            -- Write to a file
+            local file1 = io.open("multi_open.txt", "w")
+            file1:write("Content written by first handle\n")
+            file1:close()
+
+            -- Open it again and add more content
+            local file2 = io.open("multi_open.txt", "w")
+            file2:write("Content written by second handle\n")
+            file2:close()
+
+            -- Open it a third time and read
+            local file3 = io.open("multi_open.txt", "r")
+            local content = file3:read("*a")
+            print("Final content:", content)
+            file3:close()
+        "#;
+
+        let test_session = TestSession::init(cx).await;
+        let output = test_session.test_success(script, cx).await;
+        assert_eq!(
+            output,
+            "Final content:\tContent written by second handle\n\n"
+        );
+        assert!(test_session.was_marked_changed("multi_open.txt", cx));
+    }
+
+    #[gpui::test]
+    async fn test_append_mode(cx: &mut TestAppContext) {
+        let script = r#"
+            -- Test append mode
+            local file = io.open("append.txt", "w")
+            file:write("Initial content\n")
+            file:close()
+
+            -- Append more content
+            file = io.open("append.txt", "a")
+            file:write("Appended content\n")
+            file:close()
+
+            -- Add even more
+            file = io.open("append.txt", "a")
+            file:write("More appended content")
+            file:close()
+
+            -- Read back to verify
+            local read_file = io.open("append.txt", "r")
+            local content = read_file:read("*a")
+            print("Content after appends:", content)
+            read_file:close()
+        "#;
+
+        let test_session = TestSession::init(cx).await;
+        let output = test_session.test_success(script, cx).await;
+        assert_eq!(
+            output,
+            "Content after appends:\tInitial content\nAppended content\nMore appended content\n"
+        );
+        assert!(test_session.was_marked_changed("append.txt", cx));
     }
 
     #[gpui::test]
@@ -1018,7 +1103,8 @@ mod tests {
             f:close()
         "#;
 
-        let output = test_script(script, cx).await.unwrap();
+        let test_session = TestSession::init(cx).await;
+        let output = test_session.test_success(script, cx).await;
         println!("{}", &output);
         assert!(output.contains("All:\tLine 1\nLine 2\nLine 3"));
         assert!(output.contains("Line 1:\tLine 1"));
@@ -1027,46 +1113,75 @@ mod tests {
         assert!(output.contains("Line with newline length:\t7"));
         assert!(output.contains("Last char:\t10")); // LF
         assert!(output.contains("5 bytes:\tLine "));
+        assert!(test_session.was_marked_changed("multiline.txt", cx));
     }
 
     // helpers
 
-    async fn test_script(source: &str, cx: &mut TestAppContext) -> anyhow::Result<String> {
+    struct TestSession {
-        init_test(cx);
+        session: Entity<ScriptingSession>,
-        let fs = FakeFs::new(cx.executor());
-        fs.insert_tree(
-            "/",
-            json!({
-                "file1.txt": "Hello world!",
-                "file2.txt": "Goodbye moon!"
-            }),
-        )
-        .await;
-
-        let project = Project::test(fs, [Path::new("/")], cx).await;
-        let session = cx.new(|cx| ScriptingSession::new(project, cx));
-
-        let (script_id, task) =
-            session.update(cx, |session, cx| session.run_script(source.to_string(), cx));
-
-        task.await;
-
-        Ok(session.read_with(cx, |session, _cx| {
-            let script = session.get(script_id);
-            let stdout = script.stdout_snapshot();
-
-            if let ScriptState::Failed { error, .. } = &script.state {
-                panic!("Script failed:\n{}\n\n{}", error, stdout);
-            }
-
-            stdout
-        }))
     }
 
-    fn init_test(cx: &mut TestAppContext) {
+    impl TestSession {
-        let settings_store = cx.update(SettingsStore::test);
+        async fn init(cx: &mut TestAppContext) -> Self {
-        cx.set_global(settings_store);
+            let settings_store = cx.update(SettingsStore::test);
-        cx.update(Project::init_settings);
+            cx.set_global(settings_store);
-        cx.update(language::init);
+            cx.update(Project::init_settings);
+            cx.update(language::init);
+
+            let fs = FakeFs::new(cx.executor());
+            fs.insert_tree(
+                "/",
+                json!({
+                    "file1.txt": "Hello world!",
+                    "file2.txt": "Goodbye moon!"
+                }),
+            )
+            .await;
+
+            let project = Project::test(fs.clone(), [Path::new("/")], cx).await;
+            let session = cx.new(|cx| ScriptingSession::new(project, cx));
+
+            TestSession { session }
+        }
+
+        async fn test_success(&self, source: &str, cx: &mut TestAppContext) -> String {
+            let script_id = self.run_script(source, cx).await;
+
+            self.session.read_with(cx, |session, _cx| {
+                let script = session.get(script_id);
+                let stdout = script.stdout_snapshot();
+
+                if let ScriptState::Failed { error, .. } = &script.state {
+                    panic!("Script failed:\n{}\n\n{}", error, stdout);
+                }
+
+                stdout
+            })
+        }
+
+        fn was_marked_changed(&self, path_str: &str, cx: &mut TestAppContext) -> bool {
+            self.session.read_with(cx, |session, cx| {
+                let count_changed = session
+                    .changed_buffers
+                    .iter()
+                    .filter(|buffer| buffer.read(cx).file().unwrap().path().ends_with(path_str))
+                    .count();
+
+                assert!(count_changed < 2, "Multiple buffers matched for same path");
+
+                count_changed > 0
+            })
+        }
+
+        async fn run_script(&self, source: &str, cx: &mut TestAppContext) -> ScriptId {
+            let (script_id, task) = self
+                .session
+                .update(cx, |session, cx| session.run_script(source.to_string(), cx));
+
+            task.await;
+
+            script_id
+        }
     }
 }

crates/scripting_tool/src/scripting_tool.rs

@@ -1,9 +1,7 @@
-mod session;
+mod scripting_session;
 
-use project::Project;
+pub use scripting_session::*;
-use session::*;
 
-use gpui::{App, AppContext as _, Entity, Task};
 use schemars::JsonSchema;
 use serde::Deserialize;
 
@@ -24,40 +22,9 @@ impl ScriptingTool {
         serde_json::to_value(&schema).unwrap()
     }
 
-    pub fn run(
+    pub fn deserialize_input(
-        &self,
         input: serde_json::Value,
-        project: Entity<Project>,
+    ) -> Result<ScriptingToolInput, serde_json::Error> {
-        cx: &mut App,
+        serde_json::from_value(input)
-    ) -> Task<anyhow::Result<String>> {
-        let input = match serde_json::from_value::<ScriptingToolInput>(input) {
-            Err(err) => return Task::ready(Err(err.into())),
-            Ok(input) => input,
-        };
-
-        // TODO: Store a session per thread
-        let session = cx.new(|cx| ScriptingSession::new(project, cx));
-        let lua_script = input.lua_script;
-
-        let (script_id, script_task) =
-            session.update(cx, |session, cx| session.run_script(lua_script, cx));
-
-        cx.spawn(|cx| async move {
-            script_task.await;
-
-            let message = session.read_with(&cx, |session, _cx| {
-                // Using a id to get the script output seems impractical.
-                // Why not just include it in the Task result?
-                // This is because we'll later report the script state as it runs,
-                // currently not supported by the `Tool` interface.
-                session
-                    .get(script_id)
-                    .output_message_for_llm()
-                    .expect("Script shouldn't still be running")
-            })?;
-
-            drop(session);
-            Ok(message)
-        })
     }
 }