Get contents of symbols and selections

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

crates/acp_thread/src/mention.rs

@@ -1,13 +1,24 @@
-use anyhow::{Result, bail};
+use anyhow::{Context as _, Result, bail};
-use std::path::PathBuf;
+use std::{
+    ops::Range,
+    path::{Path, PathBuf},
+};
 
 #[derive(Clone, Debug, PartialEq, Eq)]
 pub enum MentionUri {
     File(PathBuf),
-    Symbol(PathBuf, String),
+    Symbol {
+        path: PathBuf,
+        name: String,
+        line_range: Range<u32>,
+    },
     Thread(String),
     TextThread(PathBuf),
     Rule(String),
+    Selection {
+        path: PathBuf,
+        line_range: Range<u32>,
+    },
 }
 
 impl MentionUri {
@@ -17,7 +28,40 @@ impl MentionUri {
         match url.scheme() {
             "file" => {
                 if let Some(fragment) = url.fragment() {
-                    Ok(Self::Symbol(path.into(), fragment.into()))
+                    let range = fragment
+                        .strip_prefix("L")
+                        .context("Line range must start with \"L\"")?;
+                    let (start, end) = range
+                        .split_once(":")
+                        .context("Line range must use colon as separator")?;
+                    let line_range = start
+                        .parse::<u32>()
+                        .context("Parsing line range start")?
+                        .checked_sub(1)
+                        .context("Line numbers should be 1-based")?
+                        ..end
+                            .parse::<u32>()
+                            .context("Parsing line range end")?
+                            .checked_sub(1)
+                            .context("Line numbers should be 1-based")?;
+                    let pairs = url.query_pairs().collect::<Vec<_>>();
+                    match pairs.as_slice() {
+                        [] => Ok(Self::Selection {
+                            path: path.into(),
+                            line_range,
+                        }),
+                        [(k, v)] => {
+                            if k != "symbol" {
+                                bail!("invalid query parameter")
+                            }
+                            Ok(Self::Symbol {
+                                name: v.to_string(),
+                                path: path.into(),
+                                line_range,
+                            })
+                        }
+                        _ => bail!("too many query pairs"),
+                    }
                 } else {
                     Ok(Self::File(path.into()))
                 }
@@ -37,11 +81,18 @@ impl MentionUri {
 
     pub fn name(&self) -> String {
         match self {
-            MentionUri::File(path) => path.file_name().unwrap().to_string_lossy().into_owned(),
+            MentionUri::File(path) => path
-            MentionUri::Symbol(_path, name) => name.clone(),
+                .file_name()
+                .unwrap_or_default()
+                .to_string_lossy()
+                .into_owned(),
+            MentionUri::Symbol { name, .. } => name.clone(),
             MentionUri::Thread(thread) => thread.to_string(),
             MentionUri::TextThread(thread) => thread.display().to_string(),
             MentionUri::Rule(rule) => rule.clone(),
+            MentionUri::Selection {
+                path, line_range, ..
+            } => selection_name(path, line_range),
         }
     }
 
@@ -57,8 +108,26 @@ impl MentionUri {
             MentionUri::File(path) => {
                 format!("file://{}", path.display())
             }
-            MentionUri::Symbol(path, name) => {
+            MentionUri::Symbol {
-                format!("file://{}#{}", path.display(), name)
+                path,
+                name,
+                line_range,
+            } => {
+                format!(
+                    "file://{}?symbol={}#L{}:{}",
+                    path.display(),
+                    name,
+                    line_range.start + 1,
+                    line_range.end + 1,
+                )
+            }
+            MentionUri::Selection { path, line_range } => {
+                format!(
+                    "file://{}#L{}:{}",
+                    path.display(),
+                    line_range.start + 1,
+                    line_range.end + 1,
+                )
             }
             MentionUri::Thread(thread) => {
                 format!("zed:///agent/thread/{}", thread)
@@ -73,13 +142,21 @@ impl MentionUri {
     }
 }
 
+pub fn selection_name(path: &Path, line_range: &Range<u32>) -> String {
+    format!(
+        "{} ({}:{})",
+        path.file_name().unwrap_or_default().display(),
+        line_range.start + 1,
+        line_range.end + 1
+    )
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
 
     #[test]
-    fn test_mention_uri_parse_and_display() {
+    fn test_parse_file_uri() {
-        // Test file URI
         let file_uri = "file:///path/to/file.rs";
         let parsed = MentionUri::parse(file_uri).unwrap();
         match &parsed {
@@ -87,20 +164,45 @@ mod tests {
             _ => panic!("Expected File variant"),
         }
         assert_eq!(parsed.to_uri(), file_uri);
+    }
 
-        // Test symbol URI
+    #[test]
-        let symbol_uri = "file:///path/to/file.rs#MySymbol";
+    fn test_parse_symbol_uri() {
+        let symbol_uri = "file:///path/to/file.rs?symbol=MySymbol#L10:20";
         let parsed = MentionUri::parse(symbol_uri).unwrap();
         match &parsed {
-            MentionUri::Symbol(path, symbol) => {
+            MentionUri::Symbol {
+                path,
+                name,
+                line_range,
+            } => {
                 assert_eq!(path.to_str().unwrap(), "/path/to/file.rs");
-                assert_eq!(symbol, "MySymbol");
+                assert_eq!(name, "MySymbol");
+                assert_eq!(line_range.start, 9);
+                assert_eq!(line_range.end, 19);
             }
             _ => panic!("Expected Symbol variant"),
         }
         assert_eq!(parsed.to_uri(), symbol_uri);
+    }
 
-        // Test thread URI
+    #[test]
+    fn test_parse_selection_uri() {
+        let selection_uri = "file:///path/to/file.rs#L5:15";
+        let parsed = MentionUri::parse(selection_uri).unwrap();
+        match &parsed {
+            MentionUri::Selection { path, line_range } => {
+                assert_eq!(path.to_str().unwrap(), "/path/to/file.rs");
+                assert_eq!(line_range.start, 4);
+                assert_eq!(line_range.end, 14);
+            }
+            _ => panic!("Expected Selection variant"),
+        }
+        assert_eq!(parsed.to_uri(), selection_uri);
+    }
+
+    #[test]
+    fn test_parse_thread_uri() {
         let thread_uri = "zed:///agent/thread/session123";
         let parsed = MentionUri::parse(thread_uri).unwrap();
         match &parsed {
@@ -108,8 +210,10 @@ mod tests {
             _ => panic!("Expected Thread variant"),
         }
         assert_eq!(parsed.to_uri(), thread_uri);
+    }
 
-        // Test rule URI
+    #[test]
+    fn test_parse_rule_uri() {
         let rule_uri = "zed:///agent/rule/my_rule";
         let parsed = MentionUri::parse(rule_uri).unwrap();
         match &parsed {
@@ -117,11 +221,50 @@ mod tests {
             _ => panic!("Expected Rule variant"),
         }
         assert_eq!(parsed.to_uri(), rule_uri);
+    }
 
-        // Test invalid scheme
+    #[test]
+    fn test_invalid_scheme() {
         assert!(MentionUri::parse("http://example.com").is_err());
+        assert!(MentionUri::parse("https://example.com").is_err());
+        assert!(MentionUri::parse("ftp://example.com").is_err());
+    }
 
-        // Test invalid zed path
+    #[test]
+    fn test_invalid_zed_path() {
         assert!(MentionUri::parse("zed:///invalid/path").is_err());
+        assert!(MentionUri::parse("zed:///agent/unknown/test").is_err());
+    }
+
+    #[test]
+    fn test_invalid_line_range_format() {
+        // Missing L prefix
+        assert!(MentionUri::parse("file:///path/to/file.rs#10:20").is_err());
+
+        // Missing colon separator
+        assert!(MentionUri::parse("file:///path/to/file.rs#L1020").is_err());
+
+        // Invalid numbers
+        assert!(MentionUri::parse("file:///path/to/file.rs#L10:abc").is_err());
+        assert!(MentionUri::parse("file:///path/to/file.rs#Labc:20").is_err());
+    }
+
+    #[test]
+    fn test_invalid_query_parameters() {
+        // Invalid query parameter name
+        assert!(MentionUri::parse("file:///path/to/file.rs#L10:20?invalid=test").is_err());
+
+        // Too many query parameters
+        assert!(
+            MentionUri::parse("file:///path/to/file.rs#L10:20?symbol=test&another=param").is_err()
+        );
+    }
+
+    #[test]
+    fn test_zero_based_line_numbers() {
+        // Test that 0-based line numbers are rejected (should be 1-based)
+        assert!(MentionUri::parse("file:///path/to/file.rs#L0:10").is_err());
+        assert!(MentionUri::parse("file:///path/to/file.rs#L1:0").is_err());
+        assert!(MentionUri::parse("file:///path/to/file.rs#L0:0").is_err());
     }
 }

crates/agent2/src/thread.rs

@@ -26,8 +26,8 @@ use schemars::{JsonSchema, Schema};
 use serde::{Deserialize, Serialize};
 use settings::{Settings, update_settings_file};
 use smol::stream::StreamExt;
-use std::fmt::Write;
 use std::{cell::RefCell, collections::BTreeMap, path::Path, rc::Rc, sync::Arc};
+use std::{fmt::Write, ops::Range};
 use util::{ResultExt, markdown::MarkdownCodeBlock};
 
 #[derive(Debug, Clone)]
@@ -1182,17 +1182,33 @@ impl AgentMessage {
                 }
                 MessageContent::Mention { uri, content } => {
                     match uri {
-                        MentionUri::File(path) | MentionUri::Symbol(path, _) => {
+                        MentionUri::File(path) => {
                             write!(
                                 &mut symbol_context,
                                 "\n{}",
                                 MarkdownCodeBlock {
-                                    tag: &codeblock_tag(&path),
+                                    tag: &codeblock_tag(&path, None),
                                     text: &content.to_string(),
                                 }
                             )
                             .ok();
                         }
+                        MentionUri::Symbol {
+                            path, line_range, ..
+                        }
+                        | MentionUri::Selection {
+                            path, line_range, ..
+                        } => {
+                            write!(
+                                &mut rules_context,
+                                "\n{}",
+                                MarkdownCodeBlock {
+                                    tag: &codeblock_tag(&path, Some(line_range)),
+                                    text: &content
+                                }
+                            )
+                            .ok();
+                        }
                         MentionUri::Thread(_session_id) => {
                             write!(&mut thread_context, "\n{}\n", content).ok();
                         }
@@ -1263,7 +1279,7 @@ impl AgentMessage {
     }
 }
 
-fn codeblock_tag(full_path: &Path) -> String {
+fn codeblock_tag(full_path: &Path, line_range: Option<&Range<u32>>) -> String {
     let mut result = String::new();
 
     if let Some(extension) = full_path.extension().and_then(|ext| ext.to_str()) {
@@ -1272,6 +1288,14 @@ fn codeblock_tag(full_path: &Path) -> String {
 
     let _ = write!(result, "{}", full_path.display());
 
+    if let Some(range) = line_range {
+        if range.start == range.end {
+            let _ = write!(result, ":{}", range.start + 1);
+        } else {
+            let _ = write!(result, ":{}-{}", range.start + 1, range.end + 1);
+        }
+    }
+
     result
 }
 

crates/agent_ui/src/acp/completion_provider.rs

@@ -3,16 +3,16 @@ use std::path::{Path, PathBuf};
 use std::sync::Arc;
 use std::sync::atomic::AtomicBool;
 
-use acp_thread::MentionUri;
+use acp_thread::{MentionUri, selection_name};
-use agent::context_store::ContextStore;
 use anyhow::{Context as _, Result};
 use collections::{HashMap, HashSet};
 use editor::display_map::CreaseId;
-use editor::{CompletionProvider, Editor, ExcerptId};
+use editor::{AnchorRangeExt, CompletionProvider, Editor, ExcerptId, ToOffset as _, ToPoint};
 use file_icons::FileIcons;
 use futures::future::try_join_all;
 use fuzzy::{StringMatch, StringMatchCandidate};
 use gpui::{App, Entity, Task, WeakEntity};
+use itertools::Itertools as _;
 use language::{Buffer, CodeLabel, HighlightId};
 use lsp::CompletionContext;
 use parking_lot::Mutex;
@@ -21,14 +21,12 @@ use project::{
 };
 use prompt_store::PromptStore;
 use rope::Point;
-use text::{Anchor, ToPoint};
+use text::{Anchor, ToPoint as _};
 use ui::prelude::*;
-use util::ResultExt as _;
 use workspace::Workspace;
 
 use agent::{
-    Thread,
+    context::RULES_ICON,
-    context::{AgentContextHandle, AgentContextKey, RULES_ICON},
     thread_store::{TextThreadStore, ThreadStore},
 };
 
@@ -40,10 +38,9 @@ use crate::context_picker::thread_context_picker::{
     ThreadContextEntry, ThreadMatch, search_threads,
 };
 use crate::context_picker::{
-    ContextPickerEntry, ContextPickerMode, RecentEntry, available_context_picker_entries,
+    ContextPickerAction, ContextPickerEntry, ContextPickerMode, RecentEntry,
-    recent_context_picker_entries,
+    available_context_picker_entries, recent_context_picker_entries, selection_ranges,
 };
-use crate::message_editor::ContextCreasesAddon;
 
 #[derive(Default)]
 pub struct MentionSet {
@@ -86,10 +83,44 @@ impl MentionSet {
                         anyhow::Ok((crease_id, Mention { uri, content }))
                     })
                 }
-                _ => {
+                MentionUri::Symbol {
-                    // TODO
+                    path, line_range, ..
-                    unimplemented!()
                 }
+                | MentionUri::Selection {
+                    path, line_range, ..
+                } => {
+                    let crease_id = *crease_id;
+                    let uri = uri.clone();
+                    let path_buf = path.clone();
+                    let line_range = line_range.clone();
+
+                    let buffer_task = project.update(cx, |project, cx| {
+                        let path = project
+                            .find_project_path(&path_buf, cx)
+                            .context("Failed to find project path")?;
+                        anyhow::Ok(project.open_buffer(path, cx))
+                    });
+
+                    cx.spawn(async move |cx| {
+                        let buffer = buffer_task?.await?;
+                        let content = buffer.read_with(cx, |buffer, _cx| {
+                            buffer
+                                .text_for_range(
+                                    Point::new(line_range.start, 0)
+                                        ..Point::new(
+                                            line_range.end,
+                                            buffer.line_len(line_range.end),
+                                        ),
+                                )
+                                .collect()
+                        })?;
+
+                        anyhow::Ok((crease_id, Mention { uri, content }))
+                    })
+                }
+                MentionUri::Thread(_) => todo!(),
+                MentionUri::TextThread(path_buf) => todo!(),
+                MentionUri::Rule(_) => todo!(),
             })
             .collect::<Vec<_>>();
 
@@ -136,8 +167,8 @@ fn search(
     cancellation_flag: Arc<AtomicBool>,
     recent_entries: Vec<RecentEntry>,
     prompt_store: Option<Entity<PromptStore>>,
-    thread_store: Option<WeakEntity<ThreadStore>>,
+    thread_store: WeakEntity<ThreadStore>,
-    text_thread_context_store: Option<WeakEntity<assistant_context::ContextStore>>,
+    text_thread_context_store: WeakEntity<assistant_context::ContextStore>,
     workspace: Entity<Workspace>,
     cx: &mut App,
 ) -> Task<Vec<Match>> {
@@ -168,9 +199,8 @@ fn search(
 
         Some(ContextPickerMode::Thread) => {
             if let Some((thread_store, context_store)) = thread_store
-                .as_ref()
+                .upgrade()
-                .and_then(|t| t.upgrade())
+                .zip(text_thread_context_store.upgrade())
-                .zip(text_thread_context_store.as_ref().and_then(|t| t.upgrade()))
             {
                 let search_threads_task = search_threads(
                     query.clone(),
@@ -240,14 +270,19 @@ fn search(
                     .collect::<Vec<_>>();
 
                 matches.extend(
-                    available_context_picker_entries(&prompt_store, &thread_store, &workspace, cx)
+                    available_context_picker_entries(
-                        .into_iter()
+                        &prompt_store,
-                        .map(|mode| {
+                        &Some(thread_store.clone()),
-                            Match::Entry(EntryMatch {
+                        &workspace,
-                                entry: mode,
+                        cx,
-                                mat: None,
+                    )
-                            })
+                    .into_iter()
-                        }),
+                    .map(|mode| {
+                        Match::Entry(EntryMatch {
+                            entry: mode,
+                            mat: None,
+                        })
+                    }),
                 );
 
                 Task::ready(matches)
@@ -257,8 +292,12 @@ fn search(
                 let search_files_task =
                     search_files(query.clone(), cancellation_flag.clone(), &workspace, cx);
 
-                let entries =
+                let entries = available_context_picker_entries(
-                    available_context_picker_entries(&prompt_store, &thread_store, &workspace, cx);
+                    &prompt_store,
+                    &Some(thread_store.clone()),
+                    &workspace,
+                    cx,
+                );
                 let entry_candidates = entries
                     .iter()
                     .enumerate()
@@ -352,117 +391,101 @@ impl ContextPickerCompletionProvider {
                 confirm: Some(Arc::new(|_, _, _| true)),
             }),
             ContextPickerEntry::Action(action) => {
-                todo!()
+                let (new_text, on_action) = match action {
-                // let (new_text, on_action) = match action {
+                    ContextPickerAction::AddSelections => {
-                //     ContextPickerAction::AddSelections => {
+                        let selections = selection_ranges(workspace, cx);
-                //         let selections = selection_ranges(workspace, cx);
 
-                //         let selection_infos = selections
+                        const PLACEHOLDER: &str = "selection ";
-                //             .iter()
-                //             .map(|(buffer, range)| {
-                //                 let full_path = buffer
-                //                     .read(cx)
-                //                     .file()
-                //                     .map(|file| file.full_path(cx))
-                //                     .unwrap_or_else(|| PathBuf::from("untitled"));
-                //                 let file_name = full_path
-                //                     .file_name()
-                //                     .unwrap_or_default()
-                //                     .to_string_lossy()
-                //                     .to_string();
-                //                 let line_range = range.to_point(&buffer.read(cx).snapshot());
 
-                //                 let link = MentionLink::for_selection(
+                        let new_text = std::iter::repeat(PLACEHOLDER)
-                //                     &file_name,
+                            .take(selections.len())
-                //                     &full_path.to_string_lossy(),
+                            .chain(std::iter::once(""))
-                //                     line_range.start.row as usize..line_range.end.row as usize,
+                            .join(" ");
-                //                 );
-                //                 (file_name, link, line_range)
-                //             })
-                //             .collect::<Vec<_>>();
 
-                //         let new_text = format!(
+                        let callback = Arc::new({
-                //             "{} ",
+                            let mention_set = mention_set.clone();
-                //             selection_infos.iter().map(|(_, link, _)| link).join(" ")
+                            let selections = selections.clone();
-                //         );
+                            move |_, window: &mut Window, cx: &mut App| {
+                                let editor = editor.clone();
+                                let mention_set = mention_set.clone();
+                                let selections = selections.clone();
+                                window.defer(cx, move |window, cx| {
+                                    let mut current_offset = 0;
 
-                //         let callback = Arc::new({
+                                    for (buffer, selection_range) in selections {
-                //             let context_store = context_store.clone();
+                                        let snapshot =
-                //             let selections = selections.clone();
+                                            editor.read(cx).buffer().read(cx).snapshot(cx);
-                //             let selection_infos = selection_infos.clone();
+                                        let Some(start) = snapshot
-                //             move |_, window: &mut Window, cx: &mut App| {
+                                            .anchor_in_excerpt(excerpt_id, source_range.start)
-                //                 context_store.update(cx, |context_store, cx| {
+                                        else {
-                //                     for (buffer, range) in &selections {
+                                            return;
-                //                         context_store.add_selection(
+                                        };
-                //                             buffer.clone(),
-                //                             range.clone(),
-                //                             cx,
-                //                         );
-                //                     }
-                //                 });
 
-                //                 let editor = editor.clone();
+                                        let offset = start.to_offset(&snapshot) + current_offset;
-                //                 let selection_infos = selection_infos.clone();
+                                        let text_len = PLACEHOLDER.len() - 1;
-                //                 window.defer(cx, move |window, cx| {
-                //                     let mut current_offset = 0;
-                //                     for (file_name, link, line_range) in selection_infos.iter() {
-                //                         let snapshot =
-                //                             editor.read(cx).buffer().read(cx).snapshot(cx);
-                //                         let Some(start) = snapshot
-                //                             .anchor_in_excerpt(excerpt_id, source_range.start)
-                //                         else {
-                //                             return;
-                //                         };
 
-                //                         let offset = start.to_offset(&snapshot) + current_offset;
+                                        let range = snapshot.anchor_after(offset)
-                //                         let text_len = link.len();
+                                            ..snapshot.anchor_after(offset + text_len);
 
-                //                         let range = snapshot.anchor_after(offset)
+                                        let path = buffer
-                //                             ..snapshot.anchor_after(offset + text_len);
+                                            .read(cx)
+                                            .file()
+                                            .map_or(PathBuf::from("untitled"), |file| {
+                                                file.path().to_path_buf()
+                                            });
 
-                //                         let crease = super::crease_for_mention(
+                                        let point_range = range.to_point(&snapshot);
-                //                             format!(
+                                        let line_range = point_range.start.row..point_range.end.row;
-                //                                 "{} ({}-{})",
+                                        let crease = crate::context_picker::crease_for_mention(
-                //                                 file_name,
+                                            selection_name(&path, &line_range).into(),
-                //                                 line_range.start.row + 1,
+                                            IconName::Reader.path().into(),
-                //                                 line_range.end.row + 1
+                                            range,
-                //                             )
+                                            editor.downgrade(),
-                //                             .into(),
+                                        );
-                //                             IconName::Reader.path().into(),
-                //                             range,
-                //                             editor.downgrade(),
-                //                         );
 
-                //                         editor.update(cx, |editor, cx| {
+                                        let [crease_id]: [_; 1] =
-                //                             editor.insert_creases(vec![crease.clone()], cx);
+                                            editor.update(cx, |editor, cx| {
-                //                             editor.fold_creases(vec![crease], false, window, cx);
+                                                let crease_ids =
-                //                         });
+                                                    editor.insert_creases(vec![crease.clone()], cx);
+                                                editor.fold_creases(
+                                                    vec![crease],
+                                                    false,
+                                                    window,
+                                                    cx,
+                                                );
+                                                crease_ids.try_into().unwrap()
+                                            });
 
-                //                         current_offset += text_len + 1;
+                                        mention_set.lock().insert(
-                //                     }
+                                            crease_id,
-                //                 });
+                                            MentionUri::Selection { path, line_range },
+                                        );
 
-                //                 false
+                                        current_offset += text_len + 1;
-                //             }
+                                    }
-                //         });
+                                });
 
-                //         (new_text, callback)
+                                false
-                //     }
+                            }
-                // };
+                        });
 
-                // Some(Completion {
+                        (new_text, callback)
-                //     replace_range: source_range.clone(),
+                    }
-                //     new_text,
+                };
-                //     label: CodeLabel::plain(action.label().to_string(), None),
+
-                //     icon_path: Some(action.icon().path().into()),
+                Some(Completion {
-                //     documentation: None,
+                    replace_range: source_range.clone(),
-                //     source: project::CompletionSource::Custom,
+                    new_text,
-                //     insert_text_mode: None,
+                    label: CodeLabel::plain(action.label().to_string(), None),
-                //     // This ensures that when a user accepts this completion, the
+                    icon_path: Some(action.icon().path().into()),
-                //     // completion menu will still be shown after "@category " is
+                    documentation: None,
-                //     // inserted
+                    source: project::CompletionSource::Custom,
-                //     confirm: Some(on_action),
+                    insert_text_mode: None,
-                // })
+                    // This ensures that when a user accepts this completion, the
+                    // completion menu will still be shown after "@category " is
+                    // inserted
+                    confirm: Some(on_action),
+                })
             }
         }
     }
@@ -636,11 +659,12 @@ impl ContextPickerCompletionProvider {
 
         let comment_id = cx.theme().syntax().highlight_id("comment").map(HighlightId);
         let mut label = CodeLabel::plain(symbol.name.clone(), None);
-        label.push_str(" ", None);
-        label.push_str(&file_name, comment_id);
-        label.push_str(&format!(" L{}", symbol.range.start.0.row + 1), comment_id);
 
-        let uri = MentionUri::Symbol(full_path.into(), symbol.name.clone());
+        let uri = MentionUri::Symbol {
+            path: full_path.into(),
+            name: symbol.name.clone(),
+            line_range: symbol.range.start.0.row..symbol.range.end.0.row,
+        };
         let new_text = format!("{} ", uri.to_link());
         let new_text_len = new_text.len();
         Some(Completion {
@@ -741,20 +765,18 @@ impl CompletionProvider for ContextPickerCompletionProvider {
         };
 
         let recent_entries = recent_context_picker_entries(
-            thread_store.clone(),
+            Some(thread_store.clone()),
-            text_thread_store.clone(),
+            Some(text_thread_store.clone()),
             workspace.clone(),
             &exclude_paths,
             &exclude_threads,
             cx,
         );
 
-        let prompt_store = thread_store.as_ref().and_then(|thread_store| {
+        let prompt_store = thread_store
-            thread_store
+            .read_with(cx, |thread_store, _cx| thread_store.prompt_store().clone())
-                .read_with(cx, |thread_store, _cx| thread_store.prompt_store().clone())
+            .ok()
-                .ok()
+            .flatten();
-                .flatten()
-        });
 
         let search_task = search(
             mode,

crates/agent_ui/src/acp/thread_view.rs

@@ -106,7 +106,7 @@ impl AcpThreadView {
         workspace: WeakEntity<Workspace>,
         project: Entity<Project>,
         thread_store: WeakEntity<ThreadStore>,
-        text_thread_store: WeaksEntity<TextThreadStore>,
+        text_thread_store: WeakEntity<TextThreadStore>,
         message_history: Rc<RefCell<MessageHistory<Vec<acp::ContentBlock>>>>,
         min_lines: usize,
         max_lines: Option<usize>,
@@ -3478,6 +3478,8 @@ mod tests {
                     Rc::new(agent),
                     workspace.downgrade(),
                     project,
+                    WeakEntity::new_invalid(),
+                    WeakEntity::new_invalid(),
                     Rc::new(RefCell::new(MessageHistory::default())),
                     1,
                     None,

crates/agent_ui/src/agent_panel.rs

@@ -965,8 +965,8 @@ impl AgentPanel {
                         server,
                         workspace.clone(),
                         project,
-                        thread_store.clone(),
+                        thread_store.downgrade(),
-                        text_thread_store.clone(),
+                        text_thread_store.downgrade(),
                         message_history,
                         MIN_EDITOR_LINES,
                         Some(MAX_EDITOR_LINES),

crates/agent_ui/src/context_picker.rs

@@ -742,7 +742,7 @@ fn add_selections_as_context(
     })
 }
 
-fn selection_ranges(
+pub(crate) fn selection_ranges(
     workspace: &Entity<Workspace>,
     cx: &mut App,
 ) -> Vec<(Entity<Buffer>, Range<text::Anchor>)> {