context.rs

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