1use crate::{assistant_panel::Conversation, LanguageModelRequestMessage, Role};
  2use gpui::{ModelContext, Subscription, Task, WeakModel};
  3use language::{Buffer, BufferSnapshot, Rope};
  4use std::{fmt::Write, path::PathBuf, time::Duration};
  5
  6use super::ContextUpdated;
  7
  8pub struct RecentBuffersContext {
  9    pub enabled: bool,
 10    pub buffers: Vec<RecentBuffer>,
 11    pub snapshot: RecentBuffersSnapshot,
 12    pub pending_message: Option<Task<()>>,
 13}
 14
 15pub struct RecentBuffer {
 16    pub buffer: WeakModel<Buffer>,
 17    pub _subscription: Subscription,
 18}
 19
 20impl Default for RecentBuffersContext {
 21    fn default() -> Self {
 22        Self {
 23            enabled: true,
 24            buffers: Vec::new(),
 25            snapshot: RecentBuffersSnapshot::default(),
 26            pending_message: None,
 27        }
 28    }
 29}
 30
 31impl RecentBuffersContext {
 32    pub fn update(&mut self, cx: &mut ModelContext<Conversation>) -> ContextUpdated {
 33        let source_buffers = self
 34            .buffers
 35            .iter()
 36            .filter_map(|recent| {
 37                let (full_path, snapshot) = recent
 38                    .buffer
 39                    .read_with(cx, |buffer, cx| {
 40                        (
 41                            buffer.file().map(|file| file.full_path(cx)),
 42                            buffer.snapshot(),
 43                        )
 44                    })
 45                    .ok()?;
 46                Some(SourceBufferSnapshot {
 47                    full_path,
 48                    model: recent.buffer.clone(),
 49                    snapshot,
 50                })
 51            })
 52            .collect::<Vec<_>>();
 53
 54        if !self.enabled || source_buffers.is_empty() {
 55            self.snapshot.message = Default::default();
 56            self.snapshot.source_buffers.clear();
 57            self.pending_message = None;
 58            cx.notify();
 59            ContextUpdated::Disabled
 60        } else {
 61            self.pending_message = Some(cx.spawn(|this, mut cx| async move {
 62                const DEBOUNCE_TIMEOUT: Duration = Duration::from_millis(100);
 63                cx.background_executor().timer(DEBOUNCE_TIMEOUT).await;
 64
 65                let message = if source_buffers.is_empty() {
 66                    Rope::new()
 67                } else {
 68                    cx.background_executor()
 69                        .spawn({
 70                            let source_buffers = source_buffers.clone();
 71                            async move { message_for_recent_buffers(source_buffers) }
 72                        })
 73                        .await
 74                };
 75                this.update(&mut cx, |this, cx| {
 76                    this.ambient_context.recent_buffers.snapshot.source_buffers = source_buffers;
 77                    this.ambient_context.recent_buffers.snapshot.message = message;
 78                    this.count_remaining_tokens(cx);
 79                    cx.notify();
 80                })
 81                .ok();
 82            }));
 83
 84            ContextUpdated::Updating
 85        }
 86    }
 87
 88    /// Returns the [`RecentBuffersContext`] as a message to the language model.
 89    pub fn to_message(&self) -> Option<LanguageModelRequestMessage> {
 90        self.enabled
 91            .then(|| LanguageModelRequestMessage {
 92                role: Role::System,
 93                content: self.snapshot.message.to_string(),
 94            })
 95            .filter(|message| !message.content.is_empty())
 96    }
 97}
 98
 99#[derive(Clone, Default, Debug)]
100pub struct RecentBuffersSnapshot {
101    pub message: Rope,
102    pub source_buffers: Vec<SourceBufferSnapshot>,
103}
104
105#[derive(Clone)]
106pub struct SourceBufferSnapshot {
107    pub full_path: Option<PathBuf>,
108    pub model: WeakModel<Buffer>,
109    pub snapshot: BufferSnapshot,
110}
111
112impl std::fmt::Debug for SourceBufferSnapshot {
113    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
114        f.debug_struct("SourceBufferSnapshot")
115            .field("full_path", &self.full_path)
116            .field("model (entity id)", &self.model.entity_id())
117            .field("snapshot (text)", &self.snapshot.text())
118            .finish()
119    }
120}
121
122fn message_for_recent_buffers(buffers: Vec<SourceBufferSnapshot>) -> Rope {
123    let mut message = String::new();
124    writeln!(
125        message,
126        "The following is a list of recent buffers that the user has opened."
127    )
128    .unwrap();
129
130    for buffer in buffers {
131        if let Some(path) = buffer.full_path {
132            writeln!(message, "```{}", path.display()).unwrap();
133        } else {
134            writeln!(message, "```untitled").unwrap();
135        }
136
137        for chunk in buffer.snapshot.chunks(0..buffer.snapshot.len(), false) {
138            message.push_str(chunk.text);
139        }
140        if !message.ends_with('\n') {
141            message.push('\n');
142        }
143        message.push_str("```\n");
144    }
145
146    Rope::from(message.as_str())
147}