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