1462 lines · 48.6 KB · Rust
RawDownload
use crate::{    CacheStatus, InvokedSlashCommandId, MessageCacheMetadata, MessageId, MessageStatus, TextThread,    TextThreadEvent, TextThreadId, TextThreadOperation, TextThreadSummary,};use anyhow::Result;use assistant_slash_command::{    ArgumentCompletion, SlashCommand, SlashCommandContent, SlashCommandEvent, SlashCommandOutput,    SlashCommandOutputSection, SlashCommandRegistry, SlashCommandResult, SlashCommandWorkingSet,};use assistant_slash_commands::FileSlashCommand;use collections::{HashMap, HashSet};use fs::FakeFs;use futures::{    channel::mpsc,    stream::{self, StreamExt},};use gpui::{App, Entity, SharedString, Task, TestAppContext, WeakEntity, prelude::*};use language::{Buffer, BufferSnapshot, LanguageRegistry, LspAdapterDelegate};use language_model::{    ConfiguredModel, LanguageModelCacheConfiguration, LanguageModelRegistry, Role,    fake_provider::{FakeLanguageModel, FakeLanguageModelProvider},};use parking_lot::Mutex;use pretty_assertions::assert_eq;use prompt_store::PromptBuilder;use rand::prelude::*;use serde_json::json;use settings::SettingsStore;use std::{    cell::RefCell,    env,    ops::Range,    path::Path,    rc::Rc,    sync::{Arc, atomic::AtomicBool},};use text::{ReplicaId, ToOffset, network::Network};use ui::{IconName, Window};use unindent::Unindent;use util::RandomCharIter;use workspace::Workspace; #[gpui::test]fn test_inserting_and_removing_messages(cx: &mut App) {    init_test(cx);     let registry = Arc::new(LanguageRegistry::test(cx.background_executor().clone()));    let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());    let text_thread = cx.new(|cx| {        TextThread::local(            registry,            None,            None,            prompt_builder.clone(),            Arc::new(SlashCommandWorkingSet::default()),            cx,        )    });    let buffer = text_thread.read(cx).buffer().clone();     let message_1 = text_thread.read(cx).message_anchors[0].clone();    assert_eq!(        messages(&text_thread, cx),        vec![(message_1.id, Role::User, 0..0)]    );     let message_2 = text_thread.update(cx, |context, cx| {        context            .insert_message_after(message_1.id, Role::Assistant, MessageStatus::Done, cx)            .unwrap()    });    assert_eq!(        messages(&text_thread, 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(&text_thread, cx),        vec![            (message_1.id, Role::User, 0..2),            (message_2.id, Role::Assistant, 2..3)        ]    );     let message_3 = text_thread.update(cx, |context, cx| {        context            .insert_message_after(message_2.id, Role::User, MessageStatus::Done, cx)            .unwrap()    });    assert_eq!(        messages(&text_thread, 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 = text_thread.update(cx, |context, cx| {        context            .insert_message_after(message_2.id, Role::User, MessageStatus::Done, cx)            .unwrap()    });    assert_eq!(        messages(&text_thread, 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(&text_thread, 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(&text_thread, 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(&text_thread, 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(&text_thread, 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 = text_thread.update(cx, |context, cx| {        context            .insert_message_after(message_1.id, Role::System, MessageStatus::Done, cx)            .unwrap()    });    assert_eq!(        messages(&text_thread, 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 App) {    init_test(cx);     let registry = Arc::new(LanguageRegistry::test(cx.background_executor().clone()));     let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());    let text_thread = cx.new(|cx| {        TextThread::local(            registry.clone(),            None,            None,            prompt_builder.clone(),            Arc::new(SlashCommandWorkingSet::default()),            cx,        )    });    let buffer = text_thread.read(cx).buffer().clone();     let message_1 = text_thread.read(cx).message_anchors[0].clone();    assert_eq!(        messages(&text_thread, 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) =        text_thread.update(cx, |text_thread, cx| text_thread.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(&text_thread, cx),        vec![            (message_1.id, Role::User, 0..4),            (message_2.id, Role::User, 4..16),        ]    );     let (_, message_3) =        text_thread.update(cx, |text_thread, cx| text_thread.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(&text_thread, 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) =        text_thread.update(cx, |text_thread, cx| text_thread.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(&text_thread, 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) =        text_thread.update(cx, |text_thread, cx| text_thread.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(&text_thread, 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) =        text_thread.update(cx, |text_thread, cx| text_thread.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(&text_thread, 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 App) {    init_test(cx);     let registry = Arc::new(LanguageRegistry::test(cx.background_executor().clone()));    let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());    let text_thread = cx.new(|cx| {        TextThread::local(            registry,            None,            None,            prompt_builder.clone(),            Arc::new(SlashCommandWorkingSet::default()),            cx,        )    });    let buffer = text_thread.read(cx).buffer().clone();     let message_1 = text_thread.read(cx).message_anchors[0].clone();    assert_eq!(        messages(&text_thread, cx),        vec![(message_1.id, Role::User, 0..0)]    );     buffer.update(cx, |buffer, cx| buffer.edit([(0..0, "aaa")], None, cx));    let message_2 = text_thread        .update(cx, |text_thread, cx| {            text_thread.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 = text_thread        .update(cx, |text_thread, cx| {            text_thread.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(&text_thread, 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(&text_thread, &[0, 4, 9], cx),        [message_1.id, message_2.id, message_3.id]    );    assert_eq!(        message_ids_for_offsets(&text_thread, &[0, 1, 11], cx),        [message_1.id, message_3.id]    );     let message_4 = text_thread        .update(cx, |text_thread, cx| {            text_thread.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(&text_thread, 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(&text_thread, &[0, 4, 8, 12], cx),        [message_1.id, message_2.id, message_3.id, message_4.id]    );     fn message_ids_for_offsets(        context: &Entity<TextThread>,        offsets: &[usize],        cx: &App,    ) -> 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) {    cx.update(init_test);     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 text_thread = cx.new(|cx| {        TextThread::local(            registry.clone(),            None,            None,            prompt_builder.clone(),            Arc::new(SlashCommandWorkingSet::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()));    text_thread.update(cx, |_, cx| {        cx.subscribe(&text_thread, {            let context_ranges = context_ranges.clone();            move |text_thread, _, event, _| {                let mut context_ranges = context_ranges.borrow_mut();                match event {                    TextThreadEvent::InvokedSlashCommandChanged { command_id } => {                        let command = text_thread.invoked_slash_command(command_id).unwrap();                        context_ranges                            .command_outputs                            .insert(*command_id, command.range.clone());                    }                    TextThreadEvent::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());                        }                    }                    TextThreadEvent::SlashCommandOutputSectionAdded { section } => {                        context_ranges.output_sections.insert(section.range.clone());                    }                    _ => {}                }            }        })        .detach();    });     let buffer = text_thread.read_with(cx, |text_thread, _| text_thread.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")],            None,            cx,        );    });    assert_text_and_context_ranges(        &buffer,        &context_ranges,        &"        /unknown src/main.rs"            .unindent(),        cx,    );

Lines 1–500 of 1462.