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: String,
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 = file.full_path(cx).to_string_lossy().into_owned();
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: String,
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
278 .full_path(&directory_path)
279 .to_string_lossy()
280 .to_string();
281
282 let file_paths = collect_files_in_path(worktree_ref, &directory_path);
283 let descendants_future = future::join_all(file_paths.into_iter().map(|path| {
284 let worktree_ref = worktree.read(cx);
285 let worktree_id = worktree_ref.id();
286 let full_path = worktree_ref.full_path(&path).to_string_lossy().into_owned();
287
288 let rel_path = path
289 .strip_prefix(&directory_path)
290 .log_err()
291 .map_or_else(|| path.clone(), |rel_path| rel_path.into());
292
293 let open_task = project.update(cx, |project, cx| {
294 project.buffer_store().update(cx, |buffer_store, cx| {
295 let project_path = ProjectPath { worktree_id, path };
296 buffer_store.open_buffer(project_path, cx)
297 })
298 });
299
300 // TODO: report load errors instead of just logging
301 let rope_task = cx.spawn(async move |cx| {
302 let buffer = open_task.await.log_err()?;
303 let rope = buffer
304 .read_with(cx, |buffer, _cx| buffer.as_rope().clone())
305 .log_err()?;
306 Some((rope, buffer))
307 });
308
309 cx.background_spawn(async move {
310 let (rope, buffer) = rope_task.await?;
311 let fenced_codeblock = MarkdownCodeBlock {
312 tag: &codeblock_tag(&full_path, None),
313 text: &rope.to_string(),
314 }
315 .to_string()
316 .into();
317 let descendant = DirectoryContextDescendant {
318 rel_path,
319 fenced_codeblock,
320 };
321 Some((descendant, buffer))
322 })
323 }));
324
325 cx.background_spawn(async move {
326 let (descendants, buffers) = descendants_future.await.into_iter().flatten().unzip();
327 let context = AgentContext::Directory(DirectoryContext {
328 handle: self,
329 full_path: directory_full_path,
330 descendants,
331 });
332 Some((context, buffers))
333 })
334 }
335}
336
337impl Display for DirectoryContext {
338 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
339 let mut is_first = true;
340 for descendant in &self.descendants {
341 if !is_first {
342 writeln!(f)?;
343 } else {
344 is_first = false;
345 }
346 write!(f, "{}", descendant.fenced_codeblock)?;
347 }
348 Ok(())
349 }
350}
351
352#[derive(Debug, Clone)]
353pub struct SymbolContextHandle {
354 pub buffer: Entity<Buffer>,
355 pub symbol: SharedString,
356 pub range: Range<Anchor>,
357 /// The range that fully contains the symbol. e.g. for function symbol, this will include not
358 /// only the signature, but also the body. Not used by `PartialEq` or `Hash` for
359 /// `AgentContextKey`.
360 pub enclosing_range: Range<Anchor>,
361 pub context_id: ContextId,
362}
363
364#[derive(Debug, Clone)]
365pub struct SymbolContext {
366 pub handle: SymbolContextHandle,
367 pub full_path: String,
368 pub line_range: Range<Point>,
369 pub text: SharedString,
370}
371
372impl SymbolContextHandle {
373 pub fn eq_for_key(&self, other: &Self) -> bool {
374 self.buffer == other.buffer && self.symbol == other.symbol && self.range == other.range
375 }
376
377 pub fn hash_for_key<H: Hasher>(&self, state: &mut H) {
378 self.buffer.hash(state);
379 self.symbol.hash(state);
380 self.range.hash(state);
381 }
382
383 pub fn full_path(&self, cx: &App) -> Option<PathBuf> {
384 Some(self.buffer.read(cx).file()?.full_path(cx))
385 }
386
387 pub fn enclosing_line_range(&self, cx: &App) -> Range<Point> {
388 self.enclosing_range
389 .to_point(&self.buffer.read(cx).snapshot())
390 }
391
392 pub fn text(&self, cx: &App) -> SharedString {
393 self.buffer
394 .read(cx)
395 .text_for_range(self.enclosing_range.clone())
396 .collect::<String>()
397 .into()
398 }
399
400 fn load(self, cx: &App) -> Task<Option<(AgentContext, Vec<Entity<Buffer>>)>> {
401 let buffer_ref = self.buffer.read(cx);
402 let Some(file) = buffer_ref.file() else {
403 log::error!("symbol context's file has no path");
404 return Task::ready(None);
405 };
406 let full_path = file.full_path(cx).to_string_lossy().into_owned();
407 let line_range = self.enclosing_range.to_point(&buffer_ref.snapshot());
408 let text = self.text(cx);
409 let buffer = self.buffer.clone();
410 let context = AgentContext::Symbol(SymbolContext {
411 handle: self,
412 full_path,
413 line_range,
414 text,
415 });
416 Task::ready(Some((context, vec![buffer])))
417 }
418}
419
420impl Display for SymbolContext {
421 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
422 let code_block = MarkdownCodeBlock {
423 tag: &codeblock_tag(&self.full_path, Some(self.line_range.clone())),
424 text: &self.text,
425 };
426 write!(f, "{code_block}",)
427 }
428}
429
430#[derive(Debug, Clone)]
431pub struct SelectionContextHandle {
432 pub buffer: Entity<Buffer>,
433 pub range: Range<Anchor>,
434 pub context_id: ContextId,
435}
436
437#[derive(Debug, Clone)]
438pub struct SelectionContext {
439 pub handle: SelectionContextHandle,
440 pub full_path: String,
441 pub line_range: Range<Point>,
442 pub text: SharedString,
443}
444
445impl SelectionContextHandle {
446 pub fn eq_for_key(&self, other: &Self) -> bool {
447 self.buffer == other.buffer && self.range == other.range
448 }
449
450 pub fn hash_for_key<H: Hasher>(&self, state: &mut H) {
451 self.buffer.hash(state);
452 self.range.hash(state);
453 }
454
455 pub fn full_path(&self, cx: &App) -> Option<PathBuf> {
456 Some(self.buffer.read(cx).file()?.full_path(cx))
457 }
458
459 pub fn line_range(&self, cx: &App) -> Range<Point> {
460 self.range.to_point(&self.buffer.read(cx).snapshot())
461 }
462
463 pub fn text(&self, cx: &App) -> SharedString {
464 self.buffer
465 .read(cx)
466 .text_for_range(self.range.clone())
467 .collect::<String>()
468 .into()
469 }
470
471 fn load(self, cx: &App) -> Task<Option<(AgentContext, Vec<Entity<Buffer>>)>> {
472 let Some(full_path) = self.full_path(cx) else {
473 log::error!("selection context's file has no path");
474 return Task::ready(None);
475 };
476 let text = self.text(cx);
477 let buffer = self.buffer.clone();
478 let context = AgentContext::Selection(SelectionContext {
479 full_path: full_path.to_string_lossy().into_owned(),
480 line_range: self.line_range(cx),
481 text,
482 handle: self,
483 });
484
485 Task::ready(Some((context, vec![buffer])))
486 }
487}
488
489impl Display for SelectionContext {
490 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
491 let code_block = MarkdownCodeBlock {
492 tag: &codeblock_tag(&self.full_path, Some(self.line_range.clone())),
493 text: &self.text,
494 };
495 write!(f, "{code_block}",)
496 }
497}
498
499#[derive(Debug, Clone)]
500pub struct FetchedUrlContext {
501 pub url: SharedString,
502 /// Text contents of the fetched url. Unlike other context types, the contents of this gets
503 /// populated when added rather than when sending the message. Not used by `PartialEq` or `Hash`
504 /// for `AgentContextKey`.
505 pub text: SharedString,
506 pub context_id: ContextId,
507}
508
509impl FetchedUrlContext {
510 pub fn eq_for_key(&self, other: &Self) -> bool {
511 self.url == other.url
512 }
513
514 pub fn hash_for_key<H: Hasher>(&self, state: &mut H) {
515 self.url.hash(state);
516 }
517
518 pub fn lookup_key(url: SharedString) -> AgentContextKey {
519 AgentContextKey(AgentContextHandle::FetchedUrl(FetchedUrlContext {
520 url,
521 text: "".into(),
522 context_id: ContextId::for_lookup(),
523 }))
524 }
525
526 pub fn load(self) -> Task<Option<(AgentContext, Vec<Entity<Buffer>>)>> {
527 Task::ready(Some((AgentContext::FetchedUrl(self), vec![])))
528 }
529}
530
531impl Display for FetchedUrlContext {
532 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
533 // TODO: Better format - url and contents are not delimited.
534 write!(f, "{}\n{}\n", self.url, self.text)
535 }
536}
537
538#[derive(Debug, Clone)]
539pub struct ThreadContextHandle {
540 pub thread: Entity<Thread>,
541 pub context_id: ContextId,
542}
543
544#[derive(Debug, Clone)]
545pub struct ThreadContext {
546 pub handle: ThreadContextHandle,
547 pub title: SharedString,
548 pub text: SharedString,
549}
550
551impl ThreadContextHandle {
552 pub fn eq_for_key(&self, other: &Self) -> bool {
553 self.thread == other.thread
554 }
555
556 pub fn hash_for_key<H: Hasher>(&self, state: &mut H) {
557 self.thread.hash(state)
558 }
559
560 pub fn title(&self, cx: &App) -> SharedString {
561 self.thread.read(cx).summary().or_default()
562 }
563
564 fn load(self, cx: &App) -> Task<Option<(AgentContext, Vec<Entity<Buffer>>)>> {
565 cx.spawn(async move |cx| {
566 let text = Thread::wait_for_detailed_summary_or_text(&self.thread, cx).await?;
567 let title = self
568 .thread
569 .read_with(cx, |thread, _cx| thread.summary().or_default())
570 .ok()?;
571 let context = AgentContext::Thread(ThreadContext {
572 title,
573 text,
574 handle: self,
575 });
576 Some((context, vec![]))
577 })
578 }
579}
580
581impl Display for ThreadContext {
582 fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
583 // TODO: Better format for this - doesn't distinguish title and contents.
584 write!(f, "{}\n{}\n", &self.title, &self.text.trim())
585 }
586}
587
588#[derive(Debug, Clone)]
589pub struct TextThreadContextHandle {
590 pub context: Entity<AssistantContext>,
591 pub context_id: ContextId,
592}
593
594#[derive(Debug, Clone)]
595pub struct TextThreadContext {
596 pub handle: TextThreadContextHandle,
597 pub title: SharedString,
598 pub text: SharedString,
599}
600
601impl TextThreadContextHandle {
602 // pub fn lookup_key() ->
603 pub fn eq_for_key(&self, other: &Self) -> bool {
604 self.context == other.context
605 }
606
607 pub fn hash_for_key<H: Hasher>(&self, state: &mut H) {
608 self.context.hash(state)
609 }
610
611 pub fn title(&self, cx: &App) -> SharedString {
612 self.context.read(cx).summary().or_default()
613 }
614
615 fn load(self, cx: &App) -> Task<Option<(AgentContext, Vec<Entity<Buffer>>)>> {
616 let title = self.title(cx);
617 let text = self.context.read(cx).to_xml(cx);
618 let context = AgentContext::TextThread(TextThreadContext {
619 title,
620 text: text.into(),
621 handle: self,
622 });
623 Task::ready(Some((context, vec![])))
624 }
625}
626
627impl Display for TextThreadContext {
628 fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
629 // TODO: escape title?
630 writeln!(f, "<text_thread title=\"{}\">", self.title)?;
631 write!(f, "{}", self.text.trim())?;
632 write!(f, "\n</text_thread>")
633 }
634}
635
636#[derive(Debug, Clone)]
637pub struct RulesContextHandle {
638 pub prompt_id: UserPromptId,
639 pub context_id: ContextId,
640}
641
642#[derive(Debug, Clone)]
643pub struct RulesContext {
644 pub handle: RulesContextHandle,
645 pub title: Option<SharedString>,
646 pub text: SharedString,
647}
648
649impl RulesContextHandle {
650 pub fn eq_for_key(&self, other: &Self) -> bool {
651 self.prompt_id == other.prompt_id
652 }
653
654 pub fn hash_for_key<H: Hasher>(&self, state: &mut H) {
655 self.prompt_id.hash(state)
656 }
657
658 pub fn lookup_key(prompt_id: UserPromptId) -> AgentContextKey {
659 AgentContextKey(AgentContextHandle::Rules(RulesContextHandle {
660 prompt_id,
661 context_id: ContextId::for_lookup(),
662 }))
663 }
664
665 pub fn load(
666 self,
667 prompt_store: &Option<Entity<PromptStore>>,
668 cx: &App,
669 ) -> Task<Option<(AgentContext, Vec<Entity<Buffer>>)>> {
670 let Some(prompt_store) = prompt_store.as_ref() else {
671 return Task::ready(None);
672 };
673 let prompt_store = prompt_store.read(cx);
674 let prompt_id = self.prompt_id.into();
675 let Some(metadata) = prompt_store.metadata(prompt_id) else {
676 return Task::ready(None);
677 };
678 let title = metadata.title;
679 let text_task = prompt_store.load(prompt_id, cx);
680 cx.background_spawn(async move {
681 // TODO: report load errors instead of just logging
682 let text = text_task.await.log_err()?.into();
683 let context = AgentContext::Rules(RulesContext {
684 handle: self,
685 title,
686 text,
687 });
688 Some((context, vec![]))
689 })
690 }
691}
692
693impl Display for RulesContext {
694 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
695 if let Some(title) = &self.title {
696 writeln!(f, "Rules title: {}", title)?;
697 }
698 let code_block = MarkdownCodeBlock {
699 tag: "",
700 text: self.text.trim(),
701 };
702 write!(f, "{code_block}")
703 }
704}
705
706#[derive(Debug, Clone)]
707pub struct ImageContext {
708 pub project_path: Option<ProjectPath>,
709 pub full_path: Option<String>,
710 pub original_image: Arc<gpui::Image>,
711 // TODO: handle this elsewhere and remove `ignore-interior-mutability` opt-out in clippy.toml
712 // needed due to a false positive of `clippy::mutable_key_type`.
713 pub image_task: Shared<Task<Option<LanguageModelImage>>>,
714 pub context_id: ContextId,
715}
716
717pub enum ImageStatus {
718 Loading,
719 Error,
720 Warning,
721 Ready,
722}
723
724impl ImageContext {
725 pub fn eq_for_key(&self, other: &Self) -> bool {
726 self.original_image.id() == other.original_image.id()
727 }
728
729 pub fn hash_for_key<H: Hasher>(&self, state: &mut H) {
730 self.original_image.id().hash(state);
731 }
732
733 pub fn image(&self) -> Option<LanguageModelImage> {
734 self.image_task.clone().now_or_never().flatten()
735 }
736
737 pub fn status(&self, model: Option<&Arc<dyn language_model::LanguageModel>>) -> ImageStatus {
738 match self.image_task.clone().now_or_never() {
739 None => ImageStatus::Loading,
740 Some(None) => ImageStatus::Error,
741 Some(Some(_)) => {
742 if model.is_some_and(|model| !model.supports_images()) {
743 ImageStatus::Warning
744 } else {
745 ImageStatus::Ready
746 }
747 }
748 }
749 }
750
751 pub fn load(self, cx: &App) -> Task<Option<(AgentContext, Vec<Entity<Buffer>>)>> {
752 cx.background_spawn(async move {
753 self.image_task.clone().await;
754 Some((AgentContext::Image(self), vec![]))
755 })
756 }
757}
758
759#[derive(Debug, Clone, Default)]
760pub struct ContextLoadResult {
761 pub loaded_context: LoadedContext,
762 pub referenced_buffers: HashSet<Entity<Buffer>>,
763}
764
765#[derive(Debug, Clone, Default)]
766pub struct LoadedContext {
767 pub contexts: Vec<AgentContext>,
768 pub text: String,
769 pub images: Vec<LanguageModelImage>,
770}
771
772impl LoadedContext {
773 pub fn is_empty(&self) -> bool {
774 self.text.is_empty() && self.images.is_empty()
775 }
776
777 pub fn add_to_request_message(&self, request_message: &mut LanguageModelRequestMessage) {
778 if !self.text.is_empty() {
779 request_message
780 .content
781 .push(MessageContent::Text(self.text.to_string()));
782 }
783
784 if !self.images.is_empty() {
785 // Some providers only support image parts after an initial text part
786 if request_message.content.is_empty() {
787 request_message
788 .content
789 .push(MessageContent::Text("Images attached by user:".to_string()));
790 }
791
792 for image in &self.images {
793 request_message
794 .content
795 .push(MessageContent::Image(image.clone()))
796 }
797 }
798 }
799}
800
801/// Loads and formats a collection of contexts.
802pub fn load_context(
803 contexts: Vec<AgentContextHandle>,
804 project: &Entity<Project>,
805 prompt_store: &Option<Entity<PromptStore>>,
806 cx: &mut App,
807) -> Task<ContextLoadResult> {
808 let load_tasks: Vec<_> = contexts
809 .into_iter()
810 .map(|context| match context {
811 AgentContextHandle::File(context) => context.load(cx),
812 AgentContextHandle::Directory(context) => context.load(project.clone(), cx),
813 AgentContextHandle::Symbol(context) => context.load(cx),
814 AgentContextHandle::Selection(context) => context.load(cx),
815 AgentContextHandle::FetchedUrl(context) => context.load(),
816 AgentContextHandle::Thread(context) => context.load(cx),
817 AgentContextHandle::TextThread(context) => context.load(cx),
818 AgentContextHandle::Rules(context) => context.load(prompt_store, cx),
819 AgentContextHandle::Image(context) => context.load(cx),
820 })
821 .collect();
822
823 cx.background_spawn(async move {
824 let load_results = future::join_all(load_tasks).await;
825
826 let mut contexts = Vec::new();
827 let mut text = String::new();
828 let mut referenced_buffers = HashSet::default();
829 for context in load_results {
830 let Some((context, buffers)) = context else {
831 continue;
832 };
833 contexts.push(context);
834 referenced_buffers.extend(buffers);
835 }
836
837 let mut file_context = Vec::new();
838 let mut directory_context = Vec::new();
839 let mut symbol_context = Vec::new();
840 let mut selection_context = Vec::new();
841 let mut fetched_url_context = Vec::new();
842 let mut thread_context = Vec::new();
843 let mut text_thread_context = Vec::new();
844 let mut rules_context = Vec::new();
845 let mut images = Vec::new();
846 for context in &contexts {
847 match context {
848 AgentContext::File(context) => file_context.push(context),
849 AgentContext::Directory(context) => directory_context.push(context),
850 AgentContext::Symbol(context) => symbol_context.push(context),
851 AgentContext::Selection(context) => selection_context.push(context),
852 AgentContext::FetchedUrl(context) => fetched_url_context.push(context),
853 AgentContext::Thread(context) => thread_context.push(context),
854 AgentContext::TextThread(context) => text_thread_context.push(context),
855 AgentContext::Rules(context) => rules_context.push(context),
856 AgentContext::Image(context) => images.extend(context.image()),
857 }
858 }
859
860 // Use empty text if there are no contexts that contribute to text (everything but image
861 // context).
862 if file_context.is_empty()
863 && directory_context.is_empty()
864 && symbol_context.is_empty()
865 && selection_context.is_empty()
866 && fetched_url_context.is_empty()
867 && thread_context.is_empty()
868 && text_thread_context.is_empty()
869 && rules_context.is_empty()
870 {
871 return ContextLoadResult {
872 loaded_context: LoadedContext {
873 contexts,
874 text,
875 images,
876 },
877 referenced_buffers,
878 };
879 }
880
881 text.push_str(
882 "\n<context>\n\
883 The following items were attached by the user. \
884 They are up-to-date and don't need to be re-read.\n\n",
885 );
886
887 if !file_context.is_empty() {
888 text.push_str("<files>");
889 for context in file_context {
890 text.push('\n');
891 let _ = write!(text, "{context}");
892 }
893 text.push_str("</files>\n");
894 }
895
896 if !directory_context.is_empty() {
897 text.push_str("<directories>");
898 for context in directory_context {
899 text.push('\n');
900 let _ = write!(text, "{context}");
901 }
902 text.push_str("</directories>\n");
903 }
904
905 if !symbol_context.is_empty() {
906 text.push_str("<symbols>");
907 for context in symbol_context {
908 text.push('\n');
909 let _ = write!(text, "{context}");
910 }
911 text.push_str("</symbols>\n");
912 }
913
914 if !selection_context.is_empty() {
915 text.push_str("<selections>");
916 for context in selection_context {
917 text.push('\n');
918 let _ = write!(text, "{context}");
919 }
920 text.push_str("</selections>\n");
921 }
922
923 if !fetched_url_context.is_empty() {
924 text.push_str("<fetched_urls>");
925 for context in fetched_url_context {
926 text.push('\n');
927 let _ = write!(text, "{context}");
928 }
929 text.push_str("</fetched_urls>\n");
930 }
931
932 if !thread_context.is_empty() {
933 text.push_str("<conversation_threads>");
934 for context in thread_context {
935 text.push('\n');
936 let _ = write!(text, "{context}");
937 }
938 text.push_str("</conversation_threads>\n");
939 }
940
941 if !text_thread_context.is_empty() {
942 text.push_str("<text_threads>");
943 for context in text_thread_context {
944 text.push('\n');
945 let _ = writeln!(text, "{context}");
946 }
947 text.push_str("<text_threads>");
948 }
949
950 if !rules_context.is_empty() {
951 text.push_str(
952 "<user_rules>\n\
953 The user has specified the following rules that should be applied:\n",
954 );
955 for context in rules_context {
956 text.push('\n');
957 let _ = write!(text, "{context}");
958 }
959 text.push_str("</user_rules>\n");
960 }
961
962 text.push_str("</context>\n");
963
964 ContextLoadResult {
965 loaded_context: LoadedContext {
966 contexts,
967 text,
968 images,
969 },
970 referenced_buffers,
971 }
972 })
973}
974
975fn collect_files_in_path(worktree: &Worktree, path: &RelPath) -> Vec<Arc<RelPath>> {
976 let mut files = Vec::new();
977
978 for entry in worktree.child_entries(path) {
979 if entry.is_dir() {
980 files.extend(collect_files_in_path(worktree, &entry.path));
981 } else if entry.is_file() {
982 files.push(entry.path.clone());
983 }
984 }
985
986 files
987}
988
989fn codeblock_tag(full_path: &str, line_range: Option<Range<Point>>) -> String {
990 let mut result = String::new();
991
992 if let Some(extension) = Path::new(full_path)
993 .extension()
994 .and_then(|ext| ext.to_str())
995 {
996 let _ = write!(result, "{} ", extension);
997 }
998
999 let _ = write!(result, "{}", full_path);
1000
1001 if let Some(range) = line_range {
1002 if range.start.row == range.end.row {
1003 let _ = write!(result, ":{}", range.start.row + 1);
1004 } else {
1005 let _ = write!(result, ":{}-{}", range.start.row + 1, range.end.row + 1);
1006 }
1007 }
1008
1009 result
1010}
1011
1012/// Wraps `AgentContext` to opt-in to `PartialEq` and `Hash` impls which use a subset of fields
1013/// needed for stable context identity.
1014#[derive(Debug, Clone, RefCast)]
1015#[repr(transparent)]
1016pub struct AgentContextKey(pub AgentContextHandle);
1017
1018impl AsRef<AgentContextHandle> for AgentContextKey {
1019 fn as_ref(&self) -> &AgentContextHandle {
1020 &self.0
1021 }
1022}
1023
1024impl Eq for AgentContextKey {}
1025
1026impl PartialEq for AgentContextKey {
1027 fn eq(&self, other: &Self) -> bool {
1028 match &self.0 {
1029 AgentContextHandle::File(context) => {
1030 if let AgentContextHandle::File(other_context) = &other.0 {
1031 return context.eq_for_key(other_context);
1032 }
1033 }
1034 AgentContextHandle::Directory(context) => {
1035 if let AgentContextHandle::Directory(other_context) = &other.0 {
1036 return context.eq_for_key(other_context);
1037 }
1038 }
1039 AgentContextHandle::Symbol(context) => {
1040 if let AgentContextHandle::Symbol(other_context) = &other.0 {
1041 return context.eq_for_key(other_context);
1042 }
1043 }
1044 AgentContextHandle::Selection(context) => {
1045 if let AgentContextHandle::Selection(other_context) = &other.0 {
1046 return context.eq_for_key(other_context);
1047 }
1048 }
1049 AgentContextHandle::FetchedUrl(context) => {
1050 if let AgentContextHandle::FetchedUrl(other_context) = &other.0 {
1051 return context.eq_for_key(other_context);
1052 }
1053 }
1054 AgentContextHandle::Thread(context) => {
1055 if let AgentContextHandle::Thread(other_context) = &other.0 {
1056 return context.eq_for_key(other_context);
1057 }
1058 }
1059 AgentContextHandle::Rules(context) => {
1060 if let AgentContextHandle::Rules(other_context) = &other.0 {
1061 return context.eq_for_key(other_context);
1062 }
1063 }
1064 AgentContextHandle::Image(context) => {
1065 if let AgentContextHandle::Image(other_context) = &other.0 {
1066 return context.eq_for_key(other_context);
1067 }
1068 }
1069 AgentContextHandle::TextThread(context) => {
1070 if let AgentContextHandle::TextThread(other_context) = &other.0 {
1071 return context.eq_for_key(other_context);
1072 }
1073 }
1074 }
1075 false
1076 }
1077}
1078
1079impl Hash for AgentContextKey {
1080 fn hash<H: Hasher>(&self, state: &mut H) {
1081 match &self.0 {
1082 AgentContextHandle::File(context) => context.hash_for_key(state),
1083 AgentContextHandle::Directory(context) => context.hash_for_key(state),
1084 AgentContextHandle::Symbol(context) => context.hash_for_key(state),
1085 AgentContextHandle::Selection(context) => context.hash_for_key(state),
1086 AgentContextHandle::FetchedUrl(context) => context.hash_for_key(state),
1087 AgentContextHandle::Thread(context) => context.hash_for_key(state),
1088 AgentContextHandle::TextThread(context) => context.hash_for_key(state),
1089 AgentContextHandle::Rules(context) => context.hash_for_key(state),
1090 AgentContextHandle::Image(context) => context.hash_for_key(state),
1091 }
1092 }
1093}
1094
1095#[cfg(test)]
1096mod tests {
1097 use super::*;
1098 use gpui::TestAppContext;
1099 use project::{FakeFs, Project};
1100 use serde_json::json;
1101 use settings::SettingsStore;
1102 use util::path;
1103
1104 fn init_test_settings(cx: &mut TestAppContext) {
1105 cx.update(|cx| {
1106 let settings_store = SettingsStore::test(cx);
1107 cx.set_global(settings_store);
1108 language::init(cx);
1109 Project::init_settings(cx);
1110 });
1111 }
1112
1113 // Helper to create a test project with test files
1114 async fn create_test_project(
1115 cx: &mut TestAppContext,
1116 files: serde_json::Value,
1117 ) -> Entity<Project> {
1118 let fs = FakeFs::new(cx.background_executor.clone());
1119 fs.insert_tree(path!("/test"), files).await;
1120 Project::test(fs, [path!("/test").as_ref()], cx).await
1121 }
1122
1123 #[gpui::test]
1124 async fn test_large_file_uses_outline(cx: &mut TestAppContext) {
1125 init_test_settings(cx);
1126
1127 // Create a large file that exceeds AUTO_OUTLINE_SIZE
1128 const LINE: &str = "Line with some text\n";
1129 let large_content = LINE.repeat(2 * (outline::AUTO_OUTLINE_SIZE / LINE.len()));
1130 let content_len = large_content.len();
1131
1132 assert!(content_len > outline::AUTO_OUTLINE_SIZE);
1133
1134 let file_context = file_context_for(large_content, cx).await;
1135
1136 assert!(
1137 file_context.is_outline,
1138 "Large file should use outline format"
1139 );
1140
1141 assert!(
1142 file_context.text.len() < content_len,
1143 "Outline should be smaller than original content"
1144 );
1145 }
1146
1147 #[gpui::test]
1148 async fn test_small_file_uses_full_content(cx: &mut TestAppContext) {
1149 init_test_settings(cx);
1150
1151 let small_content = "This is a small file.\n";
1152 let content_len = small_content.len();
1153
1154 assert!(content_len < outline::AUTO_OUTLINE_SIZE);
1155
1156 let file_context = file_context_for(small_content.to_string(), cx).await;
1157
1158 assert!(
1159 !file_context.is_outline,
1160 "Small files should not get an outline"
1161 );
1162
1163 assert_eq!(file_context.text, small_content);
1164 }
1165
1166 async fn file_context_for(content: String, cx: &mut TestAppContext) -> FileContext {
1167 // Create a test project with the file
1168 let project = create_test_project(
1169 cx,
1170 json!({
1171 "file.txt": content,
1172 }),
1173 )
1174 .await;
1175
1176 // Open the buffer
1177 let buffer_path = project
1178 .read_with(cx, |project, cx| project.find_project_path("file.txt", cx))
1179 .unwrap();
1180
1181 let buffer = project
1182 .update(cx, |project, cx| project.open_buffer(buffer_path, cx))
1183 .await
1184 .unwrap();
1185
1186 let context_handle = AgentContextHandle::File(FileContextHandle {
1187 buffer: buffer.clone(),
1188 context_id: ContextId::zero(),
1189 });
1190
1191 cx.update(|cx| load_context(vec![context_handle], &project, &None, cx))
1192 .await
1193 .loaded_context
1194 .contexts
1195 .into_iter()
1196 .find_map(|ctx| {
1197 if let AgentContext::File(file_ctx) = ctx {
1198 Some(file_ctx)
1199 } else {
1200 None
1201 }
1202 })
1203 .expect("Should have found a file context")
1204 }
1205}