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