assistant: Make scripting a first-class concept instead of a tool (#26338)

This PR makes refactors the scripting functionality to be a first-class
concept of the assistant instead of a generic tool, which will allow us
to build a more customized experience.

- The tool prompt has been slightly tweaked and is now included as a
system message in all conversations. I'm getting decent results, but now
that it isn't in the tools framework, it will probably require more
refining.

- The model will now include an `<eval ...>` tag at the end of the
message with the script. We parse this tag incrementally as it streams
in so that we can indicate that we are generating a script before we see
the closing `</eval>` tag. Later, this will help us interpret the script
as it arrives also.

- Threads now hold a `ScriptSession` entity which manages the state of
all scripts (from parsing to exited) in a centralized way, and will
later collect all script operations so they can be displayed in the UI.

- `script_tool` has been renamed to `assistant_scripting` 

- Script source now opens in a regular read-only buffer  

Note: We still need to handle persistence properly

Release Notes:

- N/A

---------

Co-authored-by: Marshall Bowers <[email protected]>

Author: agu-z

Cargo.lock

@@ -118,7 +118,7 @@ members = [
     "crates/rope",
     "crates/rpc",
     "crates/schema_generator",
-    "crates/scripting_tool",
+    "crates/assistant_scripting",
     "crates/search",
     "crates/semantic_index",
     "crates/semantic_version",
@@ -318,7 +318,7 @@ reqwest_client = { path = "crates/reqwest_client" }
 rich_text = { path = "crates/rich_text" }
 rope = { path = "crates/rope" }
 rpc = { path = "crates/rpc" }
-scripting_tool = { path = "crates/scripting_tool" }
+assistant_scripting = { path = "crates/assistant_scripting" }
 search = { path = "crates/search" }
 semantic_index = { path = "crates/semantic_index" }
 semantic_version = { path = "crates/semantic_version" }

Cargo.toml

@@ -63,6 +63,7 @@ serde.workspace = true
 serde_json.workspace = true
 settings.workspace = true
 smol.workspace = true
+assistant_scripting.workspace = true
 streaming_diff.workspace = true
 telemetry_events.workspace = true
 terminal.workspace = true

crates/assistant2/Cargo.toml

@@ -1,11 +1,12 @@
 use std::sync::Arc;
 
-use collections::HashMap;
+use assistant_scripting::{ScriptId, ScriptState};
+use collections::{HashMap, HashSet};
 use editor::{Editor, MultiBuffer};
 use gpui::{
     list, AbsoluteLength, AnyElement, App, ClickEvent, DefiniteLength, EdgesRefinement, Empty,
     Entity, Focusable, Length, ListAlignment, ListOffset, ListState, StyleRefinement, Subscription,
-    Task, TextStyleRefinement, UnderlineStyle,
+    Task, TextStyleRefinement, UnderlineStyle, WeakEntity,
 };
 use language::{Buffer, LanguageRegistry};
 use language_model::{LanguageModelRegistry, LanguageModelToolUseId, Role};
@@ -14,13 +15,15 @@ use settings::Settings as _;
 use theme::ThemeSettings;
 use ui::{prelude::*, Disclosure, KeyBinding};
 use util::ResultExt as _;
+use workspace::Workspace;
 
 use crate::thread::{MessageId, RequestKind, Thread, ThreadError, ThreadEvent};
 use crate::thread_store::ThreadStore;
 use crate::tool_use::{ToolUse, ToolUseStatus};
 use crate::ui::ContextPill;
 
 pub struct ActiveThread {
+    workspace: WeakEntity<Workspace>,
     language_registry: Arc<LanguageRegistry>,
     thread_store: Entity<ThreadStore>,
     thread: Entity<Thread>,
@@ -30,6 +33,7 @@ pub struct ActiveThread {
     rendered_messages_by_id: HashMap<MessageId, Entity<Markdown>>,
     editing_message: Option<(MessageId, EditMessageState)>,
     expanded_tool_uses: HashMap<LanguageModelToolUseId, bool>,
+    expanded_scripts: HashSet<ScriptId>,
     last_error: Option<ThreadError>,
     _subscriptions: Vec<Subscription>,
 }
@@ -40,6 +44,7 @@ struct EditMessageState {
 
 impl ActiveThread {
     pub fn new(
+        workspace: WeakEntity<Workspace>,
         thread: Entity<Thread>,
         thread_store: Entity<ThreadStore>,
         language_registry: Arc<LanguageRegistry>,
@@ -52,13 +57,15 @@ impl ActiveThread {
         ];
 
         let mut this = Self {
+            workspace,
             language_registry,
             thread_store,
             thread: thread.clone(),
             save_thread_task: None,
             messages: Vec::new(),
             rendered_messages_by_id: HashMap::default(),
             expanded_tool_uses: HashMap::default(),
+            expanded_scripts: HashSet::default(),
             list_state: ListState::new(0, ListAlignment::Bottom, px(1024.), {
                 let this = cx.entity().downgrade();
                 move |ix, window: &mut Window, cx: &mut App| {
@@ -241,7 +248,7 @@ impl ActiveThread {
 
     fn handle_thread_event(
         &mut self,
-        _: &Entity<Thread>,
+        _thread: &Entity<Thread>,
         event: &ThreadEvent,
         window: &mut Window,
         cx: &mut Context<Self>,
@@ -306,6 +313,14 @@ impl ActiveThread {
                     }
                 }
             }
+            ThreadEvent::ScriptFinished => {
+                let model_registry = LanguageModelRegistry::read_global(cx);
+                if let Some(model) = model_registry.active_model() {
+                    self.thread.update(cx, |thread, cx| {
+                        thread.send_to_model(model, RequestKind::Chat, false, cx);
+                    });
+                }
+            }
         }
     }
 
@@ -445,12 +460,16 @@ impl ActiveThread {
             return Empty.into_any();
         };
 
-        let context = self.thread.read(cx).context_for_message(message_id);
-        let tool_uses = self.thread.read(cx).tool_uses_for_message(message_id);
-        let colors = cx.theme().colors();
+        let thread = self.thread.read(cx);
+
+        let context = thread.context_for_message(message_id);
+        let tool_uses = thread.tool_uses_for_message(message_id);
 
         // Don't render user messages that are just there for returning tool results.
-        if message.role == Role::User && self.thread.read(cx).message_has_tool_results(message_id) {
+        if message.role == Role::User
+            && (thread.message_has_tool_results(message_id)
+                || thread.message_has_script_output(message_id))
+        {
             return Empty.into_any();
         }
 
@@ -463,6 +482,8 @@ impl ActiveThread {
             .filter(|(id, _)| *id == message_id)
             .map(|(_, state)| state.editor.clone());
 
+        let colors = cx.theme().colors();
+
         let message_content = v_flex()
             .child(
                 if let Some(edit_message_editor) = edit_message_editor.clone() {
@@ -597,6 +618,7 @@ impl ActiveThread {
             Role::Assistant => div()
                 .id(("message-container", ix))
                 .child(message_content)
+                .children(self.render_script(message_id, cx))
                 .map(|parent| {
                     if tool_uses.is_empty() {
                         return parent;
@@ -716,6 +738,139 @@ impl ActiveThread {
                 }),
         )
     }
+
+    fn render_script(&self, message_id: MessageId, cx: &mut Context<Self>) -> Option<AnyElement> {
+        let script = self.thread.read(cx).script_for_message(message_id, cx)?;
+
+        let is_open = self.expanded_scripts.contains(&script.id);
+        let colors = cx.theme().colors();
+
+        let element = div().px_2p5().child(
+            v_flex()
+                .gap_1()
+                .rounded_lg()
+                .border_1()
+                .border_color(colors.border)
+                .child(
+                    h_flex()
+                        .justify_between()
+                        .py_0p5()
+                        .pl_1()
+                        .pr_2()
+                        .bg(colors.editor_foreground.opacity(0.02))
+                        .when(is_open, |element| element.border_b_1().rounded_t(px(6.)))
+                        .when(!is_open, |element| element.rounded_md())
+                        .border_color(colors.border)
+                        .child(
+                            h_flex()
+                                .gap_1()
+                                .child(Disclosure::new("script-disclosure", is_open).on_click(
+                                    cx.listener({
+                                        let script_id = script.id;
+                                        move |this, _event, _window, _cx| {
+                                            if this.expanded_scripts.contains(&script_id) {
+                                                this.expanded_scripts.remove(&script_id);
+                                            } else {
+                                                this.expanded_scripts.insert(script_id);
+                                            }
+                                        }
+                                    }),
+                                ))
+                                // TODO: Generate script description
+                                .child(Label::new("Script")),
+                        )
+                        .child(
+                            h_flex()
+                                .gap_1()
+                                .child(
+                                    Label::new(match script.state {
+                                        ScriptState::Generating => "Generating",
+                                        ScriptState::Running { .. } => "Running",
+                                        ScriptState::Succeeded { .. } => "Finished",
+                                        ScriptState::Failed { .. } => "Error",
+                                    })
+                                    .size(LabelSize::XSmall)
+                                    .buffer_font(cx),
+                                )
+                                .child(
+                                    IconButton::new("view-source", IconName::Eye)
+                                        .icon_color(Color::Muted)
+                                        .disabled(matches!(script.state, ScriptState::Generating))
+                                        .on_click(cx.listener({
+                                            let source = script.source.clone();
+                                            move |this, _event, window, cx| {
+                                                this.open_script_source(source.clone(), window, cx);
+                                            }
+                                        })),
+                                ),
+                        ),
+                )
+                .when(is_open, |parent| {
+                    let stdout = script.stdout_snapshot();
+                    let error = script.error();
+
+                    parent.child(
+                        v_flex()
+                            .p_2()
+                            .bg(colors.editor_background)
+                            .gap_2()
+                            .child(if stdout.is_empty() && error.is_none() {
+                                Label::new("No output yet")
+                                    .size(LabelSize::Small)
+                                    .color(Color::Muted)
+                            } else {
+                                Label::new(stdout).size(LabelSize::Small).buffer_font(cx)
+                            })
+                            .children(script.error().map(|err| {
+                                Label::new(err.to_string())
+                                    .size(LabelSize::Small)
+                                    .color(Color::Error)
+                            })),
+                    )
+                }),
+        );
+
+        Some(element.into_any())
+    }
+
+    fn open_script_source(
+        &mut self,
+        source: SharedString,
+        window: &mut Window,
+        cx: &mut Context<'_, ActiveThread>,
+    ) {
+        let language_registry = self.language_registry.clone();
+        let workspace = self.workspace.clone();
+        let source = source.clone();
+
+        cx.spawn_in(window, |_, mut cx| async move {
+            let lua = language_registry.language_for_name("Lua").await.log_err();
+
+            workspace.update_in(&mut cx, |workspace, window, cx| {
+                let project = workspace.project().clone();
+
+                let buffer = project.update(cx, |project, cx| {
+                    project.create_local_buffer(&source.trim(), lua, cx)
+                });
+
+                let buffer = cx.new(|cx| {
+                    MultiBuffer::singleton(buffer, cx)
+                        // TODO: Generate script description
+                        .with_title("Assistant script".into())
+                });
+
+                let editor = cx.new(|cx| {
+                    let mut editor =
+                        Editor::for_multibuffer(buffer, Some(project), true, window, cx);
+                    editor.set_read_only(true);
+                    editor
+                });
+
+                workspace.add_item_to_active_pane(Box::new(editor), None, true, window, cx);
+            })
+        })
+        .detach_and_log_err(cx);
+    }
 }
 
 impl Render for ActiveThread {