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.then(|| LanguageModelRequestMessage {
 91            role: Role::System,
 92            content: self.snapshot.message.to_string(),
 93        })
 94    }
 95}
 96
 97#[derive(Clone, Default, Debug)]
 98pub struct RecentBuffersSnapshot {
 99    pub message: Rope,
100    pub source_buffers: Vec<SourceBufferSnapshot>,
101}
102
103#[derive(Clone)]
104pub struct SourceBufferSnapshot {
105    pub full_path: Option<PathBuf>,
106    pub model: WeakModel<Buffer>,
107    pub snapshot: BufferSnapshot,
108}
109
110impl std::fmt::Debug for SourceBufferSnapshot {
111    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
112        f.debug_struct("SourceBufferSnapshot")
113            .field("full_path", &self.full_path)
114            .field("model (entity id)", &self.model.entity_id())
115            .field("snapshot (text)", &self.snapshot.text())
116            .finish()
117    }
118}
119
120fn message_for_recent_buffers(buffers: Vec<SourceBufferSnapshot>) -> Rope {
121    let mut message = String::new();
122    writeln!(
123        message,
124        "The following is a list of recent buffers that the user has opened."
125    )
126    .unwrap();
127
128    for buffer in buffers {
129        if let Some(path) = buffer.full_path {
130            writeln!(message, "```{}", path.display()).unwrap();
131        } else {
132            writeln!(message, "```untitled").unwrap();
133        }
134
135        for chunk in buffer.snapshot.chunks(0..buffer.snapshot.len(), false) {
136            message.push_str(chunk.text);
137        }
138        if !message.ends_with('\n') {
139            message.push('\n');
140        }
141        message.push_str("```\n");
142    }
143
144    Rope::from(message.as_str())
145}