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