acp: Supress gemini aborted errors (#36633)

Agus Zubiaga created 3 months ago

This PR adds a temporary workaround to supress "Aborted" errors from
Gemini when cancelling generation. This won't be needed once
https://github.com/google-gemini/gemini-cli/pull/6656 is generally
available.

Release Notes:

- N/A

Change summary

Cargo.lock                         |  4 +-
Cargo.toml                         |  2 
crates/agent_servers/src/acp/v1.rs | 61 +++++++++++++++++++++++++++++--
3 files changed, 60 insertions(+), 7 deletions(-)

Detailed changes

Cargo.lock

@@ -171,9 +171,9 @@ dependencies = [
 
 [[package]]
 name = "agent-client-protocol"
-version = "0.0.28"
+version = "0.0.29"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "4c887e795097665ab95119580534e7cc1335b59e1a7fec296943e534b970f4ed"
+checksum = "89a2cd7e0bd2bb7ed27687cfcf6561b91542c1ce23e52fd54ee59b7568c9bd84"
 dependencies = [
  "anyhow",
  "futures 0.3.31",

Cargo.toml

@@ -423,7 +423,7 @@ zlog_settings = { path = "crates/zlog_settings" }
 #
 
 agentic-coding-protocol = "0.0.10"
-agent-client-protocol = "0.0.28"
+agent-client-protocol = "0.0.29"
 aho-corasick = "1.1"
 alacritty_terminal = { git = "https://github.com/zed-industries/alacritty.git", branch = "add-hush-login-flag" }
 any_vec = "0.14"

crates/agent_servers/src/acp/v1.rs

@@ -1,11 +1,12 @@
 use action_log::ActionLog;
-use agent_client_protocol::{self as acp, Agent as _};
+use agent_client_protocol::{self as acp, Agent as _, ErrorCode};
 use anyhow::anyhow;
 use collections::HashMap;
 use futures::AsyncBufReadExt as _;
 use futures::channel::oneshot;
 use futures::io::BufReader;
 use project::Project;
+use serde::Deserialize;
 use std::path::Path;
 use std::rc::Rc;
 use std::{any::Any, cell::RefCell};
@@ -27,6 +28,7 @@ pub struct AcpConnection {
 
 pub struct AcpSession {
     thread: WeakEntity<AcpThread>,
+    pending_cancel: bool,
 }
 
 const MINIMUM_SUPPORTED_VERSION: acp::ProtocolVersion = acp::V1;
@@ -171,6 +173,7 @@ impl AgentConnection for AcpConnection {
 
             let session = AcpSession {
                 thread: thread.downgrade(),
+                pending_cancel: false,
             };
             sessions.borrow_mut().insert(session_id, session);
 
@@ -202,9 +205,48 @@ impl AgentConnection for AcpConnection {
         cx: &mut App,
     ) -> Task<Result<acp::PromptResponse>> {
         let conn = self.connection.clone();
+        let sessions = self.sessions.clone();
+        let session_id = params.session_id.clone();
         cx.foreground_executor().spawn(async move {
-            let response = conn.prompt(params).await?;
-            Ok(response)
+            match conn.prompt(params).await {
+                Ok(response) => Ok(response),
+                Err(err) => {
+                    if err.code != ErrorCode::INTERNAL_ERROR.code {
+                        anyhow::bail!(err)
+                    }
+
+                    let Some(data) = &err.data else {
+                        anyhow::bail!(err)
+                    };
+
+                    // Temporary workaround until the following PR is generally available:
+                    // https://github.com/google-gemini/gemini-cli/pull/6656
+
+                    #[derive(Deserialize)]
+                    #[serde(deny_unknown_fields)]
+                    struct ErrorDetails {
+                        details: Box<str>,
+                    }
+
+                    match serde_json::from_value(data.clone()) {
+                        Ok(ErrorDetails { details }) => {
+                            if sessions
+                                .borrow()
+                                .get(&session_id)
+                                .is_some_and(|session| session.pending_cancel)
+                                && details.contains("This operation was aborted")
+                            {
+                                Ok(acp::PromptResponse {
+                                    stop_reason: acp::StopReason::Canceled,
+                                })
+                            } else {
+                                Err(anyhow!(details))
+                            }
+                        }
+                        Err(_) => Err(anyhow!(err)),
+                    }
+                }
+            }
         })
     }
 
@@ -213,12 +255,23 @@ impl AgentConnection for AcpConnection {
     }
 
     fn cancel(&self, session_id: &acp::SessionId, cx: &mut App) {
+        if let Some(session) = self.sessions.borrow_mut().get_mut(session_id) {
+            session.pending_cancel = true;
+        }
         let conn = self.connection.clone();
         let params = acp::CancelNotification {
             session_id: session_id.clone(),
         };
+        let sessions = self.sessions.clone();
+        let session_id = session_id.clone();
         cx.foreground_executor()
-            .spawn(async move { conn.cancel(params).await })
+            .spawn(async move {
+                let resp = conn.cancel(params).await;
+                if let Some(session) = sessions.borrow_mut().get_mut(&session_id) {
+                    session.pending_cancel = false;
+                }
+                resp
+            })
             .detach();
     }