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