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