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