acp.rs

  1use crate::AgentServerCommand;
  2use acp_thread::AgentConnection;
  3use acp_tools::AcpConnectionRegistry;
  4use action_log::ActionLog;
  5use agent_client_protocol::{self as acp, Agent as _, ErrorCode};
  6use anyhow::anyhow;
  7use collections::HashMap;
  8use futures::AsyncBufReadExt as _;
  9use futures::channel::oneshot;
 10use futures::io::BufReader;
 11use project::Project;
 12use serde::Deserialize;
 13use std::{any::Any, cell::RefCell};
 14use std::{path::Path, rc::Rc};
 15use thiserror::Error;
 16
 17use anyhow::{Context as _, Result};
 18use gpui::{App, AppContext as _, AsyncApp, Entity, SharedString, Task, WeakEntity};
 19
 20use acp_thread::{AcpThread, AuthRequired, LoadError};
 21
 22#[derive(Debug, Error)]
 23#[error("Unsupported version")]
 24pub struct UnsupportedVersion;
 25
 26pub struct AcpConnection {
 27    server_name: SharedString,
 28    connection: Rc<acp::ClientSideConnection>,
 29    sessions: Rc<RefCell<HashMap<acp::SessionId, AcpSession>>>,
 30    auth_methods: Vec<acp::AuthMethod>,
 31    prompt_capabilities: acp::PromptCapabilities,
 32    _io_task: Task<Result<()>>,
 33}
 34
 35pub struct AcpSession {
 36    thread: WeakEntity<AcpThread>,
 37    suppress_abort_err: bool,
 38}
 39
 40pub async fn connect(
 41    server_name: SharedString,
 42    command: AgentServerCommand,
 43    root_dir: &Path,
 44    cx: &mut AsyncApp,
 45) -> Result<Rc<dyn AgentConnection>> {
 46    let conn = AcpConnection::stdio(server_name, command.clone(), root_dir, cx).await?;
 47    Ok(Rc::new(conn) as _)
 48}
 49
 50const MINIMUM_SUPPORTED_VERSION: acp::ProtocolVersion = acp::V1;
 51
 52impl AcpConnection {
 53    pub async fn stdio(
 54        server_name: SharedString,
 55        command: AgentServerCommand,
 56        root_dir: &Path,
 57        cx: &mut AsyncApp,
 58    ) -> Result<Self> {
 59        let mut child = util::command::new_smol_command(command.path)
 60            .args(command.args.iter().map(|arg| arg.as_str()))
 61            .envs(command.env.iter().flatten())
 62            .current_dir(root_dir)
 63            .stdin(std::process::Stdio::piped())
 64            .stdout(std::process::Stdio::piped())
 65            .stderr(std::process::Stdio::piped())
 66            .kill_on_drop(true)
 67            .spawn()?;
 68
 69        let stdout = child.stdout.take().context("Failed to take stdout")?;
 70        let stdin = child.stdin.take().context("Failed to take stdin")?;
 71        let stderr = child.stderr.take().context("Failed to take stderr")?;
 72        log::trace!("Spawned (pid: {})", child.id());
 73
 74        let sessions = Rc::new(RefCell::new(HashMap::default()));
 75
 76        let client = ClientDelegate {
 77            sessions: sessions.clone(),
 78            cx: cx.clone(),
 79        };
 80        let (connection, io_task) = acp::ClientSideConnection::new(client, stdin, stdout, {
 81            let foreground_executor = cx.foreground_executor().clone();
 82            move |fut| {
 83                foreground_executor.spawn(fut).detach();
 84            }
 85        });
 86
 87        let io_task = cx.background_spawn(io_task);
 88
 89        cx.background_spawn(async move {
 90            let mut stderr = BufReader::new(stderr);
 91            let mut line = String::new();
 92            while let Ok(n) = stderr.read_line(&mut line).await
 93                && n > 0
 94            {
 95                log::warn!("agent stderr: {}", &line);
 96                line.clear();
 97            }
 98        })
 99        .detach();
100
101        cx.spawn({
102            let sessions = sessions.clone();
103            async move |cx| {
104                let status = child.status().await?;
105
106                for session in sessions.borrow().values() {
107                    session
108                        .thread
109                        .update(cx, |thread, cx| {
110                            thread.emit_load_error(LoadError::Exited { status }, cx)
111                        })
112                        .ok();
113                }
114
115                anyhow::Ok(())
116            }
117        })
118        .detach();
119
120        let connection = Rc::new(connection);
121
122        cx.update(|cx| {
123            AcpConnectionRegistry::default_global(cx).update(cx, |registry, cx| {
124                registry.set_active_connection(server_name.clone(), &connection, cx)
125            });
126        })?;
127
128        let response = connection
129            .initialize(acp::InitializeRequest {
130                protocol_version: acp::VERSION,
131                client_capabilities: acp::ClientCapabilities {
132                    fs: acp::FileSystemCapability {
133                        read_text_file: true,
134                        write_text_file: true,
135                    },
136                },
137            })
138            .await?;
139
140        if response.protocol_version < MINIMUM_SUPPORTED_VERSION {
141            return Err(UnsupportedVersion.into());
142        }
143
144        Ok(Self {
145            auth_methods: response.auth_methods,
146            connection,
147            server_name,
148            sessions,
149            prompt_capabilities: response.agent_capabilities.prompt_capabilities,
150            _io_task: io_task,
151        })
152    }
153
154    pub fn prompt_capabilities(&self) -> &acp::PromptCapabilities {
155        &self.prompt_capabilities
156    }
157}
158
159impl AgentConnection for AcpConnection {
160    fn new_thread(
161        self: Rc<Self>,
162        project: Entity<Project>,
163        cwd: &Path,
164        cx: &mut App,
165    ) -> Task<Result<Entity<AcpThread>>> {
166        let conn = self.connection.clone();
167        let sessions = self.sessions.clone();
168        let cwd = cwd.to_path_buf();
169        let context_server_store = project.read(cx).context_server_store().read(cx);
170        let mcp_servers = context_server_store
171            .configured_server_ids()
172            .iter()
173            .filter_map(|id| {
174                let configuration = context_server_store.configuration_for_server(id)?;
175                let command = configuration.command();
176                Some(acp::McpServer {
177                    name: id.0.to_string(),
178                    command: command.path.clone(),
179                    args: command.args.clone(),
180                    env: if let Some(env) = command.env.as_ref() {
181                        env.iter()
182                            .map(|(name, value)| acp::EnvVariable {
183                                name: name.clone(),
184                                value: value.clone(),
185                            })
186                            .collect()
187                    } else {
188                        vec![]
189                    },
190                })
191            })
192            .collect();
193
194        cx.spawn(async move |cx| {
195            let response = conn
196                .new_session(acp::NewSessionRequest { mcp_servers, cwd })
197                .await
198                .map_err(|err| {
199                    if err.code == acp::ErrorCode::AUTH_REQUIRED.code {
200                        let mut error = AuthRequired::new();
201
202                        if err.message != acp::ErrorCode::AUTH_REQUIRED.message {
203                            error = error.with_description(err.message);
204                        }
205
206                        anyhow!(error)
207                    } else {
208                        anyhow!(err)
209                    }
210                })?;
211
212            let session_id = response.session_id;
213            let action_log = cx.new(|_| ActionLog::new(project.clone()))?;
214            let thread = cx.new(|cx| {
215                AcpThread::new(
216                    self.server_name.clone(),
217                    self.clone(),
218                    project,
219                    action_log,
220                    session_id.clone(),
221                    // ACP doesn't currently support per-session prompt capabilities or changing capabilities dynamically.
222                    watch::Receiver::constant(self.prompt_capabilities),
223                    cx,
224                )
225            })?;
226
227            let session = AcpSession {
228                thread: thread.downgrade(),
229                suppress_abort_err: false,
230            };
231            sessions.borrow_mut().insert(session_id, session);
232
233            Ok(thread)
234        })
235    }
236
237    fn auth_methods(&self) -> &[acp::AuthMethod] {
238        &self.auth_methods
239    }
240
241    fn authenticate(&self, method_id: acp::AuthMethodId, cx: &mut App) -> Task<Result<()>> {
242        let conn = self.connection.clone();
243        cx.foreground_executor().spawn(async move {
244            let result = conn
245                .authenticate(acp::AuthenticateRequest {
246                    method_id: method_id.clone(),
247                })
248                .await?;
249
250            Ok(result)
251        })
252    }
253
254    fn prompt(
255        &self,
256        _id: Option<acp_thread::UserMessageId>,
257        params: acp::PromptRequest,
258        cx: &mut App,
259    ) -> Task<Result<acp::PromptResponse>> {
260        let conn = self.connection.clone();
261        let sessions = self.sessions.clone();
262        let session_id = params.session_id.clone();
263        cx.foreground_executor().spawn(async move {
264            let result = conn.prompt(params).await;
265
266            let mut suppress_abort_err = false;
267
268            if let Some(session) = sessions.borrow_mut().get_mut(&session_id) {
269                suppress_abort_err = session.suppress_abort_err;
270                session.suppress_abort_err = false;
271            }
272
273            match result {
274                Ok(response) => Ok(response),
275                Err(err) => {
276                    if err.code != ErrorCode::INTERNAL_ERROR.code {
277                        anyhow::bail!(err)
278                    }
279
280                    let Some(data) = &err.data else {
281                        anyhow::bail!(err)
282                    };
283
284                    // Temporary workaround until the following PR is generally available:
285                    // https://github.com/google-gemini/gemini-cli/pull/6656
286
287                    #[derive(Deserialize)]
288                    #[serde(deny_unknown_fields)]
289                    struct ErrorDetails {
290                        details: Box<str>,
291                    }
292
293                    match serde_json::from_value(data.clone()) {
294                        Ok(ErrorDetails { details }) => {
295                            if suppress_abort_err
296                                && (details.contains("This operation was aborted")
297                                    || details.contains("The user aborted a request"))
298                            {
299                                Ok(acp::PromptResponse {
300                                    stop_reason: acp::StopReason::Cancelled,
301                                })
302                            } else {
303                                Err(anyhow!(details))
304                            }
305                        }
306                        Err(_) => Err(anyhow!(err)),
307                    }
308                }
309            }
310        })
311    }
312
313    fn cancel(&self, session_id: &acp::SessionId, cx: &mut App) {
314        if let Some(session) = self.sessions.borrow_mut().get_mut(session_id) {
315            session.suppress_abort_err = true;
316        }
317        let conn = self.connection.clone();
318        let params = acp::CancelNotification {
319            session_id: session_id.clone(),
320        };
321        cx.foreground_executor()
322            .spawn(async move { conn.cancel(params).await })
323            .detach();
324    }
325
326    fn into_any(self: Rc<Self>) -> Rc<dyn Any> {
327        self
328    }
329}
330
331struct ClientDelegate {
332    sessions: Rc<RefCell<HashMap<acp::SessionId, AcpSession>>>,
333    cx: AsyncApp,
334}
335
336impl acp::Client for ClientDelegate {
337    async fn request_permission(
338        &self,
339        arguments: acp::RequestPermissionRequest,
340    ) -> Result<acp::RequestPermissionResponse, acp::Error> {
341        let cx = &mut self.cx.clone();
342        let rx = self
343            .sessions
344            .borrow()
345            .get(&arguments.session_id)
346            .context("Failed to get session")?
347            .thread
348            .update(cx, |thread, cx| {
349                thread.request_tool_call_authorization(arguments.tool_call, arguments.options, cx)
350            })?;
351
352        let result = rx?.await;
353
354        let outcome = match result {
355            Ok(option) => acp::RequestPermissionOutcome::Selected { option_id: option },
356            Err(oneshot::Canceled) => acp::RequestPermissionOutcome::Cancelled,
357        };
358
359        Ok(acp::RequestPermissionResponse { outcome })
360    }
361
362    async fn write_text_file(
363        &self,
364        arguments: acp::WriteTextFileRequest,
365    ) -> Result<(), acp::Error> {
366        let cx = &mut self.cx.clone();
367        let task = self
368            .sessions
369            .borrow()
370            .get(&arguments.session_id)
371            .context("Failed to get session")?
372            .thread
373            .update(cx, |thread, cx| {
374                thread.write_text_file(arguments.path, arguments.content, cx)
375            })?;
376
377        task.await?;
378
379        Ok(())
380    }
381
382    async fn read_text_file(
383        &self,
384        arguments: acp::ReadTextFileRequest,
385    ) -> Result<acp::ReadTextFileResponse, acp::Error> {
386        let cx = &mut self.cx.clone();
387        let task = self
388            .sessions
389            .borrow()
390            .get(&arguments.session_id)
391            .context("Failed to get session")?
392            .thread
393            .update(cx, |thread, cx| {
394                thread.read_text_file(arguments.path, arguments.line, arguments.limit, false, cx)
395            })?;
396
397        let content = task.await?;
398
399        Ok(acp::ReadTextFileResponse { content })
400    }
401
402    async fn session_notification(
403        &self,
404        notification: acp::SessionNotification,
405    ) -> Result<(), acp::Error> {
406        let cx = &mut self.cx.clone();
407        let sessions = self.sessions.borrow();
408        let session = sessions
409            .get(&notification.session_id)
410            .context("Failed to get session")?;
411
412        session.thread.update(cx, |thread, cx| {
413            thread.handle_session_update(notification.update, cx)
414        })??;
415
416        Ok(())
417    }
418}