Test mention set contents

Co-authored-by: Cole Miller <cole@zed.dev>

crates/agent/src/thread_store.rs

@@ -205,6 +205,22 @@ impl ThreadStore {
         (this, ready_rx)
     }
 
+    #[cfg(any(test, feature = "test-support"))]
+    pub fn fake(project: Entity<Project>, cx: &mut App) -> Self {
+        Self {
+            project,
+            tools: cx.new(|_| ToolWorkingSet::default()),
+            prompt_builder: Arc::new(PromptBuilder::new(None).unwrap()),
+            prompt_store: None,
+            context_server_tool_ids: HashMap::default(),
+            threads: Vec::new(),
+            project_context: SharedProjectContext::default(),
+            reload_system_prompt_tx: mpsc::channel(0).0,
+            _reload_system_prompt_task: Task::ready(()),
+            _subscriptions: vec![],
+        }
+    }
+
     fn handle_project_event(
         &mut self,
         _project: Entity<Project>,

crates/agent_ui/Cargo.toml

@@ -102,6 +102,8 @@ workspace.workspace = true
 zed_actions.workspace = true
 
 [dev-dependencies]
+agent = { workspace = true, features = ["test-support"] }
+assistant_context = { workspace = true, features = ["test-support"] }
 assistant_tools.workspace = true
 buffer_diff = { workspace = true, features = ["test-support"] }
 editor = { workspace = true, features = ["test-support"] }

crates/agent_ui/src/acp/completion_provider.rs

@@ -167,6 +167,7 @@ impl MentionSet {
     }
 }
 
+#[derive(Debug)]
 pub struct Mention {
     pub uri: MentionUri,
     pub content: String,
@@ -1166,16 +1167,16 @@ mod tests {
                 json!({
                     "editor": "",
                     "a": {
-                        "one.txt": "",
+                        "one.txt": "1",
-                        "two.txt": "",
+                        "two.txt": "2",
-                        "three.txt": "",
+                        "three.txt": "3",
-                        "four.txt": ""
+                        "four.txt": "4"
                     },
                     "b": {
-                        "five.txt": "",
+                        "five.txt": "5",
-                        "six.txt": "",
+                        "six.txt": "6",
-                        "seven.txt": "",
+                        "seven.txt": "7",
-                        "eight.txt": "",
+                        "eight.txt": "8",
                     }
                 }),
             )
@@ -1250,14 +1251,17 @@ mod tests {
 
         let mention_set = Arc::new(Mutex::new(MentionSet::default()));
 
+        let thread_store = cx.new(|cx| ThreadStore::fake(project.clone(), cx));
+        let text_thread_store = cx.new(|cx| TextThreadStore::fake(project.clone(), cx));
+
         let editor_entity = editor.downgrade();
         editor.update_in(&mut cx, |editor, window, cx| {
             window.focus(&editor.focus_handle(cx));
             editor.set_completion_provider(Some(Rc::new(ContextPickerCompletionProvider::new(
                 mention_set.clone(),
                 workspace.downgrade(),
-                WeakEntity::new_invalid(),
+                thread_store.downgrade(),
-                WeakEntity::new_invalid(),
+                text_thread_store.downgrade(),
                 editor_entity,
             ))));
         });
@@ -1328,6 +1332,28 @@ mod tests {
             );
         });
 
+        let contents = cx
+            .update(|window, cx| {
+                mention_set.lock().contents(
+                    project.clone(),
+                    thread_store.clone(),
+                    text_thread_store.clone(),
+                    window,
+                    cx,
+                )
+            })
+            .await
+            .unwrap()
+            .into_values()
+            .collect::<Vec<_>>();
+
+        assert_eq!(contents.len(), 1);
+        assert_eq!(contents[0].content, "1");
+        assert_eq!(
+            contents[0].uri.to_uri().to_string(),
+            "file:///dir/a/one.txt"
+        );
+
         cx.simulate_input(" ");
 
         editor.update(&mut cx, |editor, cx| {
@@ -1373,6 +1399,28 @@ mod tests {
 
         cx.run_until_parked();
 
+        let contents = cx
+            .update(|window, cx| {
+                mention_set.lock().contents(
+                    project.clone(),
+                    thread_store.clone(),
+                    text_thread_store.clone(),
+                    window,
+                    cx,
+                )
+            })
+            .await
+            .unwrap()
+            .into_values()
+            .collect::<Vec<_>>();
+
+        assert_eq!(contents.len(), 2);
+        let new_mention = contents
+            .iter()
+            .find(|mention| mention.uri.to_uri().to_string() == "file:///dir/b/eight.txt")
+            .unwrap();
+        assert_eq!(new_mention.content, "8");
+
         editor.update(&mut cx, |editor, cx| {
             assert_eq!(
                 editor.text(cx),
@@ -1388,44 +1436,6 @@ mod tests {
             );
         });
 
-        cx.simulate_input("\n@");
-
-        editor.update(&mut cx, |editor, cx| {
-            assert_eq!(
-                editor.text(cx),
-                "Lorem [@one.txt](file:///dir/a/one.txt)  Ipsum [@eight.txt](file:///dir/b/eight.txt) \n@"
-            );
-            assert!(editor.has_visible_completions_menu());
-            assert_eq!(
-                fold_ranges(editor, cx),
-                vec![
-                    Point::new(0, 6)..Point::new(0, 39),
-                    Point::new(0, 47)..Point::new(0, 84)
-                ]
-            );
-        });
-
-        editor.update_in(&mut cx, |editor, window, cx| {
-            editor.confirm_completion(&editor::actions::ConfirmCompletion::default(), window, cx);
-        });
-
-        cx.run_until_parked();
-
-        editor.update(&mut cx, |editor, cx| {
-            assert_eq!(
-                editor.text(cx),
-                "Lorem [@one.txt](file:///dir/a/one.txt)  Ipsum [@eight.txt](file:///dir/b/eight.txt) \n[@seven.txt](file:///dir/b/seven.txt) "
-            );
-            assert!(!editor.has_visible_completions_menu());
-            assert_eq!(
-                fold_ranges(editor, cx),
-                vec![
-                    Point::new(0, 6)..Point::new(0, 39),
-                    Point::new(0, 47)..Point::new(0, 84),
-                    Point::new(1, 0)..Point::new(1, 37)
-                ]
-            );
-        });
         let plain_text_language = Arc::new(language::Language::new(
             language::LanguageConfig {
                 name: "Plain Text".into(),
@@ -1478,8 +1488,8 @@ mod tests {
                         location: lsp::Location {
                             uri: lsp::Url::from_file_path(path!("/dir/a/one.txt")).unwrap(),
                             range: lsp::Range::new(
-                                lsp::Position::new(0, 6),
+                                lsp::Position::new(0, 0),
-                                lsp::Position::new(0, 9),
+                                lsp::Position::new(0, 1),
                             ),
                         },
                         kind: lsp::SymbolKind::CONSTANT,
@@ -1496,7 +1506,7 @@ mod tests {
         editor.update(&mut cx, |editor, cx| {
             assert_eq!(
                 editor.text(cx),
-                "Lorem [@one.txt](file:///dir/a/one.txt)  Ipsum [@eight.txt](file:///dir/b/eight.txt) \n[@seven.txt](file:///dir/b/seven.txt) @symbol "
+                "Lorem [@one.txt](file:///dir/a/one.txt)  Ipsum [@eight.txt](file:///dir/b/eight.txt) @symbol "
             );
             assert!(editor.has_visible_completions_menu());
             assert_eq!(
@@ -1511,12 +1521,36 @@ mod tests {
             editor.confirm_completion(&editor::actions::ConfirmCompletion::default(), window, cx);
         });
 
+        let contents = cx
+            .update(|window, cx| {
+                mention_set.lock().contents(
+                    project.clone(),
+                    thread_store,
+                    text_thread_store,
+                    window,
+                    cx,
+                )
+            })
+            .await
+            .unwrap()
+            .into_values()
+            .collect::<Vec<_>>();
+
+        assert_eq!(contents.len(), 3);
+        let new_mention = contents
+            .iter()
+            .find(|mention| {
+                mention.uri.to_uri().to_string() == "file:///dir/a/one.txt?symbol=MySymbol#L1:1"
+            })
+            .unwrap();
+        assert_eq!(new_mention.content, "1");
+
         cx.run_until_parked();
 
         editor.read_with(&mut cx, |editor, cx| {
             assert_eq!(
                 editor.text(cx),
-                "Lorem [@one.txt](file:///dir/a/one.txt)  Ipsum [@eight.txt](file:///dir/b/eight.txt) \n[@seven.txt](file:///dir/b/seven.txt) [@MySymbol](file:///dir/a/one.txt?symbol=MySymbol#L1:1) "
+                "Lorem [@one.txt](file:///dir/a/one.txt)  Ipsum [@eight.txt](file:///dir/b/eight.txt) [@MySymbol](file:///dir/a/one.txt?symbol=MySymbol#L1:1) "
             );
         });
     }

crates/assistant_context/Cargo.toml

@@ -11,6 +11,9 @@ workspace = true
 [lib]
 path = "src/assistant_context.rs"
 
+[features]
+test-support = []
+
 [dependencies]
 agent_settings.workspace = true
 anyhow.workspace = true

crates/assistant_context/src/context_store.rs

@@ -138,6 +138,27 @@ impl ContextStore {
         })
     }
 
+    #[cfg(any(test, feature = "test-support"))]
+    pub fn fake(project: Entity<Project>, cx: &mut Context<Self>) -> Self {
+        Self {
+            contexts: Default::default(),
+            contexts_metadata: Default::default(),
+            context_server_slash_command_ids: Default::default(),
+            host_contexts: Default::default(),
+            fs: project.read(cx).fs().clone(),
+            languages: project.read(cx).languages().clone(),
+            slash_commands: Arc::default(),
+            telemetry: project.read(cx).client().telemetry().clone(),
+            _watch_updates: Task::ready(None),
+            client: project.read(cx).client(),
+            project,
+            project_is_shared: false,
+            client_subscription: None,
+            _project_subscriptions: Default::default(),
+            prompt_builder: Arc::new(PromptBuilder::new(None).unwrap()),
+        }
+    }
+
     async fn handle_advertise_contexts(
         this: Entity<Self>,
         envelope: TypedEnvelope<proto::AdvertiseContexts>,