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                // todo! merge with last chunk if same type
230                chunks.push(chunk);
231                return;
232            }
233        }
234
235        self.entries.push(ThreadEntry {
236            id: self.next_entry_id.post_inc(),
237            content: AgentThreadEntryContent::Message(Message {
238                role: Role::Assistant,
239                chunks: vec![chunk],
240            }),
241        });
242        cx.notify();
243    }
244
245    pub fn send(&mut self, message: Message, cx: &mut Context<Self>) -> Task<Result<()>> {
246        let agent = self.agent.clone();
247        let id = self.id.clone();
248        self.push_entry(AgentThreadEntryContent::Message(message.clone()), cx);
249        cx.spawn(async move |_, cx| {
250            agent.send_thread_message(id, message, cx).await?;
251            Ok(())
252        })
253    }
254}
255
256#[cfg(test)]
257mod tests {
258    use super::*;
259    use crate::acp::AcpAgent;
260    use gpui::TestAppContext;
261    use project::FakeFs;
262    use serde_json::json;
263    use settings::SettingsStore;
264    use std::{env, path::Path, process::Stdio};
265    use util::path;
266
267    fn init_test(cx: &mut TestAppContext) {
268        env_logger::init();
269        cx.update(|cx| {
270            let settings_store = SettingsStore::test(cx);
271            cx.set_global(settings_store);
272            Project::init_settings(cx);
273            language::init(cx);
274        });
275    }
276
277    #[gpui::test]
278    async fn test_gemini(cx: &mut TestAppContext) {
279        init_test(cx);
280
281        cx.executor().allow_parking();
282
283        let fs = FakeFs::new(cx.executor());
284        fs.insert_tree(
285            path!("/private/tmp"),
286            json!({"foo": "Lorem ipsum dolor", "bar": "bar", "baz": "baz"}),
287        )
288        .await;
289        let project = Project::test(fs, [path!("/private/tmp").as_ref()], cx).await;
290        let agent = gemini_agent(project.clone(), cx.to_async()).unwrap();
291        let thread_store = ThreadStore::load(agent, project, &mut cx.to_async())
292            .await
293            .unwrap();
294        let thread = thread_store
295            .update(cx, |thread_store, cx| {
296                assert_eq!(thread_store.threads().len(), 0);
297                thread_store.create_thread(cx)
298            })
299            .await
300            .unwrap();
301        thread
302            .update(cx, |thread, cx| {
303                thread.send(
304                    Message {
305                        role: Role::User,
306                        chunks: vec![
307                            "Read the '/private/tmp/foo' file and output all of its contents."
308                                .into(),
309                        ],
310                    },
311                    cx,
312                )
313            })
314            .await
315            .unwrap();
316
317        thread.read_with(cx, |thread, _| {
318            assert!(matches!(
319                thread.entries[0].content,
320                AgentThreadEntryContent::Message(Message {
321                    role: Role::User,
322                    ..
323                })
324            ));
325            assert!(
326                thread.entries().iter().any(|entry| {
327                    entry.content
328                        == AgentThreadEntryContent::ReadFile {
329                            path: "/private/tmp/foo".into(),
330                            content: "Lorem ipsum dolor".into(),
331                        }
332                }),
333                "Thread does not contain entry. Actual: {:?}",
334                thread.entries()
335            );
336        });
337    }
338
339    pub fn gemini_agent(project: Entity<Project>, mut cx: AsyncApp) -> Result<Arc<AcpAgent>> {
340        let cli_path =
341            Path::new(env!("CARGO_MANIFEST_DIR")).join("../../../gemini-cli/packages/cli");
342        let mut command = util::command::new_smol_command("node");
343        command
344            .arg(cli_path)
345            .arg("--acp")
346            .args(["--model", "gemini-2.5-flash"])
347            .current_dir("/private/tmp")
348            .stdin(Stdio::piped())
349            .stdout(Stdio::piped())
350            .stderr(Stdio::inherit())
351            .kill_on_drop(true);
352
353        if let Ok(gemini_key) = std::env::var("GEMINI_API_KEY") {
354            command.env("GEMINI_API_KEY", gemini_key);
355        }
356
357        let child = command.spawn().unwrap();
358
359        Ok(AcpAgent::stdio(child, project, &mut cx))
360    }
361}