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