1671 lines · 52.4 KB · Rust
RawDownload
use super::{AssistantEdit, MessageCacheMetadata};use crate::{    assistant_panel, AssistantEditKind, CacheStatus, Context, ContextEvent, ContextId,    ContextOperation, InvokedSlashCommandId, MessageId, MessageStatus,};use anyhow::Result;use assistant_slash_command::{    ArgumentCompletion, SlashCommand, SlashCommandContent, SlashCommandEvent, SlashCommandOutput,    SlashCommandOutputSection, SlashCommandRegistry, SlashCommandResult, SlashCommandWorkingSet,};use assistant_slash_commands::FileSlashCommand;use assistant_tool::ToolWorkingSet;use collections::{HashMap, HashSet};use fs::FakeFs;use futures::{    channel::mpsc,    stream::{self, StreamExt},};use gpui::{prelude::*, AppContext, Model, SharedString, Task, TestAppContext, WeakView};use language::{Buffer, BufferSnapshot, LanguageRegistry, LspAdapterDelegate};use language_model::{LanguageModelCacheConfiguration, LanguageModelRegistry, Role};use parking_lot::Mutex;use pretty_assertions::assert_eq;use project::Project;use prompt_library::PromptBuilder;use rand::prelude::*;use serde_json::json;use settings::SettingsStore;use std::{    cell::RefCell,    env,    ops::Range,    path::Path,    rc::Rc,    sync::{atomic::AtomicBool, Arc},};use text::{network::Network, OffsetRangeExt as _, ReplicaId, ToOffset};use ui::{IconName, WindowContext};use unindent::Unindent;use util::{    test::{generate_marked_text, marked_text_ranges},    RandomCharIter,};use workspace::Workspace; #[gpui::test]fn test_inserting_and_removing_messages(cx: &mut AppContext) {    let settings_store = SettingsStore::test(cx);    LanguageModelRegistry::test(cx);    cx.set_global(settings_store);    assistant_panel::init(cx);    let registry = Arc::new(LanguageRegistry::test(cx.background_executor().clone()));    let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());    let context = cx.new_model(|cx| {        Context::local(            registry,            None,            None,            prompt_builder.clone(),            Arc::new(SlashCommandWorkingSet::default()),            Arc::new(ToolWorkingSet::default()),            cx,        )    });    let buffer = context.read(cx).buffer.clone();     let message_1 = context.read(cx).message_anchors[0].clone();    assert_eq!(        messages(&context, cx),        vec![(message_1.id, Role::User, 0..0)]    );     let message_2 = context.update(cx, |context, cx| {        context            .insert_message_after(message_1.id, Role::Assistant, MessageStatus::Done, cx)            .unwrap()    });    assert_eq!(        messages(&context, cx),        vec![            (message_1.id, Role::User, 0..1),            (message_2.id, Role::Assistant, 1..1)        ]    );     buffer.update(cx, |buffer, cx| {        buffer.edit([(0..0, "1"), (1..1, "2")], None, cx)    });    assert_eq!(        messages(&context, cx),        vec![            (message_1.id, Role::User, 0..2),            (message_2.id, Role::Assistant, 2..3)        ]    );     let message_3 = context.update(cx, |context, cx| {        context            .insert_message_after(message_2.id, Role::User, MessageStatus::Done, cx)            .unwrap()    });    assert_eq!(        messages(&context, cx),        vec![            (message_1.id, Role::User, 0..2),            (message_2.id, Role::Assistant, 2..4),            (message_3.id, Role::User, 4..4)        ]    );     let message_4 = context.update(cx, |context, cx| {        context            .insert_message_after(message_2.id, Role::User, MessageStatus::Done, cx)            .unwrap()    });    assert_eq!(        messages(&context, cx),        vec![            (message_1.id, Role::User, 0..2),            (message_2.id, Role::Assistant, 2..4),            (message_4.id, Role::User, 4..5),            (message_3.id, Role::User, 5..5),        ]    );     buffer.update(cx, |buffer, cx| {        buffer.edit([(4..4, "C"), (5..5, "D")], None, cx)    });    assert_eq!(        messages(&context, cx),        vec![            (message_1.id, Role::User, 0..2),            (message_2.id, Role::Assistant, 2..4),            (message_4.id, Role::User, 4..6),            (message_3.id, Role::User, 6..7),        ]    );     // Deleting across message boundaries merges the messages.    buffer.update(cx, |buffer, cx| buffer.edit([(1..4, "")], None, cx));    assert_eq!(        messages(&context, cx),        vec![            (message_1.id, Role::User, 0..3),            (message_3.id, Role::User, 3..4),        ]    );     // Undoing the deletion should also undo the merge.    buffer.update(cx, |buffer, cx| buffer.undo(cx));    assert_eq!(        messages(&context, cx),        vec![            (message_1.id, Role::User, 0..2),            (message_2.id, Role::Assistant, 2..4),            (message_4.id, Role::User, 4..6),            (message_3.id, Role::User, 6..7),        ]    );     // Redoing the deletion should also redo the merge.    buffer.update(cx, |buffer, cx| buffer.redo(cx));    assert_eq!(        messages(&context, cx),        vec![            (message_1.id, Role::User, 0..3),            (message_3.id, Role::User, 3..4),        ]    );     // Ensure we can still insert after a merged message.    let message_5 = context.update(cx, |context, cx| {        context            .insert_message_after(message_1.id, Role::System, MessageStatus::Done, cx)            .unwrap()    });    assert_eq!(        messages(&context, cx),        vec![            (message_1.id, Role::User, 0..3),            (message_5.id, Role::System, 3..4),            (message_3.id, Role::User, 4..5)        ]    );} #[gpui::test]fn test_message_splitting(cx: &mut AppContext) {    let settings_store = SettingsStore::test(cx);    cx.set_global(settings_store);    LanguageModelRegistry::test(cx);    assistant_panel::init(cx);    let registry = Arc::new(LanguageRegistry::test(cx.background_executor().clone()));     let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());    let context = cx.new_model(|cx| {        Context::local(            registry.clone(),            None,            None,            prompt_builder.clone(),            Arc::new(SlashCommandWorkingSet::default()),            Arc::new(ToolWorkingSet::default()),            cx,        )    });    let buffer = context.read(cx).buffer.clone();     let message_1 = context.read(cx).message_anchors[0].clone();    assert_eq!(        messages(&context, cx),        vec![(message_1.id, Role::User, 0..0)]    );     buffer.update(cx, |buffer, cx| {        buffer.edit([(0..0, "aaa\nbbb\nccc\nddd\n")], None, cx)    });     let (_, message_2) = context.update(cx, |context, cx| context.split_message(3..3, cx));    let message_2 = message_2.unwrap();     // We recycle newlines in the middle of a split message    assert_eq!(buffer.read(cx).text(), "aaa\nbbb\nccc\nddd\n");    assert_eq!(        messages(&context, cx),        vec![            (message_1.id, Role::User, 0..4),            (message_2.id, Role::User, 4..16),        ]    );     let (_, message_3) = context.update(cx, |context, cx| context.split_message(3..3, cx));    let message_3 = message_3.unwrap();     // We don't recycle newlines at the end of a split message    assert_eq!(buffer.read(cx).text(), "aaa\n\nbbb\nccc\nddd\n");    assert_eq!(        messages(&context, cx),        vec![            (message_1.id, Role::User, 0..4),            (message_3.id, Role::User, 4..5),            (message_2.id, Role::User, 5..17),        ]    );     let (_, message_4) = context.update(cx, |context, cx| context.split_message(9..9, cx));    let message_4 = message_4.unwrap();    assert_eq!(buffer.read(cx).text(), "aaa\n\nbbb\nccc\nddd\n");    assert_eq!(        messages(&context, cx),        vec![            (message_1.id, Role::User, 0..4),            (message_3.id, Role::User, 4..5),            (message_2.id, Role::User, 5..9),            (message_4.id, Role::User, 9..17),        ]    );     let (_, message_5) = context.update(cx, |context, cx| context.split_message(9..9, cx));    let message_5 = message_5.unwrap();    assert_eq!(buffer.read(cx).text(), "aaa\n\nbbb\n\nccc\nddd\n");    assert_eq!(        messages(&context, cx),        vec![            (message_1.id, Role::User, 0..4),            (message_3.id, Role::User, 4..5),            (message_2.id, Role::User, 5..9),            (message_4.id, Role::User, 9..10),            (message_5.id, Role::User, 10..18),        ]    );     let (message_6, message_7) =        context.update(cx, |context, cx| context.split_message(14..16, cx));    let message_6 = message_6.unwrap();    let message_7 = message_7.unwrap();    assert_eq!(buffer.read(cx).text(), "aaa\n\nbbb\n\nccc\ndd\nd\n");    assert_eq!(        messages(&context, cx),        vec![            (message_1.id, Role::User, 0..4),            (message_3.id, Role::User, 4..5),            (message_2.id, Role::User, 5..9),            (message_4.id, Role::User, 9..10),            (message_5.id, Role::User, 10..14),            (message_6.id, Role::User, 14..17),            (message_7.id, Role::User, 17..19),        ]    );} #[gpui::test]fn test_messages_for_offsets(cx: &mut AppContext) {    let settings_store = SettingsStore::test(cx);    LanguageModelRegistry::test(cx);    cx.set_global(settings_store);    assistant_panel::init(cx);    let registry = Arc::new(LanguageRegistry::test(cx.background_executor().clone()));    let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());    let context = cx.new_model(|cx| {        Context::local(            registry,            None,            None,            prompt_builder.clone(),            Arc::new(SlashCommandWorkingSet::default()),            Arc::new(ToolWorkingSet::default()),            cx,        )    });    let buffer = context.read(cx).buffer.clone();     let message_1 = context.read(cx).message_anchors[0].clone();    assert_eq!(        messages(&context, cx),        vec![(message_1.id, Role::User, 0..0)]    );     buffer.update(cx, |buffer, cx| buffer.edit([(0..0, "aaa")], None, cx));    let message_2 = context        .update(cx, |context, cx| {            context.insert_message_after(message_1.id, Role::User, MessageStatus::Done, cx)        })        .unwrap();    buffer.update(cx, |buffer, cx| buffer.edit([(4..4, "bbb")], None, cx));     let message_3 = context        .update(cx, |context, cx| {            context.insert_message_after(message_2.id, Role::User, MessageStatus::Done, cx)        })        .unwrap();    buffer.update(cx, |buffer, cx| buffer.edit([(8..8, "ccc")], None, cx));     assert_eq!(buffer.read(cx).text(), "aaa\nbbb\nccc");    assert_eq!(        messages(&context, cx),        vec![            (message_1.id, Role::User, 0..4),            (message_2.id, Role::User, 4..8),            (message_3.id, Role::User, 8..11)        ]    );     assert_eq!(        message_ids_for_offsets(&context, &[0, 4, 9], cx),        [message_1.id, message_2.id, message_3.id]    );    assert_eq!(        message_ids_for_offsets(&context, &[0, 1, 11], cx),        [message_1.id, message_3.id]    );     let message_4 = context        .update(cx, |context, cx| {            context.insert_message_after(message_3.id, Role::User, MessageStatus::Done, cx)        })        .unwrap();    assert_eq!(buffer.read(cx).text(), "aaa\nbbb\nccc\n");    assert_eq!(        messages(&context, cx),        vec![            (message_1.id, Role::User, 0..4),            (message_2.id, Role::User, 4..8),            (message_3.id, Role::User, 8..12),            (message_4.id, Role::User, 12..12)        ]    );    assert_eq!(        message_ids_for_offsets(&context, &[0, 4, 8, 12], cx),        [message_1.id, message_2.id, message_3.id, message_4.id]    );     fn message_ids_for_offsets(        context: &Model<Context>,        offsets: &[usize],        cx: &AppContext,    ) -> Vec<MessageId> {        context            .read(cx)            .messages_for_offsets(offsets.iter().copied(), cx)            .into_iter()            .map(|message| message.id)            .collect()    }} #[gpui::test]async fn test_slash_commands(cx: &mut TestAppContext) {    let settings_store = cx.update(SettingsStore::test);    cx.set_global(settings_store);    cx.update(LanguageModelRegistry::test);    cx.update(Project::init_settings);    cx.update(assistant_panel::init);    let fs = FakeFs::new(cx.background_executor.clone());     fs.insert_tree(        "/test",        json!({            "src": {                "lib.rs": "fn one() -> usize { 1 }",                "main.rs": "                    use crate::one;                    fn main() { one(); }                ".unindent(),            }        }),    )    .await;     let slash_command_registry = cx.update(SlashCommandRegistry::default_global);    slash_command_registry.register_command(FileSlashCommand, false);     let registry = Arc::new(LanguageRegistry::test(cx.executor()));    let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());    let context = cx.new_model(|cx| {        Context::local(            registry.clone(),            None,            None,            prompt_builder.clone(),            Arc::new(SlashCommandWorkingSet::default()),            Arc::new(ToolWorkingSet::default()),            cx,        )    });     #[derive(Default)]    struct ContextRanges {        parsed_commands: HashSet<Range<language::Anchor>>,        command_outputs: HashMap<InvokedSlashCommandId, Range<language::Anchor>>,        output_sections: HashSet<Range<language::Anchor>>,    }     let context_ranges = Rc::new(RefCell::new(ContextRanges::default()));    context.update(cx, |_, cx| {        cx.subscribe(&context, {            let context_ranges = context_ranges.clone();            move |context, _, event, _| {                let mut context_ranges = context_ranges.borrow_mut();                match event {                    ContextEvent::InvokedSlashCommandChanged { command_id } => {                        let command = context.invoked_slash_command(command_id).unwrap();                        context_ranges                            .command_outputs                            .insert(*command_id, command.range.clone());                    }                    ContextEvent::ParsedSlashCommandsUpdated { removed, updated } => {                        for range in removed {                            context_ranges.parsed_commands.remove(range);                        }                        for command in updated {                            context_ranges                                .parsed_commands                                .insert(command.source_range.clone());                        }                    }                    ContextEvent::SlashCommandOutputSectionAdded { section } => {                        context_ranges.output_sections.insert(section.range.clone());                    }                    _ => {}                }            }        })        .detach();    });     let buffer = context.read_with(cx, |context, _| context.buffer.clone());     // Insert a slash command    buffer.update(cx, |buffer, cx| {        buffer.edit([(0..0, "/file src/lib.rs")], None, cx);    });    assert_text_and_context_ranges(        &buffer,        &context_ranges,        &"        «/file src/lib.rs»"            .unindent(),        cx,    );     // Edit the argument of the slash command.    buffer.update(cx, |buffer, cx| {        let edit_offset = buffer.text().find("lib.rs").unwrap();        buffer.edit([(edit_offset..edit_offset + "lib".len(), "main")], None, cx);    });    assert_text_and_context_ranges(        &buffer,        &context_ranges,        &"        «/file src/main.rs»"            .unindent(),        cx,    );     // Edit the name of the slash command, using one that doesn't exist.    buffer.update(cx, |buffer, cx| {        let edit_offset = buffer.text().find("/file").unwrap();        buffer.edit(            [(edit_offset..edit_offset + "/file".len(), "/unknown")],

Lines 1–500 of 1671.