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