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 full_path: Option<Arc<Path>>,
738 pub original_image: Arc<gpui::Image>,
739 // TODO: handle this elsewhere and remove `ignore-interior-mutability` opt-out in clippy.toml
740 // needed due to a false positive of `clippy::mutable_key_type`.
741 pub image_task: Shared<Task<Option<LanguageModelImage>>>,
742 pub context_id: ContextId,
743}
744
745pub enum ImageStatus {
746 Loading,
747 Error,
748 Ready,
749}
750
751impl ImageContext {
752 pub fn eq_for_key(&self, other: &Self) -> bool {
753 self.original_image.id() == other.original_image.id()
754 }
755
756 pub fn hash_for_key<H: Hasher>(&self, state: &mut H) {
757 self.original_image.id().hash(state);
758 }
759
760 pub fn image(&self) -> Option<LanguageModelImage> {
761 self.image_task.clone().now_or_never().flatten()
762 }
763
764 pub fn status(&self) -> ImageStatus {
765 match self.image_task.clone().now_or_never() {
766 None => ImageStatus::Loading,
767 Some(None) => ImageStatus::Error,
768 Some(Some(_)) => ImageStatus::Ready,
769 }
770 }
771
772 pub fn load(self, cx: &App) -> Task<Option<(AgentContext, Vec<Entity<Buffer>>)>> {
773 cx.background_spawn(async move {
774 self.image_task.clone().await;
775 Some((AgentContext::Image(self), vec![]))
776 })
777 }
778}
779
780#[derive(Debug, Clone, Default)]
781pub struct ContextLoadResult {
782 pub loaded_context: LoadedContext,
783 pub referenced_buffers: HashSet<Entity<Buffer>>,
784}
785
786#[derive(Debug, Clone, Default)]
787pub struct LoadedContext {
788 pub contexts: Vec<AgentContext>,
789 pub text: String,
790 pub images: Vec<LanguageModelImage>,
791}
792
793impl LoadedContext {
794 pub fn is_empty(&self) -> bool {
795 self.text.is_empty() && self.images.is_empty()
796 }
797
798 pub fn add_to_request_message(&self, request_message: &mut LanguageModelRequestMessage) {
799 if !self.text.is_empty() {
800 request_message
801 .content
802 .push(MessageContent::Text(self.text.to_string()));
803 }
804
805 if !self.images.is_empty() {
806 // Some providers only support image parts after an initial text part
807 if request_message.content.is_empty() {
808 request_message
809 .content
810 .push(MessageContent::Text("Images attached by user:".to_string()));
811 }
812
813 for image in &self.images {
814 request_message
815 .content
816 .push(MessageContent::Image(image.clone()))
817 }
818 }
819 }
820}
821
822/// Loads and formats a collection of contexts.
823pub fn load_context(
824 contexts: Vec<AgentContextHandle>,
825 project: &Entity<Project>,
826 prompt_store: &Option<Entity<PromptStore>>,
827 cx: &mut App,
828) -> Task<ContextLoadResult> {
829 let load_tasks: Vec<_> = contexts
830 .into_iter()
831 .map(|context| match context {
832 AgentContextHandle::File(context) => context.load(cx),
833 AgentContextHandle::Directory(context) => context.load(project.clone(), cx),
834 AgentContextHandle::Symbol(context) => context.load(cx),
835 AgentContextHandle::Selection(context) => context.load(cx),
836 AgentContextHandle::FetchedUrl(context) => context.load(),
837 AgentContextHandle::Thread(context) => context.load(cx),
838 AgentContextHandle::TextThread(context) => context.load(cx),
839 AgentContextHandle::Rules(context) => context.load(prompt_store, cx),
840 AgentContextHandle::Image(context) => context.load(cx),
841 })
842 .collect();
843
844 cx.background_spawn(async move {
845 let load_results = future::join_all(load_tasks).await;
846
847 let mut contexts = Vec::new();
848 let mut text = String::new();
849 let mut referenced_buffers = HashSet::default();
850 for context in load_results {
851 let Some((context, buffers)) = context else {
852 continue;
853 };
854 contexts.push(context);
855 referenced_buffers.extend(buffers);
856 }
857
858 let mut file_context = Vec::new();
859 let mut directory_context = Vec::new();
860 let mut symbol_context = Vec::new();
861 let mut selection_context = Vec::new();
862 let mut fetched_url_context = Vec::new();
863 let mut thread_context = Vec::new();
864 let mut text_thread_context = Vec::new();
865 let mut rules_context = Vec::new();
866 let mut images = Vec::new();
867 for context in &contexts {
868 match context {
869 AgentContext::File(context) => file_context.push(context),
870 AgentContext::Directory(context) => directory_context.push(context),
871 AgentContext::Symbol(context) => symbol_context.push(context),
872 AgentContext::Selection(context) => selection_context.push(context),
873 AgentContext::FetchedUrl(context) => fetched_url_context.push(context),
874 AgentContext::Thread(context) => thread_context.push(context),
875 AgentContext::TextThread(context) => text_thread_context.push(context),
876 AgentContext::Rules(context) => rules_context.push(context),
877 AgentContext::Image(context) => images.extend(context.image()),
878 }
879 }
880
881 // Use empty text if there are no contexts that contribute to text (everything but image
882 // context).
883 if file_context.is_empty()
884 && directory_context.is_empty()
885 && symbol_context.is_empty()
886 && selection_context.is_empty()
887 && fetched_url_context.is_empty()
888 && thread_context.is_empty()
889 && text_thread_context.is_empty()
890 && rules_context.is_empty()
891 {
892 return ContextLoadResult {
893 loaded_context: LoadedContext {
894 contexts,
895 text,
896 images,
897 },
898 referenced_buffers,
899 };
900 }
901
902 text.push_str(
903 "\n<context>\n\
904 The following items were attached by the user. \
905 They are up-to-date and don't need to be re-read.\n\n",
906 );
907
908 if !file_context.is_empty() {
909 text.push_str("<files>");
910 for context in file_context {
911 text.push('\n');
912 let _ = write!(text, "{context}");
913 }
914 text.push_str("</files>\n");
915 }
916
917 if !directory_context.is_empty() {
918 text.push_str("<directories>");
919 for context in directory_context {
920 text.push('\n');
921 let _ = write!(text, "{context}");
922 }
923 text.push_str("</directories>\n");
924 }
925
926 if !symbol_context.is_empty() {
927 text.push_str("<symbols>");
928 for context in symbol_context {
929 text.push('\n');
930 let _ = write!(text, "{context}");
931 }
932 text.push_str("</symbols>\n");
933 }
934
935 if !selection_context.is_empty() {
936 text.push_str("<selections>");
937 for context in selection_context {
938 text.push('\n');
939 let _ = write!(text, "{context}");
940 }
941 text.push_str("</selections>\n");
942 }
943
944 if !fetched_url_context.is_empty() {
945 text.push_str("<fetched_urls>");
946 for context in fetched_url_context {
947 text.push('\n');
948 let _ = write!(text, "{context}");
949 }
950 text.push_str("</fetched_urls>\n");
951 }
952
953 if !thread_context.is_empty() {
954 text.push_str("<conversation_threads>");
955 for context in thread_context {
956 text.push('\n');
957 let _ = write!(text, "{context}");
958 }
959 text.push_str("</conversation_threads>\n");
960 }
961
962 if !text_thread_context.is_empty() {
963 text.push_str("<text_threads>");
964 for context in text_thread_context {
965 text.push('\n');
966 let _ = writeln!(text, "{context}");
967 }
968 text.push_str("<text_threads>");
969 }
970
971 if !rules_context.is_empty() {
972 text.push_str(
973 "<user_rules>\n\
974 The user has specified the following rules that should be applied:\n",
975 );
976 for context in rules_context {
977 text.push('\n');
978 let _ = write!(text, "{context}");
979 }
980 text.push_str("</user_rules>\n");
981 }
982
983 text.push_str("</context>\n");
984
985 ContextLoadResult {
986 loaded_context: LoadedContext {
987 contexts,
988 text,
989 images,
990 },
991 referenced_buffers,
992 }
993 })
994}
995
996fn collect_files_in_path(worktree: &Worktree, path: &Path) -> Vec<Arc<Path>> {
997 let mut files = Vec::new();
998
999 for entry in worktree.child_entries(path) {
1000 if entry.is_dir() {
1001 files.extend(collect_files_in_path(worktree, &entry.path));
1002 } else if entry.is_file() {
1003 files.push(entry.path.clone());
1004 }
1005 }
1006
1007 files
1008}
1009
1010fn codeblock_tag(full_path: &Path, line_range: Option<Range<Point>>) -> String {
1011 let mut result = String::new();
1012
1013 if let Some(extension) = full_path.extension().and_then(|ext| ext.to_str()) {
1014 let _ = write!(result, "{} ", extension);
1015 }
1016
1017 let _ = write!(result, "{}", full_path.display());
1018
1019 if let Some(range) = line_range {
1020 if range.start.row == range.end.row {
1021 let _ = write!(result, ":{}", range.start.row + 1);
1022 } else {
1023 let _ = write!(result, ":{}-{}", range.start.row + 1, range.end.row + 1);
1024 }
1025 }
1026
1027 result
1028}
1029
1030/// Wraps `AgentContext` to opt-in to `PartialEq` and `Hash` impls which use a subset of fields
1031/// needed for stable context identity.
1032#[derive(Debug, Clone, RefCast)]
1033#[repr(transparent)]
1034pub struct AgentContextKey(pub AgentContextHandle);
1035
1036impl AsRef<AgentContextHandle> for AgentContextKey {
1037 fn as_ref(&self) -> &AgentContextHandle {
1038 &self.0
1039 }
1040}
1041
1042impl Eq for AgentContextKey {}
1043
1044impl PartialEq for AgentContextKey {
1045 fn eq(&self, other: &Self) -> bool {
1046 match &self.0 {
1047 AgentContextHandle::File(context) => {
1048 if let AgentContextHandle::File(other_context) = &other.0 {
1049 return context.eq_for_key(other_context);
1050 }
1051 }
1052 AgentContextHandle::Directory(context) => {
1053 if let AgentContextHandle::Directory(other_context) = &other.0 {
1054 return context.eq_for_key(other_context);
1055 }
1056 }
1057 AgentContextHandle::Symbol(context) => {
1058 if let AgentContextHandle::Symbol(other_context) = &other.0 {
1059 return context.eq_for_key(other_context);
1060 }
1061 }
1062 AgentContextHandle::Selection(context) => {
1063 if let AgentContextHandle::Selection(other_context) = &other.0 {
1064 return context.eq_for_key(other_context);
1065 }
1066 }
1067 AgentContextHandle::FetchedUrl(context) => {
1068 if let AgentContextHandle::FetchedUrl(other_context) = &other.0 {
1069 return context.eq_for_key(other_context);
1070 }
1071 }
1072 AgentContextHandle::Thread(context) => {
1073 if let AgentContextHandle::Thread(other_context) = &other.0 {
1074 return context.eq_for_key(other_context);
1075 }
1076 }
1077 AgentContextHandle::Rules(context) => {
1078 if let AgentContextHandle::Rules(other_context) = &other.0 {
1079 return context.eq_for_key(other_context);
1080 }
1081 }
1082 AgentContextHandle::Image(context) => {
1083 if let AgentContextHandle::Image(other_context) = &other.0 {
1084 return context.eq_for_key(other_context);
1085 }
1086 }
1087 AgentContextHandle::TextThread(context) => {
1088 if let AgentContextHandle::TextThread(other_context) = &other.0 {
1089 return context.eq_for_key(other_context);
1090 }
1091 }
1092 }
1093 false
1094 }
1095}
1096
1097impl Hash for AgentContextKey {
1098 fn hash<H: Hasher>(&self, state: &mut H) {
1099 match &self.0 {
1100 AgentContextHandle::File(context) => context.hash_for_key(state),
1101 AgentContextHandle::Directory(context) => context.hash_for_key(state),
1102 AgentContextHandle::Symbol(context) => context.hash_for_key(state),
1103 AgentContextHandle::Selection(context) => context.hash_for_key(state),
1104 AgentContextHandle::FetchedUrl(context) => context.hash_for_key(state),
1105 AgentContextHandle::Thread(context) => context.hash_for_key(state),
1106 AgentContextHandle::TextThread(context) => context.hash_for_key(state),
1107 AgentContextHandle::Rules(context) => context.hash_for_key(state),
1108 AgentContextHandle::Image(context) => context.hash_for_key(state),
1109 }
1110 }
1111}
1112
1113#[derive(Default)]
1114pub struct ContextCreasesAddon {
1115 creases: HashMap<AgentContextKey, Vec<(CreaseId, SharedString)>>,
1116 _subscription: Option<Subscription>,
1117}
1118
1119impl Addon for ContextCreasesAddon {
1120 fn to_any(&self) -> &dyn std::any::Any {
1121 self
1122 }
1123
1124 fn to_any_mut(&mut self) -> Option<&mut dyn std::any::Any> {
1125 Some(self)
1126 }
1127}
1128
1129impl ContextCreasesAddon {
1130 pub fn new() -> Self {
1131 Self {
1132 creases: HashMap::default(),
1133 _subscription: None,
1134 }
1135 }
1136
1137 pub fn add_creases(
1138 &mut self,
1139 context_store: &Entity<ContextStore>,
1140 key: AgentContextKey,
1141 creases: impl IntoIterator<Item = (CreaseId, SharedString)>,
1142 cx: &mut Context<Editor>,
1143 ) {
1144 self.creases.entry(key).or_default().extend(creases);
1145 self._subscription = Some(cx.subscribe(
1146 &context_store,
1147 |editor, _, event, cx| match event {
1148 ContextStoreEvent::ContextRemoved(key) => {
1149 let Some(this) = editor.addon_mut::<Self>() else {
1150 return;
1151 };
1152 let (crease_ids, replacement_texts): (Vec<_>, Vec<_>) = this
1153 .creases
1154 .remove(key)
1155 .unwrap_or_default()
1156 .into_iter()
1157 .unzip();
1158 let ranges = editor
1159 .remove_creases(crease_ids, cx)
1160 .into_iter()
1161 .map(|(_, range)| range)
1162 .collect::<Vec<_>>();
1163 editor.unfold_ranges(&ranges, false, false, cx);
1164 editor.edit(ranges.into_iter().zip(replacement_texts), cx);
1165 cx.notify();
1166 }
1167 },
1168 ))
1169 }
1170
1171 pub fn into_inner(self) -> HashMap<AgentContextKey, Vec<(CreaseId, SharedString)>> {
1172 self.creases
1173 }
1174}
1175
1176#[cfg(test)]
1177mod tests {
1178 use super::*;
1179 use gpui::TestAppContext;
1180 use project::{FakeFs, Project};
1181 use serde_json::json;
1182 use settings::SettingsStore;
1183 use util::path;
1184
1185 fn init_test_settings(cx: &mut TestAppContext) {
1186 cx.update(|cx| {
1187 let settings_store = SettingsStore::test(cx);
1188 cx.set_global(settings_store);
1189 language::init(cx);
1190 Project::init_settings(cx);
1191 });
1192 }
1193
1194 // Helper to create a test project with test files
1195 async fn create_test_project(
1196 cx: &mut TestAppContext,
1197 files: serde_json::Value,
1198 ) -> Entity<Project> {
1199 let fs = FakeFs::new(cx.background_executor.clone());
1200 fs.insert_tree(path!("/test"), files).await;
1201 Project::test(fs, [path!("/test").as_ref()], cx).await
1202 }
1203
1204 #[gpui::test]
1205 async fn test_large_file_uses_outline(cx: &mut TestAppContext) {
1206 init_test_settings(cx);
1207
1208 // Create a large file that exceeds AUTO_OUTLINE_SIZE
1209 const LINE: &str = "Line with some text\n";
1210 let large_content = LINE.repeat(2 * (outline::AUTO_OUTLINE_SIZE / LINE.len()));
1211 let content_len = large_content.len();
1212
1213 assert!(content_len > outline::AUTO_OUTLINE_SIZE);
1214
1215 let file_context = file_context_for(large_content, cx).await;
1216
1217 assert!(
1218 file_context.is_outline,
1219 "Large file should use outline format"
1220 );
1221
1222 assert!(
1223 file_context.text.len() < content_len,
1224 "Outline should be smaller than original content"
1225 );
1226 }
1227
1228 #[gpui::test]
1229 async fn test_small_file_uses_full_content(cx: &mut TestAppContext) {
1230 init_test_settings(cx);
1231
1232 let small_content = "This is a small file.\n";
1233 let content_len = small_content.len();
1234
1235 assert!(content_len < outline::AUTO_OUTLINE_SIZE);
1236
1237 let file_context = file_context_for(small_content.to_string(), cx).await;
1238
1239 assert!(
1240 !file_context.is_outline,
1241 "Small files should not get an outline"
1242 );
1243
1244 assert_eq!(file_context.text, small_content);
1245 }
1246
1247 async fn file_context_for(content: String, cx: &mut TestAppContext) -> FileContext {
1248 // Create a test project with the file
1249 let project = create_test_project(
1250 cx,
1251 json!({
1252 "file.txt": content,
1253 }),
1254 )
1255 .await;
1256
1257 // Open the buffer
1258 let buffer_path = project
1259 .read_with(cx, |project, cx| project.find_project_path("file.txt", cx))
1260 .unwrap();
1261
1262 let buffer = project
1263 .update(cx, |project, cx| project.open_buffer(buffer_path, cx))
1264 .await
1265 .unwrap();
1266
1267 let context_handle = AgentContextHandle::File(FileContextHandle {
1268 buffer: buffer.clone(),
1269 context_id: ContextId::zero(),
1270 });
1271
1272 cx.update(|cx| load_context(vec![context_handle], &project, &None, cx))
1273 .await
1274 .loaded_context
1275 .contexts
1276 .into_iter()
1277 .find_map(|ctx| {
1278 if let AgentContext::File(file_ctx) = ctx {
1279 Some(file_ctx)
1280 } else {
1281 None
1282 }
1283 })
1284 .expect("Should have found a file context")
1285 }
1286}