agent2.rs

  1mod acp;
  2mod thread_element;
  3
  4use anyhow::Result;
  5use async_trait::async_trait;
  6use chrono::{DateTime, Utc};
  7use gpui::{AppContext, AsyncApp, Context, Entity, SharedString, Task};
  8use project::Project;
  9use std::{ops::Range, path::PathBuf, sync::Arc};
 10
 11pub use acp::AcpAgent;
 12pub use thread_element::ThreadElement;
 13
 14#[async_trait(?Send)]
 15pub trait Agent: 'static {
 16    async fn threads(&self, cx: &mut AsyncApp) -> Result<Vec<AgentThreadSummary>>;
 17    async fn create_thread(self: Arc<Self>, cx: &mut AsyncApp) -> Result<Entity<Thread>>;
 18    async fn open_thread(&self, id: ThreadId, cx: &mut AsyncApp) -> Result<Entity<Thread>>;
 19    async fn thread_entries(
 20        &self,
 21        id: ThreadId,
 22        cx: &mut AsyncApp,
 23    ) -> Result<Vec<AgentThreadEntryContent>>;
 24    async fn send_thread_message(
 25        &self,
 26        thread_id: ThreadId,
 27        message: Message,
 28        cx: &mut AsyncApp,
 29    ) -> Result<()>;
 30}
 31
 32#[derive(Debug, Clone, PartialEq, Eq, Hash)]
 33pub struct ThreadId(SharedString);
 34
 35#[derive(Copy, Clone, Debug, PartialEq, Eq)]
 36pub struct FileVersion(u64);
 37
 38#[derive(Debug)]
 39pub struct AgentThreadSummary {
 40    pub id: ThreadId,
 41    pub title: String,
 42    pub created_at: DateTime<Utc>,
 43}
 44
 45#[derive(Clone, Debug, PartialEq, Eq)]
 46pub struct FileContent {
 47    pub path: PathBuf,
 48    pub version: FileVersion,
 49    pub content: SharedString,
 50}
 51
 52#[derive(Copy, Clone, Debug, Eq, PartialEq)]
 53pub enum Role {
 54    User,
 55    Assistant,
 56}
 57
 58#[derive(Clone, Debug, Eq, PartialEq)]
 59pub struct Message {
 60    pub role: Role,
 61    pub chunks: Vec<MessageChunk>,
 62}
 63
 64#[derive(Clone, Debug, Eq, PartialEq)]
 65pub enum MessageChunk {
 66    Text {
 67        chunk: SharedString,
 68    },
 69    File {
 70        content: FileContent,
 71    },
 72    Directory {
 73        path: PathBuf,
 74        contents: Vec<FileContent>,
 75    },
 76    Symbol {
 77        path: PathBuf,
 78        range: Range<u64>,
 79        version: FileVersion,
 80        name: SharedString,
 81        content: SharedString,
 82    },
 83    Fetch {
 84        url: SharedString,
 85        content: SharedString,
 86    },
 87}
 88
 89impl From<&str> for MessageChunk {
 90    fn from(chunk: &str) -> Self {
 91        MessageChunk::Text {
 92            chunk: chunk.to_string().into(),
 93        }
 94    }
 95}
 96
 97#[derive(Clone, Debug, Eq, PartialEq)]
 98pub enum AgentThreadEntryContent {
 99    Message(Message),
100    ReadFile { path: PathBuf, content: String },
101}
102
103#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
104pub struct ThreadEntryId(usize);
105
106impl ThreadEntryId {
107    pub fn post_inc(&mut self) -> Self {
108        let id = *self;
109        self.0 += 1;
110        id
111    }
112}
113
114#[derive(Debug)]
115pub struct ThreadEntry {
116    pub id: ThreadEntryId,
117    pub content: AgentThreadEntryContent,
118}
119
120pub struct ThreadStore {
121    threads: Vec<AgentThreadSummary>,
122    agent: Arc<dyn Agent>,
123    project: Entity<Project>,
124}
125
126impl ThreadStore {
127    pub async fn load(
128        agent: Arc<dyn Agent>,
129        project: Entity<Project>,
130        cx: &mut AsyncApp,
131    ) -> Result<Entity<Self>> {
132        let threads = agent.threads(cx).await?;
133        cx.new(|_cx| Self {
134            threads,
135            agent,
136            project,
137        })
138    }
139
140    /// Returns the threads in reverse chronological order.
141    pub fn threads(&self) -> &[AgentThreadSummary] {
142        &self.threads
143    }
144
145    /// Opens a thread with the given ID.
146    pub fn open_thread(
147        &self,
148        id: ThreadId,
149        cx: &mut Context<Self>,
150    ) -> Task<Result<Entity<Thread>>> {
151        let agent = self.agent.clone();
152        cx.spawn(async move |_, cx| agent.open_thread(id, cx).await)
153    }
154
155    /// Creates a new thread.
156    pub fn create_thread(&self, cx: &mut Context<Self>) -> Task<Result<Entity<Thread>>> {
157        let agent = self.agent.clone();
158        cx.spawn(async move |_, cx| agent.create_thread(cx).await)
159    }
160}
161
162pub struct Thread {
163    id: ThreadId,
164    next_entry_id: ThreadEntryId,
165    entries: Vec<ThreadEntry>,
166    agent: Arc<dyn Agent>,
167    title: SharedString,
168    project: Entity<Project>,
169}
170
171impl Thread {
172    pub async fn load(
173        agent: Arc<dyn Agent>,
174        thread_id: ThreadId,
175        project: Entity<Project>,
176        cx: &mut AsyncApp,
177    ) -> Result<Entity<Self>> {
178        let entries = agent.thread_entries(thread_id.clone(), cx).await?;
179        cx.new(|cx| Self::new(agent, thread_id, entries, project, cx))
180    }
181
182    pub fn new(
183        agent: Arc<dyn Agent>,
184        thread_id: ThreadId,
185        entries: Vec<AgentThreadEntryContent>,
186        project: Entity<Project>,
187        _: &mut Context<Self>,
188    ) -> Self {
189        let mut next_entry_id = ThreadEntryId(0);
190        Self {
191            title: "A new agent2 thread".into(),
192            entries: entries
193                .into_iter()
194                .map(|entry| ThreadEntry {
195                    id: next_entry_id.post_inc(),
196                    content: entry,
197                })
198                .collect(),
199            agent,
200            id: thread_id,
201            next_entry_id,
202            project,
203        }
204    }
205
206    pub fn title(&self) -> SharedString {
207        self.title.clone()
208    }
209
210    pub fn entries(&self) -> &[ThreadEntry] {
211        &self.entries
212    }
213
214    pub fn push_entry(&mut self, entry: AgentThreadEntryContent, cx: &mut Context<Self>) {
215        self.entries.push(ThreadEntry {
216            id: self.next_entry_id.post_inc(),
217            content: entry,
218        });
219        cx.notify();
220    }
221
222    pub fn push_assistant_chunk(&mut self, chunk: MessageChunk, cx: &mut Context<Self>) {
223        if let Some(last_entry) = self.entries.last_mut() {
224            if let AgentThreadEntryContent::Message(Message {
225                ref mut chunks,
226                role: Role::Assistant,
227            }) = last_entry.content
228            {
229                chunks.push(chunk);
230                return;
231            }
232        }
233
234        self.entries.push(ThreadEntry {
235            id: self.next_entry_id.post_inc(),
236            content: AgentThreadEntryContent::Message(Message {
237                role: Role::Assistant,
238                chunks: vec![chunk],
239            }),
240        });
241        cx.notify();
242    }
243
244    pub fn send(&mut self, message: Message, cx: &mut Context<Self>) -> Task<Result<()>> {
245        let agent = self.agent.clone();
246        let id = self.id.clone();
247        self.push_entry(AgentThreadEntryContent::Message(message.clone()), cx);
248        cx.spawn(async move |_, cx| {
249            agent.send_thread_message(id, message, cx).await?;
250            Ok(())
251        })
252    }
253}
254
255#[cfg(test)]
256mod tests {
257    use super::*;
258    use crate::acp::AcpAgent;
259    use gpui::TestAppContext;
260    use project::FakeFs;
261    use serde_json::json;
262    use settings::SettingsStore;
263    use std::{env, path::Path, process::Stdio};
264    use util::path;
265
266    fn init_test(cx: &mut TestAppContext) {
267        env_logger::init();
268        cx.update(|cx| {
269            let settings_store = SettingsStore::test(cx);
270            cx.set_global(settings_store);
271            Project::init_settings(cx);
272            language::init(cx);
273        });
274    }
275
276    #[gpui::test]
277    async fn test_gemini(cx: &mut TestAppContext) {
278        init_test(cx);
279
280        cx.executor().allow_parking();
281
282        let fs = FakeFs::new(cx.executor());
283        fs.insert_tree(
284            path!("/private/tmp"),
285            json!({"foo": "Lorem ipsum dolor", "bar": "bar", "baz": "baz"}),
286        )
287        .await;
288        let project = Project::test(fs, [path!("/private/tmp").as_ref()], cx).await;
289        let agent = gemini_agent(project.clone(), cx.to_async()).unwrap();
290        let thread_store = ThreadStore::load(agent, project, &mut cx.to_async())
291            .await
292            .unwrap();
293        let thread = thread_store
294            .update(cx, |thread_store, cx| {
295                assert_eq!(thread_store.threads().len(), 0);
296                thread_store.create_thread(cx)
297            })
298            .await
299            .unwrap();
300        thread
301            .update(cx, |thread, cx| {
302                thread.send(
303                    Message {
304                        role: Role::User,
305                        chunks: vec![
306                            "Read the '/private/tmp/foo' file and output all of its contents."
307                                .into(),
308                        ],
309                    },
310                    cx,
311                )
312            })
313            .await
314            .unwrap();
315
316        thread.read_with(cx, |thread, _| {
317            assert!(matches!(
318                thread.entries[0].content,
319                AgentThreadEntryContent::Message(Message {
320                    role: Role::User,
321                    ..
322                })
323            ));
324            assert!(
325                thread.entries().iter().any(|entry| {
326                    entry.content
327                        == AgentThreadEntryContent::ReadFile {
328                            path: "/private/tmp/foo".into(),
329                            content: "Lorem ipsum dolor".into(),
330                        }
331                }),
332                "Thread does not contain entry. Actual: {:?}",
333                thread.entries()
334            );
335        });
336    }
337
338    pub fn gemini_agent(project: Entity<Project>, mut cx: AsyncApp) -> Result<Arc<AcpAgent>> {
339        let cli_path =
340            Path::new(env!("CARGO_MANIFEST_DIR")).join("../../../gemini-cli/packages/cli");
341        let mut command = util::command::new_smol_command("node");
342        command
343            .arg(cli_path)
344            .arg("--acp")
345            .args(["--model", "gemini-2.5-flash"])
346            .current_dir("/private/tmp")
347            .stdin(Stdio::piped())
348            .stdout(Stdio::piped())
349            .stderr(Stdio::inherit())
350            .kill_on_drop(true);
351
352        if let Ok(gemini_key) = std::env::var("GEMINI_API_KEY") {
353            command.env("GEMINI_API_KEY", gemini_key);
354        }
355
356        let child = command.spawn().unwrap();
357
358        Ok(AcpAgent::stdio(child, project, &mut cx))
359    }
360}