787 lines · 31.4 KB · Rust
RawDownload
use assistant_tool::{ToolRegistry, ToolWorkingSet};use collections::HashMap;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, WeakEntity,};use language::{Buffer, Language, LanguageRegistry};use language_model::{LanguageModelRegistry, LanguageModelToolUseId, Role};use markdown::{Markdown, MarkdownStyle};use settings::Settings as _;use std::sync::Arc;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>,    tools: Arc<ToolWorkingSet>,    thread_store: Entity<ThreadStore>,    thread: Entity<Thread>,    save_thread_task: Option<Task<()>>,    messages: Vec<MessageId>,    list_state: ListState,    rendered_messages_by_id: HashMap<MessageId, Entity<Markdown>>,    editing_message: Option<(MessageId, EditMessageState)>,    expanded_tool_uses: HashMap<LanguageModelToolUseId, bool>,    last_error: Option<ThreadError>,    lua_language: Option<Arc<Language>>, // Used for syntax highlighting in the Lua script tool    _subscriptions: Vec<Subscription>,} struct EditMessageState {    editor: Entity<Editor>,} impl ActiveThread {    pub fn new(        thread: Entity<Thread>,        thread_store: Entity<ThreadStore>,        workspace: WeakEntity<Workspace>,        language_registry: Arc<LanguageRegistry>,        tools: Arc<ToolWorkingSet>,        window: &mut Window,        cx: &mut Context<Self>,    ) -> Self {        let subscriptions = vec![            cx.observe(&thread, |_, _, cx| cx.notify()),            cx.subscribe_in(&thread, window, Self::handle_thread_event),        ];         let mut this = Self {            workspace,            language_registry,            tools,            thread_store,            thread: thread.clone(),            save_thread_task: None,            messages: Vec::new(),            rendered_messages_by_id: HashMap::default(),            expanded_tool_uses: HashMap::default(),            list_state: ListState::new(0, ListAlignment::Bottom, px(1024.), {                let this = cx.entity().downgrade();                move |ix, window: &mut Window, cx: &mut App| {                    this.update(cx, |this, cx| this.render_message(ix, window, cx))                        .unwrap()                }            }),            editing_message: None,            last_error: None,            lua_language: None,            _subscriptions: subscriptions,        };         // Initialize the Lua language in the background, for syntax highlighting.        let language_registry = this.language_registry.clone();        cx.spawn(|this, mut cx| async move {            if let Ok(lua_language) = language_registry.language_for_name("Lua").await {                this.update(&mut cx, |this, _| {                    this.lua_language = Some(lua_language);                })?;            }             Ok::<_, anyhow::Error>(())        })        .detach();         for message in thread.read(cx).messages().cloned().collect::<Vec<_>>() {            this.push_message(&message.id, message.text.clone(), window, cx);        }         this    }     pub fn thread(&self) -> &Entity<Thread> {        &self.thread    }     pub fn is_empty(&self) -> bool {        self.messages.is_empty()    }     pub fn summary(&self, cx: &App) -> Option<SharedString> {        self.thread.read(cx).summary()    }     pub fn summary_or_default(&self, cx: &App) -> SharedString {        self.thread.read(cx).summary_or_default()    }     pub fn cancel_last_completion(&mut self, cx: &mut App) -> bool {        self.last_error.take();        self.thread            .update(cx, |thread, _cx| thread.cancel_last_completion())    }     pub fn last_error(&self) -> Option<ThreadError> {        self.last_error.clone()    }     pub fn clear_last_error(&mut self) {        self.last_error.take();    }     fn push_message(        &mut self,        id: &MessageId,        text: String,        window: &mut Window,        cx: &mut Context<Self>,    ) {        let old_len = self.messages.len();        self.messages.push(*id);        self.list_state.splice(old_len..old_len, 1);         let markdown = self.render_markdown(text.into(), window, cx);        self.rendered_messages_by_id.insert(*id, markdown);        self.list_state.scroll_to(ListOffset {            item_ix: old_len,            offset_in_item: Pixels(0.0),        });    }     fn edited_message(        &mut self,        id: &MessageId,        text: String,        window: &mut Window,        cx: &mut Context<Self>,    ) {        let Some(index) = self.messages.iter().position(|message_id| message_id == id) else {            return;        };        self.list_state.splice(index..index + 1, 1);        let markdown = self.render_markdown(text.into(), window, cx);        self.rendered_messages_by_id.insert(*id, markdown);    }     fn deleted_message(&mut self, id: &MessageId) {        let Some(index) = self.messages.iter().position(|message_id| message_id == id) else {            return;        };        self.messages.remove(index);        self.list_state.splice(index..index + 1, 0);        self.rendered_messages_by_id.remove(id);    }     fn render_markdown(        &self,        text: SharedString,        window: &Window,        cx: &mut Context<Self>,    ) -> Entity<Markdown> {        let theme_settings = ThemeSettings::get_global(cx);        let colors = cx.theme().colors();        let ui_font_size = TextSize::Default.rems(cx);        let buffer_font_size = TextSize::Small.rems(cx);        let mut text_style = window.text_style();         text_style.refine(&TextStyleRefinement {            font_family: Some(theme_settings.ui_font.family.clone()),            font_size: Some(ui_font_size.into()),            color: Some(cx.theme().colors().text),            ..Default::default()        });         let markdown_style = MarkdownStyle {            base_text_style: text_style,            syntax: cx.theme().syntax().clone(),            selection_background_color: cx.theme().players().local().selection,            code_block_overflow_x_scroll: true,            table_overflow_x_scroll: true,            code_block: StyleRefinement {                margin: EdgesRefinement {                    top: Some(Length::Definite(rems(0.).into())),                    left: Some(Length::Definite(rems(0.).into())),                    right: Some(Length::Definite(rems(0.).into())),                    bottom: Some(Length::Definite(rems(0.5).into())),                },                padding: EdgesRefinement {                    top: Some(DefiniteLength::Absolute(AbsoluteLength::Pixels(Pixels(8.)))),                    left: Some(DefiniteLength::Absolute(AbsoluteLength::Pixels(Pixels(8.)))),                    right: Some(DefiniteLength::Absolute(AbsoluteLength::Pixels(Pixels(8.)))),                    bottom: Some(DefiniteLength::Absolute(AbsoluteLength::Pixels(Pixels(8.)))),                },                background: Some(colors.editor_background.into()),                border_color: Some(colors.border_variant),                border_widths: EdgesRefinement {                    top: Some(AbsoluteLength::Pixels(Pixels(1.))),                    left: Some(AbsoluteLength::Pixels(Pixels(1.))),                    right: Some(AbsoluteLength::Pixels(Pixels(1.))),                    bottom: Some(AbsoluteLength::Pixels(Pixels(1.))),                },                text: Some(TextStyleRefinement {                    font_family: Some(theme_settings.buffer_font.family.clone()),                    font_size: Some(buffer_font_size.into()),                    ..Default::default()                }),                ..Default::default()            },            inline_code: TextStyleRefinement {                font_family: Some(theme_settings.buffer_font.family.clone()),                font_size: Some(buffer_font_size.into()),                background_color: Some(colors.editor_foreground.opacity(0.1)),                ..Default::default()            },            link: TextStyleRefinement {                background_color: Some(colors.editor_foreground.opacity(0.025)),                underline: Some(UnderlineStyle {                    color: Some(colors.text_accent.opacity(0.5)),                    thickness: px(1.),                    ..Default::default()                }),                ..Default::default()            },            ..Default::default()        };         cx.new(|cx| {            Markdown::new(                text,                markdown_style,                Some(self.language_registry.clone()),                None,                cx,            )        })    }     fn handle_thread_event(        &mut self,        _: &Entity<Thread>,        event: &ThreadEvent,        window: &mut Window,        cx: &mut Context<Self>,    ) {        match event {            ThreadEvent::ShowError(error) => {                self.last_error = Some(error.clone());            }            ThreadEvent::StreamedCompletion | ThreadEvent::SummaryChanged => {                self.save_thread(cx);            }            ThreadEvent::StreamedAssistantText(message_id, text) => {                if let Some(markdown) = self.rendered_messages_by_id.get_mut(&message_id) {                    markdown.update(cx, |markdown, cx| {                        markdown.append(text, cx);                    });                }            }            ThreadEvent::MessageAdded(message_id) => {                if let Some(message_text) = self                    .thread                    .read(cx)                    .message(*message_id)                    .map(|message| message.text.clone())                {                    self.push_message(message_id, message_text, window, cx);                }                 self.save_thread(cx);                cx.notify();            }            ThreadEvent::MessageEdited(message_id) => {                if let Some(message_text) = self                    .thread                    .read(cx)                    .message(*message_id)                    .map(|message| message.text.clone())                {                    self.edited_message(message_id, message_text, window, cx);                }                 self.save_thread(cx);                cx.notify();            }            ThreadEvent::MessageDeleted(message_id) => {                self.deleted_message(message_id);                self.save_thread(cx);                cx.notify();            }            ThreadEvent::UsePendingTools => {                let thread = self.thread.read(cx);                let thread_id = thread.id().0.clone();                let pending_tool_uses = thread                    .pending_tool_uses()                    .into_iter()                    .filter(|tool_use| tool_use.status.is_idle())                    .cloned()                    .collect::<Vec<_>>();                 for tool_use in pending_tool_uses {                    if let Some(tool) = self.tools.tool(&tool_use.name, cx) {                        let task = tool.run(                            tool_use.input,                            thread_id.clone(),                            self.workspace.clone(),                            window,                            cx,                        );                         self.thread.update(cx, |thread, cx| {                            thread.insert_tool_output(tool_use.id.clone(), task, cx);                        });                    }                }            }            ThreadEvent::ToolFinished { .. } => {                let all_tools_finished = self                    .thread                    .read(cx)                    .pending_tool_uses()                    .into_iter()                    .all(|tool_use| tool_use.status.is_error());                if all_tools_finished {                    let model_registry = LanguageModelRegistry::read_global(cx);                    if let Some(model) = model_registry.active_model() {                        self.thread.update(cx, |thread, cx| {                            // Insert a user message to contain the tool results.                            thread.insert_user_message(                                // TODO: Sending up a user message without any content results in the model sending back                                // responses that also don't have any content. We currently don't handle this case well,                                // so for now we provide some text to keep the model on track.                                "Here are the tool results.",                                Vec::new(),                                cx,                            );                            thread.send_to_model(model, RequestKind::Chat, true, cx);                        });                    }                }            }        }    }     /// Spawns a task to save the active thread.    ///    /// Only one task to save the thread will be in flight at a time.    fn save_thread(&mut self, cx: &mut Context<Self>) {        let thread = self.thread.clone();        self.save_thread_task = Some(cx.spawn(|this, mut cx| async move {            let task = this                .update(&mut cx, |this, cx| {                    this.thread_store                        .update(cx, |thread_store, cx| thread_store.save_thread(&thread, cx))                })                .ok();             if let Some(task) = task {                task.await.log_err();            }        }));    }     fn start_editing_message(        &mut self,        message_id: MessageId,        message_text: String,        window: &mut Window,        cx: &mut Context<Self>,    ) {        let buffer = cx.new(|cx| {            MultiBuffer::singleton(cx.new(|cx| Buffer::local(message_text.clone(), cx)), cx)        });        let editor = cx.new(|cx| {            let mut editor = Editor::new(                editor::EditorMode::AutoHeight { max_lines: 8 },                buffer,                None,                false,                window,                cx,            );            editor.focus_handle(cx).focus(window);            editor.move_to_end(&editor::actions::MoveToEnd, window, cx);            editor        });        self.editing_message = Some((            message_id,            EditMessageState {                editor: editor.clone(),            },        ));        cx.notify();    }     fn cancel_editing_message(&mut self, _: &menu::Cancel, _: &mut Window, cx: &mut Context<Self>) {        self.editing_message.take();        cx.notify();    }     fn confirm_editing_message(        &mut self,        _: &menu::Confirm,        _: &mut Window,        cx: &mut Context<Self>,    ) {        let Some((message_id, state)) = self.editing_message.take() else {            return;        };        let edited_text = state.editor.read(cx).text(cx);        self.thread.update(cx, |thread, cx| {            thread.edit_message(message_id, Role::User, edited_text, cx);            for message_id in self.messages_after(message_id) {                thread.delete_message(*message_id, cx);            }        });         let provider = LanguageModelRegistry::read_global(cx).active_provider();        if provider            .as_ref()            .map_or(false, |provider| provider.must_accept_terms(cx))        {            cx.notify();            return;        }        let model_registry = LanguageModelRegistry::read_global(cx);        let Some(model) = model_registry.active_model() else {            return;        };         self.thread.update(cx, |thread, cx| {            thread.send_to_model(model, RequestKind::Chat, false, cx)        });        cx.notify();    }     fn last_user_message(&self, cx: &Context<Self>) -> Option<MessageId> {        self.messages            .iter()            .rev()            .find(|message_id| {                self.thread                    .read(cx)                    .message(**message_id)                    .map_or(false, |message| message.role == Role::User)            })            .cloned()    }     fn messages_after(&self, message_id: MessageId) -> &[MessageId] {        self.messages            .iter()            .position(|id| *id == message_id)            .map(|index| &self.messages[index + 1..])            .unwrap_or(&[])    }     fn handle_cancel_click(&mut self, _: &ClickEvent, window: &mut Window, cx: &mut Context<Self>) {        self.cancel_editing_message(&menu::Cancel, window, cx);    }     fn handle_regenerate_click(        &mut self,        _: &ClickEvent,        window: &mut Window,        cx: &mut Context<Self>,    ) {        self.confirm_editing_message(&menu::Confirm, window, cx);    }     fn render_message(&self, ix: usize, window: &mut Window, cx: &mut Context<Self>) -> AnyElement {        let message_id = self.messages[ix];        let Some(message) = self.thread.read(cx).message(message_id) else {            return Empty.into_any();        };         let Some(markdown) = self.rendered_messages_by_id.get(&message_id) else {            return Empty.into_any();        };         let context = self.thread.read(cx).context_for_message(message_id);

Lines 1–500 of 787.