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