1438 lines · 44.6 KB · Rust
RawDownload
use super::{MessageCacheMetadata, WorkflowStepEdit};use crate::{    assistant_panel, prompt_library, slash_command::file_command, CacheStatus, Context,    ContextEvent, ContextId, ContextOperation, MessageId, MessageStatus, PromptBuilder,    WorkflowStepEditKind,};use anyhow::Result;use assistant_slash_command::{    ArgumentCompletion, SlashCommand, SlashCommandOutput, SlashCommandOutputSection,    SlashCommandRegistry,};use collections::HashSet;use fs::FakeFs;use gpui::{AppContext, Model, SharedString, Task, TestAppContext, WeakView};use language::{Buffer, LanguageRegistry, LspAdapterDelegate};use language_model::{LanguageModelCacheConfiguration, LanguageModelRegistry, Role};use parking_lot::Mutex;use project::Project;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};use ui::{Context as _, 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(), 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, None, None, prompt_builder.clone(), 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(), 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(file_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(), cx));     let output_ranges = Rc::new(RefCell::new(HashSet::default()));    context.update(cx, |_, cx| {        cx.subscribe(&context, {            let ranges = output_ranges.clone();            move |_, _, event, _| match event {                ContextEvent::PendingSlashCommandsUpdated { removed, updated } => {                    for range in removed {                        ranges.borrow_mut().remove(range);                    }                    for command in updated {                        ranges.borrow_mut().insert(command.source_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_output_ranges(        &buffer,        &output_ranges.borrow(),        "        «/file src/lib.rs»        "        .unindent()        .trim_end(),        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_output_ranges(        &buffer,        &output_ranges.borrow(),        "        «/file src/main.rs»        "        .unindent()        .trim_end(),        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_output_ranges(        &buffer,        &output_ranges.borrow(),        "        /unknown src/main.rs        "        .unindent()        .trim_end(),        cx,    );     #[track_caller]    fn assert_text_and_output_ranges(        buffer: &Model<Buffer>,        ranges: &HashSet<Range<language::Anchor>>,        expected_marked_text: &str,        cx: &mut TestAppContext,    ) {        let (expected_text, expected_ranges) = marked_text_ranges(expected_marked_text, false);        let (actual_text, actual_ranges) = buffer.update(cx, |buffer, _| {            let mut ranges = ranges                .iter()                .map(|range| range.to_offset(buffer))                .collect::<Vec<_>>();            ranges.sort_by_key(|a| a.start);            (buffer.text(), ranges)        });         assert_eq!(actual_text, expected_text);        assert_eq!(actual_ranges, expected_ranges);    }} #[gpui::test]async fn test_workflow_step_parsing(cx: &mut TestAppContext) {    cx.update(prompt_library::init);    let settings_store = cx.update(SettingsStore::test);    cx.set_global(settings_store);    cx.update(language::init);    cx.update(Project::init_settings);    let fs = FakeFs::new(cx.executor());    let project = Project::test(fs, [Path::new("/root")], cx).await;    cx.update(LanguageModelRegistry::test);     cx.update(assistant_panel::init);    let registry = Arc::new(LanguageRegistry::test(cx.executor()));     // Create a new context    let prompt_builder = Arc::new(PromptBuilder::new(None).unwrap());    let context = cx.new_model(|cx| {        Context::local(            registry.clone(),            Some(project),            None,            prompt_builder.clone(),            cx,

Lines 1–500 of 1438.