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