1use crate::context::{
2 AgentContextHandle, AgentContextKey, ContextId, ContextKind, DirectoryContextHandle,
3 FetchedUrlContext, FileContextHandle, ImageContext, RulesContextHandle, SelectionContextHandle,
4 SymbolContextHandle, TextThreadContextHandle, ThreadContextHandle,
5};
6use agent_client_protocol as acp;
7use anyhow::{Context as _, Result, anyhow};
8use assistant_text_thread::TextThread;
9use collections::{HashSet, IndexSet};
10use futures::{self, FutureExt};
11use gpui::{App, Context, Entity, EventEmitter, Image, SharedString, Task, WeakEntity};
12use language::{Buffer, File as _};
13use language_model::LanguageModelImage;
14use project::{
15 Project, ProjectItem, ProjectPath, Symbol, image_store::is_image_file,
16 lsp_store::SymbolLocation,
17};
18use prompt_store::UserPromptId;
19use ref_cast::RefCast as _;
20use std::{
21 ops::Range,
22 path::{Path, PathBuf},
23 sync::Arc,
24};
25use text::{Anchor, OffsetRangeExt};
26
27pub struct ContextStore {
28 project: WeakEntity<Project>,
29 next_context_id: ContextId,
30 context_set: IndexSet<AgentContextKey>,
31 context_thread_ids: HashSet<acp::SessionId>,
32 context_text_thread_paths: HashSet<Arc<Path>>,
33}
34
35pub enum ContextStoreEvent {
36 ContextRemoved(AgentContextKey),
37}
38
39impl EventEmitter<ContextStoreEvent> for ContextStore {}
40
41impl ContextStore {
42 pub fn new(project: WeakEntity<Project>) -> Self {
43 Self {
44 project,
45 next_context_id: ContextId::zero(),
46 context_set: IndexSet::default(),
47 context_thread_ids: HashSet::default(),
48 context_text_thread_paths: HashSet::default(),
49 }
50 }
51
52 pub fn context(&self) -> impl Iterator<Item = &AgentContextHandle> {
53 self.context_set.iter().map(|entry| entry.as_ref())
54 }
55
56 pub fn clear(&mut self, cx: &mut Context<Self>) {
57 self.context_set.clear();
58 self.context_thread_ids.clear();
59 cx.notify();
60 }
61
62 pub fn add_file_from_path(
63 &mut self,
64 project_path: ProjectPath,
65 remove_if_exists: bool,
66 cx: &mut Context<Self>,
67 ) -> Task<Result<Option<AgentContextHandle>>> {
68 let Some(project) = self.project.upgrade() else {
69 return Task::ready(Err(anyhow!("failed to read project")));
70 };
71
72 if is_image_file(&project, &project_path, cx) {
73 self.add_image_from_path(project_path, remove_if_exists, cx)
74 } else {
75 cx.spawn(async move |this, cx| {
76 let open_buffer_task = project.update(cx, |project, cx| {
77 project.open_buffer(project_path.clone(), cx)
78 })?;
79 let buffer = open_buffer_task.await?;
80 this.update(cx, |this, cx| {
81 this.add_file_from_buffer(&project_path, buffer, remove_if_exists, cx)
82 })
83 })
84 }
85 }
86
87 pub fn add_file_from_buffer(
88 &mut self,
89 project_path: &ProjectPath,
90 buffer: Entity<Buffer>,
91 remove_if_exists: bool,
92 cx: &mut Context<Self>,
93 ) -> Option<AgentContextHandle> {
94 let context_id = self.next_context_id.post_inc();
95 let context = AgentContextHandle::File(FileContextHandle { buffer, context_id });
96
97 if let Some(key) = self.context_set.get(AgentContextKey::ref_cast(&context)) {
98 if remove_if_exists {
99 self.remove_context(&context, cx);
100 None
101 } else {
102 Some(key.as_ref().clone())
103 }
104 } else if self.path_included_in_directory(project_path, cx).is_some() {
105 None
106 } else {
107 self.insert_context(context.clone(), cx);
108 Some(context)
109 }
110 }
111
112 pub fn add_directory(
113 &mut self,
114 project_path: &ProjectPath,
115 remove_if_exists: bool,
116 cx: &mut Context<Self>,
117 ) -> Result<Option<AgentContextHandle>> {
118 let project = self.project.upgrade().context("failed to read project")?;
119 let entry_id = project
120 .read(cx)
121 .entry_for_path(project_path, cx)
122 .map(|entry| entry.id)
123 .context("no entry found for directory context")?;
124
125 let context_id = self.next_context_id.post_inc();
126 let context = AgentContextHandle::Directory(DirectoryContextHandle {
127 entry_id,
128 context_id,
129 });
130
131 let context =
132 if let Some(existing) = self.context_set.get(AgentContextKey::ref_cast(&context)) {
133 if remove_if_exists {
134 self.remove_context(&context, cx);
135 None
136 } else {
137 Some(existing.as_ref().clone())
138 }
139 } else {
140 self.insert_context(context.clone(), cx);
141 Some(context)
142 };
143
144 anyhow::Ok(context)
145 }
146
147 pub fn add_symbol(
148 &mut self,
149 buffer: Entity<Buffer>,
150 symbol: SharedString,
151 range: Range<Anchor>,
152 enclosing_range: Range<Anchor>,
153 remove_if_exists: bool,
154 cx: &mut Context<Self>,
155 ) -> (Option<AgentContextHandle>, bool) {
156 let context_id = self.next_context_id.post_inc();
157 let context = AgentContextHandle::Symbol(SymbolContextHandle {
158 buffer,
159 symbol,
160 range,
161 enclosing_range,
162 context_id,
163 });
164
165 if let Some(key) = self.context_set.get(AgentContextKey::ref_cast(&context)) {
166 let handle = if remove_if_exists {
167 self.remove_context(&context, cx);
168 None
169 } else {
170 Some(key.as_ref().clone())
171 };
172 return (handle, false);
173 }
174
175 let included = self.insert_context(context.clone(), cx);
176 (Some(context), included)
177 }
178
179 pub fn add_thread(
180 &mut self,
181 thread: Entity<agent::Thread>,
182 remove_if_exists: bool,
183 cx: &mut Context<Self>,
184 ) -> Option<AgentContextHandle> {
185 let context_id = self.next_context_id.post_inc();
186 let context = AgentContextHandle::Thread(ThreadContextHandle { thread, context_id });
187
188 if let Some(existing) = self.context_set.get(AgentContextKey::ref_cast(&context)) {
189 if remove_if_exists {
190 self.remove_context(&context, cx);
191 None
192 } else {
193 Some(existing.as_ref().clone())
194 }
195 } else {
196 self.insert_context(context.clone(), cx);
197 Some(context)
198 }
199 }
200
201 pub fn add_text_thread(
202 &mut self,
203 text_thread: Entity<TextThread>,
204 remove_if_exists: bool,
205 cx: &mut Context<Self>,
206 ) -> Option<AgentContextHandle> {
207 let context_id = self.next_context_id.post_inc();
208 let context = AgentContextHandle::TextThread(TextThreadContextHandle {
209 text_thread,
210 context_id,
211 });
212
213 if let Some(existing) = self.context_set.get(AgentContextKey::ref_cast(&context)) {
214 if remove_if_exists {
215 self.remove_context(&context, cx);
216 None
217 } else {
218 Some(existing.as_ref().clone())
219 }
220 } else {
221 self.insert_context(context.clone(), cx);
222 Some(context)
223 }
224 }
225
226 pub fn add_rules(
227 &mut self,
228 prompt_id: UserPromptId,
229 remove_if_exists: bool,
230 cx: &mut Context<ContextStore>,
231 ) -> Option<AgentContextHandle> {
232 let context_id = self.next_context_id.post_inc();
233 let context = AgentContextHandle::Rules(RulesContextHandle {
234 prompt_id,
235 context_id,
236 });
237
238 if let Some(existing) = self.context_set.get(AgentContextKey::ref_cast(&context)) {
239 if remove_if_exists {
240 self.remove_context(&context, cx);
241 None
242 } else {
243 Some(existing.as_ref().clone())
244 }
245 } else {
246 self.insert_context(context.clone(), cx);
247 Some(context)
248 }
249 }
250
251 pub fn add_fetched_url(
252 &mut self,
253 url: String,
254 text: impl Into<SharedString>,
255 cx: &mut Context<ContextStore>,
256 ) -> AgentContextHandle {
257 let context = AgentContextHandle::FetchedUrl(FetchedUrlContext {
258 url: url.into(),
259 text: text.into(),
260 context_id: self.next_context_id.post_inc(),
261 });
262
263 self.insert_context(context.clone(), cx);
264 context
265 }
266
267 pub fn add_image_from_path(
268 &mut self,
269 project_path: ProjectPath,
270 remove_if_exists: bool,
271 cx: &mut Context<ContextStore>,
272 ) -> Task<Result<Option<AgentContextHandle>>> {
273 let project = self.project.clone();
274 cx.spawn(async move |this, cx| {
275 let open_image_task = project.update(cx, |project, cx| {
276 project.open_image(project_path.clone(), cx)
277 })?;
278 let image_item = open_image_task.await?;
279
280 this.update(cx, |this, cx| {
281 let item = image_item.read(cx);
282 this.insert_image(
283 Some(item.project_path(cx)),
284 Some(item.file.full_path(cx).to_string_lossy().into_owned()),
285 item.image.clone(),
286 remove_if_exists,
287 cx,
288 )
289 })
290 })
291 }
292
293 pub fn add_image_instance(&mut self, image: Arc<Image>, cx: &mut Context<ContextStore>) {
294 self.insert_image(None, None, image, false, cx);
295 }
296
297 fn insert_image(
298 &mut self,
299 project_path: Option<ProjectPath>,
300 full_path: Option<String>,
301 image: Arc<Image>,
302 remove_if_exists: bool,
303 cx: &mut Context<ContextStore>,
304 ) -> Option<AgentContextHandle> {
305 let image_task = LanguageModelImage::from_image(image.clone(), cx).shared();
306 let context = AgentContextHandle::Image(ImageContext {
307 project_path,
308 full_path,
309 original_image: image,
310 image_task,
311 context_id: self.next_context_id.post_inc(),
312 });
313 if self.has_context(&context) && remove_if_exists {
314 self.remove_context(&context, cx);
315 return None;
316 }
317
318 self.insert_context(context.clone(), cx);
319 Some(context)
320 }
321
322 pub fn add_selection(
323 &mut self,
324 buffer: Entity<Buffer>,
325 range: Range<Anchor>,
326 cx: &mut Context<ContextStore>,
327 ) {
328 let context_id = self.next_context_id.post_inc();
329 let context = AgentContextHandle::Selection(SelectionContextHandle {
330 buffer,
331 range,
332 context_id,
333 });
334 self.insert_context(context, cx);
335 }
336
337 pub fn add_suggested_context(
338 &mut self,
339 suggested: &SuggestedContext,
340 cx: &mut Context<ContextStore>,
341 ) {
342 match suggested {
343 SuggestedContext::File {
344 buffer,
345 icon_path: _,
346 name: _,
347 } => {
348 if let Some(buffer) = buffer.upgrade() {
349 let context_id = self.next_context_id.post_inc();
350 self.insert_context(
351 AgentContextHandle::File(FileContextHandle { buffer, context_id }),
352 cx,
353 );
354 };
355 }
356 SuggestedContext::TextThread {
357 text_thread,
358 name: _,
359 } => {
360 if let Some(text_thread) = text_thread.upgrade() {
361 let context_id = self.next_context_id.post_inc();
362 self.insert_context(
363 AgentContextHandle::TextThread(TextThreadContextHandle {
364 text_thread,
365 context_id,
366 }),
367 cx,
368 );
369 }
370 }
371 }
372 }
373
374 fn insert_context(&mut self, context: AgentContextHandle, cx: &mut Context<Self>) -> bool {
375 match &context {
376 // AgentContextHandle::Thread(thread_context) => {
377 // if let Some(thread_store) = self.thread_store.clone() {
378 // thread_context.thread.update(cx, |thread, cx| {
379 // thread.start_generating_detailed_summary_if_needed(thread_store, cx);
380 // });
381 // self.context_thread_ids
382 // .insert(thread_context.thread.read(cx).id().clone());
383 // } else {
384 // return false;
385 // }
386 // }
387 AgentContextHandle::TextThread(text_thread_context) => {
388 self.context_text_thread_paths
389 .extend(text_thread_context.text_thread.read(cx).path().cloned());
390 }
391 _ => {}
392 }
393 let inserted = self.context_set.insert(AgentContextKey(context));
394 if inserted {
395 cx.notify();
396 }
397 inserted
398 }
399
400 pub fn remove_context(&mut self, context: &AgentContextHandle, cx: &mut Context<Self>) {
401 if let Some((_, key)) = self
402 .context_set
403 .shift_remove_full(AgentContextKey::ref_cast(context))
404 {
405 match context {
406 AgentContextHandle::Thread(thread_context) => {
407 self.context_thread_ids
408 .remove(thread_context.thread.read(cx).id());
409 }
410 AgentContextHandle::TextThread(text_thread_context) => {
411 if let Some(path) = text_thread_context.text_thread.read(cx).path() {
412 self.context_text_thread_paths.remove(path);
413 }
414 }
415 _ => {}
416 }
417 cx.emit(ContextStoreEvent::ContextRemoved(key));
418 cx.notify();
419 }
420 }
421
422 pub fn has_context(&mut self, context: &AgentContextHandle) -> bool {
423 self.context_set
424 .contains(AgentContextKey::ref_cast(context))
425 }
426
427 /// Returns whether this file path is already included directly in the context, or if it will be
428 /// included in the context via a directory.
429 pub fn file_path_included(&self, path: &ProjectPath, cx: &App) -> Option<FileInclusion> {
430 let project = self.project.upgrade()?.read(cx);
431 self.context().find_map(|context| match context {
432 AgentContextHandle::File(file_context) => {
433 FileInclusion::check_file(file_context, path, cx)
434 }
435 AgentContextHandle::Image(image_context) => {
436 FileInclusion::check_image(image_context, path)
437 }
438 AgentContextHandle::Directory(directory_context) => {
439 FileInclusion::check_directory(directory_context, path, project, cx)
440 }
441 _ => None,
442 })
443 }
444
445 pub fn path_included_in_directory(
446 &self,
447 path: &ProjectPath,
448 cx: &App,
449 ) -> Option<FileInclusion> {
450 let project = self.project.upgrade()?.read(cx);
451 self.context().find_map(|context| match context {
452 AgentContextHandle::Directory(directory_context) => {
453 FileInclusion::check_directory(directory_context, path, project, cx)
454 }
455 _ => None,
456 })
457 }
458
459 pub fn includes_symbol(&self, symbol: &Symbol, cx: &App) -> bool {
460 self.context().any(|context| match context {
461 AgentContextHandle::Symbol(context) => {
462 if context.symbol != symbol.name {
463 return false;
464 }
465 let buffer = context.buffer.read(cx);
466 let Some(context_path) = buffer.project_path(cx) else {
467 return false;
468 };
469 if symbol.path != SymbolLocation::InProject(context_path) {
470 return false;
471 }
472 let context_range = context.range.to_point_utf16(&buffer.snapshot());
473 context_range.start == symbol.range.start.0
474 && context_range.end == symbol.range.end.0
475 }
476 _ => false,
477 })
478 }
479
480 pub fn includes_thread(&self, thread_id: &acp::SessionId) -> bool {
481 self.context_thread_ids.contains(thread_id)
482 }
483
484 pub fn includes_text_thread(&self, path: &Arc<Path>) -> bool {
485 self.context_text_thread_paths.contains(path)
486 }
487
488 pub fn includes_user_rules(&self, prompt_id: UserPromptId) -> bool {
489 self.context_set
490 .contains(&RulesContextHandle::lookup_key(prompt_id))
491 }
492
493 pub fn includes_url(&self, url: impl Into<SharedString>) -> bool {
494 self.context_set
495 .contains(&FetchedUrlContext::lookup_key(url.into()))
496 }
497
498 pub fn get_url_context(&self, url: SharedString) -> Option<AgentContextHandle> {
499 self.context_set
500 .get(&FetchedUrlContext::lookup_key(url))
501 .map(|key| key.as_ref().clone())
502 }
503
504 pub fn file_paths(&self, cx: &App) -> HashSet<ProjectPath> {
505 self.context()
506 .filter_map(|context| match context {
507 AgentContextHandle::File(file) => {
508 let buffer = file.buffer.read(cx);
509 buffer.project_path(cx)
510 }
511 AgentContextHandle::Directory(_)
512 | AgentContextHandle::Symbol(_)
513 | AgentContextHandle::Thread(_)
514 | AgentContextHandle::Selection(_)
515 | AgentContextHandle::FetchedUrl(_)
516 | AgentContextHandle::TextThread(_)
517 | AgentContextHandle::Rules(_)
518 | AgentContextHandle::Image(_) => None,
519 })
520 .collect()
521 }
522
523 pub fn thread_ids(&self) -> &HashSet<acp::SessionId> {
524 &self.context_thread_ids
525 }
526}
527
528#[derive(Clone)]
529pub enum SuggestedContext {
530 File {
531 name: SharedString,
532 icon_path: Option<SharedString>,
533 buffer: WeakEntity<Buffer>,
534 },
535 TextThread {
536 name: SharedString,
537 text_thread: WeakEntity<TextThread>,
538 },
539}
540
541impl SuggestedContext {
542 pub fn name(&self) -> &SharedString {
543 match self {
544 Self::File { name, .. } => name,
545 Self::TextThread { name, .. } => name,
546 }
547 }
548
549 pub fn icon_path(&self) -> Option<SharedString> {
550 match self {
551 Self::File { icon_path, .. } => icon_path.clone(),
552 Self::TextThread { .. } => None,
553 }
554 }
555
556 pub fn kind(&self) -> ContextKind {
557 match self {
558 Self::File { .. } => ContextKind::File,
559 Self::TextThread { .. } => ContextKind::TextThread,
560 }
561 }
562}
563
564pub enum FileInclusion {
565 Direct,
566 InDirectory { full_path: PathBuf },
567}
568
569impl FileInclusion {
570 fn check_file(file_context: &FileContextHandle, path: &ProjectPath, cx: &App) -> Option<Self> {
571 let file_path = file_context.buffer.read(cx).project_path(cx)?;
572 if path == &file_path {
573 Some(FileInclusion::Direct)
574 } else {
575 None
576 }
577 }
578
579 fn check_image(image_context: &ImageContext, path: &ProjectPath) -> Option<Self> {
580 let image_path = image_context.project_path.as_ref()?;
581 if path == image_path {
582 Some(FileInclusion::Direct)
583 } else {
584 None
585 }
586 }
587
588 fn check_directory(
589 directory_context: &DirectoryContextHandle,
590 path: &ProjectPath,
591 project: &Project,
592 cx: &App,
593 ) -> Option<Self> {
594 let worktree = project
595 .worktree_for_entry(directory_context.entry_id, cx)?
596 .read(cx);
597 let entry = worktree.entry_for_id(directory_context.entry_id)?;
598 let directory_path = ProjectPath {
599 worktree_id: worktree.id(),
600 path: entry.path.clone(),
601 };
602 if path.starts_with(&directory_path) {
603 if path == &directory_path {
604 Some(FileInclusion::Direct)
605 } else {
606 Some(FileInclusion::InDirectory {
607 full_path: worktree.full_path(&entry.path),
608 })
609 }
610 } else {
611 None
612 }
613 }
614}