| @@ -1,34 +1,391 @@ |
| 1 | use std::{path::Path, rc::Rc}; | | |
| 2 | | | |
| 3 | use crate::AgentServerCommand; | 1 | use crate::AgentServerCommand; |
| 4 | use acp_thread::AgentConnection; | 2 | use acp_thread::AgentConnection; |
| 5 | use anyhow::Result; | 3 | use acp_tools::AcpConnectionRegistry; |
| 6 | use gpui::AsyncApp; | 4 | use action_log::ActionLog; |
| | 5 | use agent_client_protocol::{self as acp, Agent as _, ErrorCode}; |
| | 6 | use anyhow::anyhow; |
| | 7 | use collections::HashMap; |
| | 8 | use futures::AsyncBufReadExt as _; |
| | 9 | use futures::channel::oneshot; |
| | 10 | use futures::io::BufReader; |
| | 11 | use project::Project; |
| | 12 | use serde::Deserialize; |
| | 13 | use std::{any::Any, cell::RefCell}; |
| | 14 | use std::{path::Path, rc::Rc}; |
| 7 | use thiserror::Error; | 15 | use thiserror::Error; |
| 8 | | 16 | |
| 9 | mod v0; | 17 | use anyhow::{Context as _, Result}; |
| 10 | mod v1; | 18 | use gpui::{App, AppContext as _, AsyncApp, Entity, Task, WeakEntity}; |
| | 19 | |
| | 20 | use acp_thread::{AcpThread, AuthRequired, LoadError}; |
| 11 | | 21 | |
| 12 | #[derive(Debug, Error)] | 22 | #[derive(Debug, Error)] |
| 13 | #[error("Unsupported version")] | 23 | #[error("Unsupported version")] |
| 14 | pub struct UnsupportedVersion; | 24 | pub struct UnsupportedVersion; |
| 15 | | 25 | |
| | 26 | pub struct AcpConnection { |
| | 27 | server_name: &'static str, |
| | 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 | |
| | 35 | pub struct AcpSession { |
| | 36 | thread: WeakEntity<AcpThread>, |
| | 37 | suppress_abort_err: bool, |
| | 38 | } |
| | 39 | |
| 16 | pub async fn connect( | 40 | pub async fn connect( |
| 17 | server_name: &'static str, | 41 | server_name: &'static str, |
| 18 | command: AgentServerCommand, | 42 | command: AgentServerCommand, |
| 19 | root_dir: &Path, | 43 | root_dir: &Path, |
| 20 | cx: &mut AsyncApp, | 44 | cx: &mut AsyncApp, |
| 21 | ) -> Result<Rc<dyn AgentConnection>> { | 45 | ) -> Result<Rc<dyn AgentConnection>> { |
| 22 | let conn = v1::AcpConnection::stdio(server_name, command.clone(), root_dir, cx).await; | 46 | let conn = AcpConnection::stdio(server_name, command.clone(), root_dir, cx).await?; |
| 23 | | 47 | Ok(Rc::new(conn) as _) |
| 24 | match conn { | 48 | } |
| 25 | Ok(conn) => Ok(Rc::new(conn) as _), | 49 | |
| 26 | Err(err) if err.is::<UnsupportedVersion>() => { | 50 | const MINIMUM_SUPPORTED_VERSION: acp::ProtocolVersion = acp::V1; |
| 27 | // Consider re-using initialize response and subprocess when adding another version here | 51 | |
| 28 | let conn: Rc<dyn AgentConnection> = | 52 | impl AcpConnection { |
| 29 | Rc::new(v0::AcpConnection::stdio(server_name, command, root_dir, cx).await?); | 53 | pub async fn stdio( |
| 30 | Ok(conn) | 54 | server_name: &'static str, |
| | 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, &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 | |
| | 155 | impl AgentConnection for AcpConnection { |
| | 156 | fn new_thread( |
| | 157 | self: Rc<Self>, |
| | 158 | project: Entity<Project>, |
| | 159 | cwd: &Path, |
| | 160 | cx: &mut App, |
| | 161 | ) -> Task<Result<Entity<AcpThread>>> { |
| | 162 | let conn = self.connection.clone(); |
| | 163 | let sessions = self.sessions.clone(); |
| | 164 | let cwd = cwd.to_path_buf(); |
| | 165 | cx.spawn(async move |cx| { |
| | 166 | let response = conn |
| | 167 | .new_session(acp::NewSessionRequest { |
| | 168 | mcp_servers: vec![], |
| | 169 | cwd, |
| | 170 | }) |
| | 171 | .await |
| | 172 | .map_err(|err| { |
| | 173 | if err.code == acp::ErrorCode::AUTH_REQUIRED.code { |
| | 174 | let mut error = AuthRequired::new(); |
| | 175 | |
| | 176 | if err.message != acp::ErrorCode::AUTH_REQUIRED.message { |
| | 177 | error = error.with_description(err.message); |
| | 178 | } |
| | 179 | |
| | 180 | anyhow!(error) |
| | 181 | } else { |
| | 182 | anyhow!(err) |
| | 183 | } |
| | 184 | })?; |
| | 185 | |
| | 186 | let session_id = response.session_id; |
| | 187 | let action_log = cx.new(|_| ActionLog::new(project.clone()))?; |
| | 188 | let thread = cx.new(|_cx| { |
| | 189 | AcpThread::new( |
| | 190 | self.server_name, |
| | 191 | self.clone(), |
| | 192 | project, |
| | 193 | action_log, |
| | 194 | session_id.clone(), |
| | 195 | ) |
| | 196 | })?; |
| | 197 | |
| | 198 | let session = AcpSession { |
| | 199 | thread: thread.downgrade(), |
| | 200 | suppress_abort_err: false, |
| | 201 | }; |
| | 202 | sessions.borrow_mut().insert(session_id, session); |
| | 203 | |
| | 204 | Ok(thread) |
| | 205 | }) |
| | 206 | } |
| | 207 | |
| | 208 | fn auth_methods(&self) -> &[acp::AuthMethod] { |
| | 209 | &self.auth_methods |
| | 210 | } |
| | 211 | |
| | 212 | fn authenticate(&self, method_id: acp::AuthMethodId, cx: &mut App) -> Task<Result<()>> { |
| | 213 | let conn = self.connection.clone(); |
| | 214 | cx.foreground_executor().spawn(async move { |
| | 215 | let result = conn |
| | 216 | .authenticate(acp::AuthenticateRequest { |
| | 217 | method_id: method_id.clone(), |
| | 218 | }) |
| | 219 | .await?; |
| | 220 | |
| | 221 | Ok(result) |
| | 222 | }) |
| | 223 | } |
| | 224 | |
| | 225 | fn prompt( |
| | 226 | &self, |
| | 227 | _id: Option<acp_thread::UserMessageId>, |
| | 228 | params: acp::PromptRequest, |
| | 229 | cx: &mut App, |
| | 230 | ) -> Task<Result<acp::PromptResponse>> { |
| | 231 | let conn = self.connection.clone(); |
| | 232 | let sessions = self.sessions.clone(); |
| | 233 | let session_id = params.session_id.clone(); |
| | 234 | cx.foreground_executor().spawn(async move { |
| | 235 | let result = conn.prompt(params).await; |
| | 236 | |
| | 237 | let mut suppress_abort_err = false; |
| | 238 | |
| | 239 | if let Some(session) = sessions.borrow_mut().get_mut(&session_id) { |
| | 240 | suppress_abort_err = session.suppress_abort_err; |
| | 241 | session.suppress_abort_err = false; |
| | 242 | } |
| | 243 | |
| | 244 | match result { |
| | 245 | Ok(response) => Ok(response), |
| | 246 | Err(err) => { |
| | 247 | if err.code != ErrorCode::INTERNAL_ERROR.code { |
| | 248 | anyhow::bail!(err) |
| | 249 | } |
| | 250 | |
| | 251 | let Some(data) = &err.data else { |
| | 252 | anyhow::bail!(err) |
| | 253 | }; |
| | 254 | |
| | 255 | // Temporary workaround until the following PR is generally available: |
| | 256 | // https://github.com/google-gemini/gemini-cli/pull/6656 |
| | 257 | |
| | 258 | #[derive(Deserialize)] |
| | 259 | #[serde(deny_unknown_fields)] |
| | 260 | struct ErrorDetails { |
| | 261 | details: Box<str>, |
| | 262 | } |
| | 263 | |
| | 264 | match serde_json::from_value(data.clone()) { |
| | 265 | Ok(ErrorDetails { details }) => { |
| | 266 | if suppress_abort_err && details.contains("This operation was aborted") |
| | 267 | { |
| | 268 | Ok(acp::PromptResponse { |
| | 269 | stop_reason: acp::StopReason::Cancelled, |
| | 270 | }) |
| | 271 | } else { |
| | 272 | Err(anyhow!(details)) |
| | 273 | } |
| | 274 | } |
| | 275 | Err(_) => Err(anyhow!(err)), |
| | 276 | } |
| | 277 | } |
| | 278 | } |
| | 279 | }) |
| | 280 | } |
| | 281 | |
| | 282 | fn prompt_capabilities(&self) -> acp::PromptCapabilities { |
| | 283 | self.prompt_capabilities |
| | 284 | } |
| | 285 | |
| | 286 | fn cancel(&self, session_id: &acp::SessionId, cx: &mut App) { |
| | 287 | if let Some(session) = self.sessions.borrow_mut().get_mut(session_id) { |
| | 288 | session.suppress_abort_err = true; |
| 31 | } | 289 | } |
| 32 | Err(err) => Err(err), | 290 | let conn = self.connection.clone(); |
| | 291 | let params = acp::CancelNotification { |
| | 292 | session_id: session_id.clone(), |
| | 293 | }; |
| | 294 | cx.foreground_executor() |
| | 295 | .spawn(async move { conn.cancel(params).await }) |
| | 296 | .detach(); |
| | 297 | } |
| | 298 | |
| | 299 | fn into_any(self: Rc<Self>) -> Rc<dyn Any> { |
| | 300 | self |
| | 301 | } |
| | 302 | } |
| | 303 | |
| | 304 | struct ClientDelegate { |
| | 305 | sessions: Rc<RefCell<HashMap<acp::SessionId, AcpSession>>>, |
| | 306 | cx: AsyncApp, |
| | 307 | } |
| | 308 | |
| | 309 | impl acp::Client for ClientDelegate { |
| | 310 | async fn request_permission( |
| | 311 | &self, |
| | 312 | arguments: acp::RequestPermissionRequest, |
| | 313 | ) -> Result<acp::RequestPermissionResponse, acp::Error> { |
| | 314 | let cx = &mut self.cx.clone(); |
| | 315 | let rx = self |
| | 316 | .sessions |
| | 317 | .borrow() |
| | 318 | .get(&arguments.session_id) |
| | 319 | .context("Failed to get session")? |
| | 320 | .thread |
| | 321 | .update(cx, |thread, cx| { |
| | 322 | thread.request_tool_call_authorization(arguments.tool_call, arguments.options, cx) |
| | 323 | })?; |
| | 324 | |
| | 325 | let result = rx?.await; |
| | 326 | |
| | 327 | let outcome = match result { |
| | 328 | Ok(option) => acp::RequestPermissionOutcome::Selected { option_id: option }, |
| | 329 | Err(oneshot::Canceled) => acp::RequestPermissionOutcome::Cancelled, |
| | 330 | }; |
| | 331 | |
| | 332 | Ok(acp::RequestPermissionResponse { outcome }) |
| | 333 | } |
| | 334 | |
| | 335 | async fn write_text_file( |
| | 336 | &self, |
| | 337 | arguments: acp::WriteTextFileRequest, |
| | 338 | ) -> Result<(), acp::Error> { |
| | 339 | let cx = &mut self.cx.clone(); |
| | 340 | let task = self |
| | 341 | .sessions |
| | 342 | .borrow() |
| | 343 | .get(&arguments.session_id) |
| | 344 | .context("Failed to get session")? |
| | 345 | .thread |
| | 346 | .update(cx, |thread, cx| { |
| | 347 | thread.write_text_file(arguments.path, arguments.content, cx) |
| | 348 | })?; |
| | 349 | |
| | 350 | task.await?; |
| | 351 | |
| | 352 | Ok(()) |
| | 353 | } |
| | 354 | |
| | 355 | async fn read_text_file( |
| | 356 | &self, |
| | 357 | arguments: acp::ReadTextFileRequest, |
| | 358 | ) -> Result<acp::ReadTextFileResponse, acp::Error> { |
| | 359 | let cx = &mut self.cx.clone(); |
| | 360 | let task = self |
| | 361 | .sessions |
| | 362 | .borrow() |
| | 363 | .get(&arguments.session_id) |
| | 364 | .context("Failed to get session")? |
| | 365 | .thread |
| | 366 | .update(cx, |thread, cx| { |
| | 367 | thread.read_text_file(arguments.path, arguments.line, arguments.limit, false, cx) |
| | 368 | })?; |
| | 369 | |
| | 370 | let content = task.await?; |
| | 371 | |
| | 372 | Ok(acp::ReadTextFileResponse { content }) |
| | 373 | } |
| | 374 | |
| | 375 | async fn session_notification( |
| | 376 | &self, |
| | 377 | notification: acp::SessionNotification, |
| | 378 | ) -> Result<(), acp::Error> { |
| | 379 | let cx = &mut self.cx.clone(); |
| | 380 | let sessions = self.sessions.borrow(); |
| | 381 | let session = sessions |
| | 382 | .get(¬ification.session_id) |
| | 383 | .context("Failed to get session")?; |
| | 384 | |
| | 385 | session.thread.update(cx, |thread, cx| { |
| | 386 | thread.handle_session_update(notification.update, cx) |
| | 387 | })??; |
| | 388 | |
| | 389 | Ok(()) |
| 33 | } | 390 | } |
| 34 | } | 391 | } |