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