context.rs

  1use std::fmt::{self, Display, Formatter, Write as _};
  2use std::hash::{Hash, Hasher};
  3use std::path::PathBuf;
  4use std::{ops::Range, path::Path, sync::Arc};
  5
  6use collections::HashSet;
  7use futures::future;
  8use futures::{FutureExt, future::Shared};
  9use gpui::{App, AppContext as _, Entity, SharedString, Task};
 10use language::Buffer;
 11use language_model::{LanguageModelImage, LanguageModelRequestMessage, MessageContent};
 12use project::{Project, ProjectEntryId, ProjectPath, Worktree};
 13use prompt_store::{PromptStore, UserPromptId};
 14use ref_cast::RefCast;
 15use rope::Point;
 16use text::{Anchor, OffsetRangeExt as _};
 17use ui::{ElementId, IconName};
 18use util::markdown::MarkdownCodeBlock;
 19use util::{ResultExt as _, post_inc};
 20
 21use crate::thread::Thread;
 22
 23pub const RULES_ICON: IconName = IconName::Context;
 24
 25pub enum ContextKind {
 26    File,
 27    Directory,
 28    Symbol,
 29    Selection,
 30    FetchedUrl,
 31    Thread,
 32    Rules,
 33    Image,
 34}
 35
 36impl ContextKind {
 37    pub fn icon(&self) -> IconName {
 38        match self {
 39            ContextKind::File => IconName::File,
 40            ContextKind::Directory => IconName::Folder,
 41            ContextKind::Symbol => IconName::Code,
 42            ContextKind::Selection => IconName::Context,
 43            ContextKind::FetchedUrl => IconName::Globe,
 44            ContextKind::Thread => IconName::MessageBubbles,
 45            ContextKind::Rules => RULES_ICON,
 46            ContextKind::Image => IconName::Image,
 47        }
 48    }
 49}
 50
 51/// Handle for context that can be attached to a user message.
 52///
 53/// This uses IDs that are stable enough for tracking renames and identifying when context has
 54/// already been added to the thread. To use this in a set, wrap it in `AgentContextKey` to opt in
 55/// to `PartialEq` and `Hash` impls that use the subset of the fields used for this stable identity.
 56#[derive(Debug, Clone)]
 57pub enum AgentContextHandle {
 58    File(FileContextHandle),
 59    Directory(DirectoryContextHandle),
 60    Symbol(SymbolContextHandle),
 61    Selection(SelectionContextHandle),
 62    FetchedUrl(FetchedUrlContext),
 63    Thread(ThreadContextHandle),
 64    Rules(RulesContextHandle),
 65    Image(ImageContext),
 66}
 67
 68impl AgentContextHandle {
 69    fn id(&self) -> ContextId {
 70        match self {
 71            Self::File(context) => context.context_id,
 72            Self::Directory(context) => context.context_id,
 73            Self::Symbol(context) => context.context_id,
 74            Self::Selection(context) => context.context_id,
 75            Self::FetchedUrl(context) => context.context_id,
 76            Self::Thread(context) => context.context_id,
 77            Self::Rules(context) => context.context_id,
 78            Self::Image(context) => context.context_id,
 79        }
 80    }
 81
 82    pub fn element_id(&self, name: SharedString) -> ElementId {
 83        ElementId::NamedInteger(name, self.id().0)
 84    }
 85}
 86
 87/// Loaded context that can be attached to a user message. This can be thought of as a
 88/// snapshot of the context along with an `AgentContextHandle`.
 89#[derive(Debug, Clone)]
 90pub enum AgentContext {
 91    File(FileContext),
 92    Directory(DirectoryContext),
 93    Symbol(SymbolContext),
 94    Selection(SelectionContext),
 95    FetchedUrl(FetchedUrlContext),
 96    Thread(ThreadContext),
 97    Rules(RulesContext),
 98    Image(ImageContext),
 99}
100
101impl AgentContext {
102    pub fn handle(&self) -> AgentContextHandle {
103        match self {
104            AgentContext::File(context) => AgentContextHandle::File(context.handle.clone()),
105            AgentContext::Directory(context) => {
106                AgentContextHandle::Directory(context.handle.clone())
107            }
108            AgentContext::Symbol(context) => AgentContextHandle::Symbol(context.handle.clone()),
109            AgentContext::Selection(context) => {
110                AgentContextHandle::Selection(context.handle.clone())
111            }
112            AgentContext::FetchedUrl(context) => AgentContextHandle::FetchedUrl(context.clone()),
113            AgentContext::Thread(context) => AgentContextHandle::Thread(context.handle.clone()),
114            AgentContext::Rules(context) => AgentContextHandle::Rules(context.handle.clone()),
115            AgentContext::Image(context) => AgentContextHandle::Image(context.clone()),
116        }
117    }
118}
119
120/// ID created at time of context add, for use in ElementId. This is not the stable identity of a
121/// context, instead that's handled by the `PartialEq` and `Hash` impls of `AgentContextKey`.
122#[derive(Debug, Copy, Clone)]
123pub struct ContextId(u64);
124
125impl ContextId {
126    pub fn zero() -> Self {
127        ContextId(0)
128    }
129
130    fn for_lookup() -> Self {
131        ContextId(u64::MAX)
132    }
133
134    pub fn post_inc(&mut self) -> Self {
135        Self(post_inc(&mut self.0))
136    }
137}
138
139/// File context provides the entire contents of a file.
140///
141/// This holds an `Entity<Buffer>` so that file path renames affect its display and so that it can
142/// be opened even if the file has been deleted. An alternative might be to use `ProjectEntryId`,
143/// but then when deleted there is no path info or ability to open.
144#[derive(Debug, Clone)]
145pub struct FileContextHandle {
146    pub buffer: Entity<Buffer>,
147    pub context_id: ContextId,
148}
149
150#[derive(Debug, Clone)]
151pub struct FileContext {
152    pub handle: FileContextHandle,
153    pub full_path: Arc<Path>,
154    pub text: SharedString,
155}
156
157impl FileContextHandle {
158    pub fn eq_for_key(&self, other: &Self) -> bool {
159        self.buffer == other.buffer
160    }
161
162    pub fn hash_for_key<H: Hasher>(&self, state: &mut H) {
163        self.buffer.hash(state)
164    }
165
166    pub fn project_path(&self, cx: &App) -> Option<ProjectPath> {
167        let file = self.buffer.read(cx).file()?;
168        Some(ProjectPath {
169            worktree_id: file.worktree_id(cx),
170            path: file.path().clone(),
171        })
172    }
173
174    fn load(self, cx: &App) -> Task<Option<(AgentContext, Vec<Entity<Buffer>>)>> {
175        let buffer_ref = self.buffer.read(cx);
176        let Some(file) = buffer_ref.file() else {
177            log::error!("file context missing path");
178            return Task::ready(None);
179        };
180        let full_path = file.full_path(cx);
181        let rope = buffer_ref.as_rope().clone();
182        let buffer = self.buffer.clone();
183        cx.background_spawn(async move {
184            let context = AgentContext::File(FileContext {
185                handle: self,
186                full_path: full_path.into(),
187                text: rope.to_string().into(),
188            });
189            Some((context, vec![buffer]))
190        })
191    }
192}
193
194impl Display for FileContext {
195    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
196        write!(
197            f,
198            "{}",
199            MarkdownCodeBlock {
200                tag: &codeblock_tag(&self.full_path, None),
201                text: &self.text,
202            }
203        )
204    }
205}
206
207/// Directory contents provides the entire contents of text files in a directory.
208///
209/// This has a `ProjectEntryId` so that it follows renames.
210#[derive(Debug, Clone)]
211pub struct DirectoryContextHandle {
212    pub entry_id: ProjectEntryId,
213    pub context_id: ContextId,
214}
215
216#[derive(Debug, Clone)]
217pub struct DirectoryContext {
218    pub handle: DirectoryContextHandle,
219    pub full_path: Arc<Path>,
220    pub descendants: Vec<DirectoryContextDescendant>,
221}
222
223#[derive(Debug, Clone)]
224pub struct DirectoryContextDescendant {
225    /// Path within the directory.
226    pub rel_path: Arc<Path>,
227    pub fenced_codeblock: SharedString,
228}
229
230impl DirectoryContextHandle {
231    pub fn eq_for_key(&self, other: &Self) -> bool {
232        self.entry_id == other.entry_id
233    }
234
235    pub fn hash_for_key<H: Hasher>(&self, state: &mut H) {
236        self.entry_id.hash(state)
237    }
238
239    fn load(
240        self,
241        project: Entity<Project>,
242        cx: &mut App,
243    ) -> Task<Option<(AgentContext, Vec<Entity<Buffer>>)>> {
244        let Some(worktree) = project.read(cx).worktree_for_entry(self.entry_id, cx) else {
245            return Task::ready(None);
246        };
247        let worktree_ref = worktree.read(cx);
248        let Some(entry) = worktree_ref.entry_for_id(self.entry_id) else {
249            return Task::ready(None);
250        };
251        if entry.is_file() {
252            log::error!("DirectoryContext unexpectedly refers to a file.");
253            return Task::ready(None);
254        }
255
256        let directory_path = entry.path.clone();
257        let directory_full_path = worktree_ref.full_path(&directory_path).into();
258
259        let file_paths = collect_files_in_path(worktree_ref, &directory_path);
260        let descendants_future = future::join_all(file_paths.into_iter().map(|path| {
261            let worktree_ref = worktree.read(cx);
262            let worktree_id = worktree_ref.id();
263            let full_path = worktree_ref.full_path(&path);
264
265            let rel_path = path
266                .strip_prefix(&directory_path)
267                .log_err()
268                .map_or_else(|| path.clone(), |rel_path| rel_path.into());
269
270            let open_task = project.update(cx, |project, cx| {
271                project.buffer_store().update(cx, |buffer_store, cx| {
272                    let project_path = ProjectPath { worktree_id, path };
273                    buffer_store.open_buffer(project_path, cx)
274                })
275            });
276
277            // TODO: report load errors instead of just logging
278            let rope_task = cx.spawn(async move |cx| {
279                let buffer = open_task.await.log_err()?;
280                let rope = buffer
281                    .read_with(cx, |buffer, _cx| buffer.as_rope().clone())
282                    .log_err()?;
283                Some((rope, buffer))
284            });
285
286            cx.background_spawn(async move {
287                let (rope, buffer) = rope_task.await?;
288                let fenced_codeblock = MarkdownCodeBlock {
289                    tag: &codeblock_tag(&full_path, None),
290                    text: &rope.to_string(),
291                }
292                .to_string()
293                .into();
294                let descendant = DirectoryContextDescendant {
295                    rel_path,
296                    fenced_codeblock,
297                };
298                Some((descendant, buffer))
299            })
300        }));
301
302        cx.background_spawn(async move {
303            let (descendants, buffers) = descendants_future.await.into_iter().flatten().unzip();
304            let context = AgentContext::Directory(DirectoryContext {
305                handle: self,
306                full_path: directory_full_path,
307                descendants,
308            });
309            Some((context, buffers))
310        })
311    }
312}
313
314impl Display for DirectoryContext {
315    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
316        let mut is_first = true;
317        for descendant in &self.descendants {
318            if !is_first {
319                write!(f, "\n")?;
320            } else {
321                is_first = false;
322            }
323            write!(f, "{}", descendant.fenced_codeblock)?;
324        }
325        Ok(())
326    }
327}
328
329#[derive(Debug, Clone)]
330pub struct SymbolContextHandle {
331    pub buffer: Entity<Buffer>,
332    pub symbol: SharedString,
333    pub range: Range<Anchor>,
334    /// The range that fully contains the symbol. e.g. for function symbol, this will include not
335    /// only the signature, but also the body. Not used by `PartialEq` or `Hash` for
336    /// `AgentContextKey`.
337    pub enclosing_range: Range<Anchor>,
338    pub context_id: ContextId,
339}
340
341#[derive(Debug, Clone)]
342pub struct SymbolContext {
343    pub handle: SymbolContextHandle,
344    pub full_path: Arc<Path>,
345    pub line_range: Range<Point>,
346    pub text: SharedString,
347}
348
349impl SymbolContextHandle {
350    pub fn eq_for_key(&self, other: &Self) -> bool {
351        self.buffer == other.buffer && self.symbol == other.symbol && self.range == other.range
352    }
353
354    pub fn hash_for_key<H: Hasher>(&self, state: &mut H) {
355        self.buffer.hash(state);
356        self.symbol.hash(state);
357        self.range.hash(state);
358    }
359
360    pub fn full_path(&self, cx: &App) -> Option<PathBuf> {
361        Some(self.buffer.read(cx).file()?.full_path(cx))
362    }
363
364    pub fn enclosing_line_range(&self, cx: &App) -> Range<Point> {
365        self.enclosing_range
366            .to_point(&self.buffer.read(cx).snapshot())
367    }
368
369    pub fn text(&self, cx: &App) -> SharedString {
370        self.buffer
371            .read(cx)
372            .text_for_range(self.enclosing_range.clone())
373            .collect::<String>()
374            .into()
375    }
376
377    fn load(self, cx: &App) -> Task<Option<(AgentContext, Vec<Entity<Buffer>>)>> {
378        let buffer_ref = self.buffer.read(cx);
379        let Some(file) = buffer_ref.file() else {
380            log::error!("symbol context's file has no path");
381            return Task::ready(None);
382        };
383        let full_path = file.full_path(cx).into();
384        let line_range = self.enclosing_range.to_point(&buffer_ref.snapshot());
385        let text = self.text(cx);
386        let buffer = self.buffer.clone();
387        let context = AgentContext::Symbol(SymbolContext {
388            handle: self,
389            full_path,
390            line_range,
391            text,
392        });
393        Task::ready(Some((context, vec![buffer])))
394    }
395}
396
397impl Display for SymbolContext {
398    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
399        let code_block = MarkdownCodeBlock {
400            tag: &codeblock_tag(&self.full_path, Some(self.line_range.clone())),
401            text: &self.text,
402        };
403        write!(f, "{code_block}",)
404    }
405}
406
407#[derive(Debug, Clone)]
408pub struct SelectionContextHandle {
409    pub buffer: Entity<Buffer>,
410    pub range: Range<Anchor>,
411    pub context_id: ContextId,
412}
413
414#[derive(Debug, Clone)]
415pub struct SelectionContext {
416    pub handle: SelectionContextHandle,
417    pub full_path: Arc<Path>,
418    pub line_range: Range<Point>,
419    pub text: SharedString,
420}
421
422impl SelectionContextHandle {
423    pub fn eq_for_key(&self, other: &Self) -> bool {
424        self.buffer == other.buffer && self.range == other.range
425    }
426
427    pub fn hash_for_key<H: Hasher>(&self, state: &mut H) {
428        self.buffer.hash(state);
429        self.range.hash(state);
430    }
431
432    pub fn full_path(&self, cx: &App) -> Option<PathBuf> {
433        Some(self.buffer.read(cx).file()?.full_path(cx))
434    }
435
436    pub fn line_range(&self, cx: &App) -> Range<Point> {
437        self.range.to_point(&self.buffer.read(cx).snapshot())
438    }
439
440    pub fn text(&self, cx: &App) -> SharedString {
441        self.buffer
442            .read(cx)
443            .text_for_range(self.range.clone())
444            .collect::<String>()
445            .into()
446    }
447
448    fn load(self, cx: &App) -> Task<Option<(AgentContext, Vec<Entity<Buffer>>)>> {
449        let Some(full_path) = self.full_path(cx) else {
450            log::error!("selection context's file has no path");
451            return Task::ready(None);
452        };
453        let text = self.text(cx);
454        let buffer = self.buffer.clone();
455        let context = AgentContext::Selection(SelectionContext {
456            full_path: full_path.into(),
457            line_range: self.line_range(cx),
458            text,
459            handle: self,
460        });
461
462        Task::ready(Some((context, vec![buffer])))
463    }
464}
465
466impl Display for SelectionContext {
467    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
468        let code_block = MarkdownCodeBlock {
469            tag: &codeblock_tag(&self.full_path, Some(self.line_range.clone())),
470            text: &self.text,
471        };
472        write!(f, "{code_block}",)
473    }
474}
475
476#[derive(Debug, Clone)]
477pub struct FetchedUrlContext {
478    pub url: SharedString,
479    /// Text contents of the fetched url. Unlike other context types, the contents of this gets
480    /// populated when added rather than when sending the message. Not used by `PartialEq` or `Hash`
481    /// for `AgentContextKey`.
482    pub text: SharedString,
483    pub context_id: ContextId,
484}
485
486impl FetchedUrlContext {
487    pub fn eq_for_key(&self, other: &Self) -> bool {
488        self.url == other.url
489    }
490
491    pub fn hash_for_key<H: Hasher>(&self, state: &mut H) {
492        self.url.hash(state);
493    }
494
495    pub fn lookup_key(url: SharedString) -> AgentContextKey {
496        AgentContextKey(AgentContextHandle::FetchedUrl(FetchedUrlContext {
497            url,
498            text: "".into(),
499            context_id: ContextId::for_lookup(),
500        }))
501    }
502
503    pub fn load(self) -> Task<Option<(AgentContext, Vec<Entity<Buffer>>)>> {
504        Task::ready(Some((AgentContext::FetchedUrl(self), vec![])))
505    }
506}
507
508impl Display for FetchedUrlContext {
509    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
510        // TODO: Better format - url and contents are not delimited.
511        write!(f, "{}\n{}\n", self.url, self.text)
512    }
513}
514
515#[derive(Debug, Clone)]
516pub struct ThreadContextHandle {
517    pub thread: Entity<Thread>,
518    pub context_id: ContextId,
519}
520
521#[derive(Debug, Clone)]
522pub struct ThreadContext {
523    pub handle: ThreadContextHandle,
524    pub title: SharedString,
525    pub text: SharedString,
526}
527
528impl ThreadContextHandle {
529    pub fn eq_for_key(&self, other: &Self) -> bool {
530        self.thread == other.thread
531    }
532
533    pub fn hash_for_key<H: Hasher>(&self, state: &mut H) {
534        self.thread.hash(state)
535    }
536
537    pub fn title(&self, cx: &App) -> SharedString {
538        self.thread
539            .read(cx)
540            .summary()
541            .unwrap_or_else(|| "New thread".into())
542    }
543
544    fn load(self, cx: &App) -> Task<Option<(AgentContext, Vec<Entity<Buffer>>)>> {
545        let context = AgentContext::Thread(ThreadContext {
546            title: self.title(cx),
547            text: self.thread.read(cx).latest_detailed_summary_or_text(),
548            handle: self,
549        });
550        Task::ready(Some((context, vec![])))
551    }
552}
553
554impl Display for ThreadContext {
555    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
556        // TODO: Better format for this - doesn't distinguish title and contents.
557        write!(f, "{}\n{}\n", &self.title, &self.text.trim())
558    }
559}
560
561#[derive(Debug, Clone)]
562pub struct RulesContextHandle {
563    pub prompt_id: UserPromptId,
564    pub context_id: ContextId,
565}
566
567#[derive(Debug, Clone)]
568pub struct RulesContext {
569    pub handle: RulesContextHandle,
570    pub title: Option<SharedString>,
571    pub text: SharedString,
572}
573
574impl RulesContextHandle {
575    pub fn eq_for_key(&self, other: &Self) -> bool {
576        self.prompt_id == other.prompt_id
577    }
578
579    pub fn hash_for_key<H: Hasher>(&self, state: &mut H) {
580        self.prompt_id.hash(state)
581    }
582
583    pub fn lookup_key(prompt_id: UserPromptId) -> AgentContextKey {
584        AgentContextKey(AgentContextHandle::Rules(RulesContextHandle {
585            prompt_id,
586            context_id: ContextId::for_lookup(),
587        }))
588    }
589
590    pub fn load(
591        self,
592        prompt_store: &Option<Entity<PromptStore>>,
593        cx: &App,
594    ) -> Task<Option<(AgentContext, Vec<Entity<Buffer>>)>> {
595        let Some(prompt_store) = prompt_store.as_ref() else {
596            return Task::ready(None);
597        };
598        let prompt_store = prompt_store.read(cx);
599        let prompt_id = self.prompt_id.into();
600        let Some(metadata) = prompt_store.metadata(prompt_id) else {
601            return Task::ready(None);
602        };
603        let title = metadata.title;
604        let text_task = prompt_store.load(prompt_id, cx);
605        cx.background_spawn(async move {
606            // TODO: report load errors instead of just logging
607            let text = text_task.await.log_err()?.into();
608            let context = AgentContext::Rules(RulesContext {
609                handle: self,
610                title,
611                text,
612            });
613            Some((context, vec![]))
614        })
615    }
616}
617
618impl Display for RulesContext {
619    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
620        if let Some(title) = &self.title {
621            write!(f, "Rules title: {}\n", title)?;
622        }
623        let code_block = MarkdownCodeBlock {
624            tag: "",
625            text: self.text.trim(),
626        };
627        write!(f, "{code_block}")
628    }
629}
630
631#[derive(Debug, Clone)]
632pub struct ImageContext {
633    pub project_path: Option<ProjectPath>,
634    pub original_image: Arc<gpui::Image>,
635    // TODO: handle this elsewhere and remove `ignore-interior-mutability` opt-out in clippy.toml
636    // needed due to a false positive of `clippy::mutable_key_type`.
637    pub image_task: Shared<Task<Option<LanguageModelImage>>>,
638    pub context_id: ContextId,
639}
640
641pub enum ImageStatus {
642    Loading,
643    Error,
644    Ready,
645}
646
647impl ImageContext {
648    pub fn eq_for_key(&self, other: &Self) -> bool {
649        self.original_image.id == other.original_image.id
650    }
651
652    pub fn hash_for_key<H: Hasher>(&self, state: &mut H) {
653        self.original_image.id.hash(state);
654    }
655
656    pub fn image(&self) -> Option<LanguageModelImage> {
657        self.image_task.clone().now_or_never().flatten()
658    }
659
660    pub fn status(&self) -> ImageStatus {
661        match self.image_task.clone().now_or_never() {
662            None => ImageStatus::Loading,
663            Some(None) => ImageStatus::Error,
664            Some(Some(_)) => ImageStatus::Ready,
665        }
666    }
667
668    pub fn load(self, cx: &App) -> Task<Option<(AgentContext, Vec<Entity<Buffer>>)>> {
669        cx.background_spawn(async move {
670            self.image_task.clone().await;
671            Some((AgentContext::Image(self), vec![]))
672        })
673    }
674}
675
676#[derive(Debug, Clone, Default)]
677pub struct ContextLoadResult {
678    pub loaded_context: LoadedContext,
679    pub referenced_buffers: HashSet<Entity<Buffer>>,
680}
681
682#[derive(Debug, Clone, Default)]
683pub struct LoadedContext {
684    pub contexts: Vec<AgentContext>,
685    pub text: String,
686    pub images: Vec<LanguageModelImage>,
687}
688
689impl LoadedContext {
690    pub fn is_empty(&self) -> bool {
691        self.text.is_empty() && self.images.is_empty()
692    }
693
694    pub fn add_to_request_message(&self, request_message: &mut LanguageModelRequestMessage) {
695        if !self.text.is_empty() {
696            request_message
697                .content
698                .push(MessageContent::Text(self.text.to_string()));
699        }
700
701        if !self.images.is_empty() {
702            // Some providers only support image parts after an initial text part
703            if request_message.content.is_empty() {
704                request_message
705                    .content
706                    .push(MessageContent::Text("Images attached by user:".to_string()));
707            }
708
709            for image in &self.images {
710                request_message
711                    .content
712                    .push(MessageContent::Image(image.clone()))
713            }
714        }
715    }
716}
717
718/// Loads and formats a collection of contexts.
719pub fn load_context(
720    contexts: Vec<AgentContextHandle>,
721    project: &Entity<Project>,
722    prompt_store: &Option<Entity<PromptStore>>,
723    cx: &mut App,
724) -> Task<ContextLoadResult> {
725    let mut load_tasks = Vec::new();
726
727    for context in contexts.iter().cloned() {
728        match context {
729            AgentContextHandle::File(context) => load_tasks.push(context.load(cx)),
730            AgentContextHandle::Directory(context) => {
731                load_tasks.push(context.load(project.clone(), cx))
732            }
733            AgentContextHandle::Symbol(context) => load_tasks.push(context.load(cx)),
734            AgentContextHandle::Selection(context) => load_tasks.push(context.load(cx)),
735            AgentContextHandle::FetchedUrl(context) => load_tasks.push(context.load()),
736            AgentContextHandle::Thread(context) => load_tasks.push(context.load(cx)),
737            AgentContextHandle::Rules(context) => load_tasks.push(context.load(prompt_store, cx)),
738            AgentContextHandle::Image(context) => load_tasks.push(context.load(cx)),
739        }
740    }
741
742    cx.background_spawn(async move {
743        let load_results = future::join_all(load_tasks).await;
744
745        let mut contexts = Vec::new();
746        let mut text = String::new();
747        let mut referenced_buffers = HashSet::default();
748        for context in load_results {
749            let Some((context, buffers)) = context else {
750                continue;
751            };
752            contexts.push(context);
753            referenced_buffers.extend(buffers);
754        }
755
756        let mut file_context = Vec::new();
757        let mut directory_context = Vec::new();
758        let mut symbol_context = Vec::new();
759        let mut selection_context = Vec::new();
760        let mut fetched_url_context = Vec::new();
761        let mut thread_context = Vec::new();
762        let mut rules_context = Vec::new();
763        let mut images = Vec::new();
764        for context in &contexts {
765            match context {
766                AgentContext::File(context) => file_context.push(context),
767                AgentContext::Directory(context) => directory_context.push(context),
768                AgentContext::Symbol(context) => symbol_context.push(context),
769                AgentContext::Selection(context) => selection_context.push(context),
770                AgentContext::FetchedUrl(context) => fetched_url_context.push(context),
771                AgentContext::Thread(context) => thread_context.push(context),
772                AgentContext::Rules(context) => rules_context.push(context),
773                AgentContext::Image(context) => images.extend(context.image()),
774            }
775        }
776
777        if file_context.is_empty()
778            && directory_context.is_empty()
779            && symbol_context.is_empty()
780            && selection_context.is_empty()
781            && fetched_url_context.is_empty()
782            && thread_context.is_empty()
783            && rules_context.is_empty()
784        {
785            return ContextLoadResult {
786                loaded_context: LoadedContext {
787                    contexts,
788                    text,
789                    images,
790                },
791                referenced_buffers,
792            };
793        }
794
795        text.push_str(
796            "\n<context>\n\
797            The following items were attached by the user. \
798            You don't need to use other tools to read them.\n\n",
799        );
800
801        if !file_context.is_empty() {
802            text.push_str("<files>");
803            for context in file_context {
804                text.push('\n');
805                let _ = write!(text, "{context}");
806            }
807            text.push_str("</files>\n");
808        }
809
810        if !directory_context.is_empty() {
811            text.push_str("<directories>");
812            for context in directory_context {
813                text.push('\n');
814                let _ = write!(text, "{context}");
815            }
816            text.push_str("</directories>\n");
817        }
818
819        if !symbol_context.is_empty() {
820            text.push_str("<symbols>");
821            for context in symbol_context {
822                text.push('\n');
823                let _ = write!(text, "{context}");
824            }
825            text.push_str("</symbols>\n");
826        }
827
828        if !selection_context.is_empty() {
829            text.push_str("<selections>");
830            for context in selection_context {
831                text.push('\n');
832                let _ = write!(text, "{context}");
833            }
834            text.push_str("</selections>\n");
835        }
836
837        if !fetched_url_context.is_empty() {
838            text.push_str("<fetched_urls>");
839            for context in fetched_url_context {
840                text.push('\n');
841                let _ = write!(text, "{context}");
842            }
843            text.push_str("</fetched_urls>\n");
844        }
845
846        if !thread_context.is_empty() {
847            text.push_str("<conversation_threads>");
848            for context in thread_context {
849                text.push('\n');
850                let _ = write!(text, "{context}");
851            }
852            text.push_str("</conversation_threads>\n");
853        }
854
855        if !rules_context.is_empty() {
856            text.push_str(
857                "<user_rules>\n\
858                The user has specified the following rules that should be applied:\n",
859            );
860            for context in rules_context {
861                text.push('\n');
862                let _ = write!(text, "{context}");
863            }
864            text.push_str("</user_rules>\n");
865        }
866
867        text.push_str("</context>\n");
868
869        ContextLoadResult {
870            loaded_context: LoadedContext {
871                contexts,
872                text,
873                images,
874            },
875            referenced_buffers,
876        }
877    })
878}
879
880fn collect_files_in_path(worktree: &Worktree, path: &Path) -> Vec<Arc<Path>> {
881    let mut files = Vec::new();
882
883    for entry in worktree.child_entries(path) {
884        if entry.is_dir() {
885            files.extend(collect_files_in_path(worktree, &entry.path));
886        } else if entry.is_file() {
887            files.push(entry.path.clone());
888        }
889    }
890
891    files
892}
893
894fn codeblock_tag(full_path: &Path, line_range: Option<Range<Point>>) -> String {
895    let mut result = String::new();
896
897    if let Some(extension) = full_path.extension().and_then(|ext| ext.to_str()) {
898        let _ = write!(result, "{} ", extension);
899    }
900
901    let _ = write!(result, "{}", full_path.display());
902
903    if let Some(range) = line_range {
904        if range.start.row == range.end.row {
905            let _ = write!(result, ":{}", range.start.row + 1);
906        } else {
907            let _ = write!(result, ":{}-{}", range.start.row + 1, range.end.row + 1);
908        }
909    }
910
911    result
912}
913
914/// Wraps `AgentContext` to opt-in to `PartialEq` and `Hash` impls which use a subset of fields
915/// needed for stable context identity.
916#[derive(Debug, Clone, RefCast)]
917#[repr(transparent)]
918pub struct AgentContextKey(pub AgentContextHandle);
919
920impl AsRef<AgentContextHandle> for AgentContextKey {
921    fn as_ref(&self) -> &AgentContextHandle {
922        &self.0
923    }
924}
925
926impl Eq for AgentContextKey {}
927
928impl PartialEq for AgentContextKey {
929    fn eq(&self, other: &Self) -> bool {
930        match &self.0 {
931            AgentContextHandle::File(context) => {
932                if let AgentContextHandle::File(other_context) = &other.0 {
933                    return context.eq_for_key(other_context);
934                }
935            }
936            AgentContextHandle::Directory(context) => {
937                if let AgentContextHandle::Directory(other_context) = &other.0 {
938                    return context.eq_for_key(other_context);
939                }
940            }
941            AgentContextHandle::Symbol(context) => {
942                if let AgentContextHandle::Symbol(other_context) = &other.0 {
943                    return context.eq_for_key(other_context);
944                }
945            }
946            AgentContextHandle::Selection(context) => {
947                if let AgentContextHandle::Selection(other_context) = &other.0 {
948                    return context.eq_for_key(other_context);
949                }
950            }
951            AgentContextHandle::FetchedUrl(context) => {
952                if let AgentContextHandle::FetchedUrl(other_context) = &other.0 {
953                    return context.eq_for_key(other_context);
954                }
955            }
956            AgentContextHandle::Thread(context) => {
957                if let AgentContextHandle::Thread(other_context) = &other.0 {
958                    return context.eq_for_key(other_context);
959                }
960            }
961            AgentContextHandle::Rules(context) => {
962                if let AgentContextHandle::Rules(other_context) = &other.0 {
963                    return context.eq_for_key(other_context);
964                }
965            }
966            AgentContextHandle::Image(context) => {
967                if let AgentContextHandle::Image(other_context) = &other.0 {
968                    return context.eq_for_key(other_context);
969                }
970            }
971        }
972        false
973    }
974}
975
976impl Hash for AgentContextKey {
977    fn hash<H: Hasher>(&self, state: &mut H) {
978        match &self.0 {
979            AgentContextHandle::File(context) => context.hash_for_key(state),
980            AgentContextHandle::Directory(context) => context.hash_for_key(state),
981            AgentContextHandle::Symbol(context) => context.hash_for_key(state),
982            AgentContextHandle::Selection(context) => context.hash_for_key(state),
983            AgentContextHandle::FetchedUrl(context) => context.hash_for_key(state),
984            AgentContextHandle::Thread(context) => context.hash_for_key(state),
985            AgentContextHandle::Rules(context) => context.hash_for_key(state),
986            AgentContextHandle::Image(context) => context.hash_for_key(state),
987        }
988    }
989}