fix: Clearer error message when connection is broken before messages are sent (#89)

* fix avoid hanging forever when send fails

* Fix formatting

* removed the unnecessary wait

* Return earlier

---------

Co-authored-by: Ben Brandt <benjamin.j.brandt@gmail.com>

Author: anaslimem

src/agent-client-protocol/src/lib.rs

@@ -81,7 +81,7 @@ impl Agent for ClientSideConnection {
             .request(
                 AGENT_METHOD_NAMES.initialize,
                 Some(ClientRequest::InitializeRequest(args)),
-            )
+            )?
             .await
     }
 
@@ -90,7 +90,7 @@ impl Agent for ClientSideConnection {
             .request::<Option<_>>(
                 AGENT_METHOD_NAMES.authenticate,
                 Some(ClientRequest::AuthenticateRequest(args)),
-            )
+            )?
             .await
             .map(Option::unwrap_or_default)
     }
@@ -101,7 +101,7 @@ impl Agent for ClientSideConnection {
             .request::<Option<_>>(
                 AGENT_METHOD_NAMES.logout,
                 Some(ClientRequest::LogoutRequest(args)),
-            )
+            )?
             .await
             .map(Option::unwrap_or_default)
     }
@@ -111,7 +111,7 @@ impl Agent for ClientSideConnection {
             .request(
                 AGENT_METHOD_NAMES.session_new,
                 Some(ClientRequest::NewSessionRequest(args)),
-            )
+            )?
             .await
     }
 
@@ -120,7 +120,7 @@ impl Agent for ClientSideConnection {
             .request::<Option<_>>(
                 AGENT_METHOD_NAMES.session_load,
                 Some(ClientRequest::LoadSessionRequest(args)),
-            )
+            )?
             .await
             .map(Option::unwrap_or_default)
     }
@@ -133,7 +133,7 @@ impl Agent for ClientSideConnection {
             .request(
                 AGENT_METHOD_NAMES.session_set_mode,
                 Some(ClientRequest::SetSessionModeRequest(args)),
-            )
+            )?
             .await
     }
 
@@ -142,7 +142,7 @@ impl Agent for ClientSideConnection {
             .request(
                 AGENT_METHOD_NAMES.session_prompt,
                 Some(ClientRequest::PromptRequest(args)),
-            )
+            )?
             .await
     }
 
@@ -162,7 +162,7 @@ impl Agent for ClientSideConnection {
             .request(
                 AGENT_METHOD_NAMES.session_set_model,
                 Some(ClientRequest::SetSessionModelRequest(args)),
-            )
+            )?
             .await
     }
 
@@ -171,7 +171,7 @@ impl Agent for ClientSideConnection {
             .request(
                 AGENT_METHOD_NAMES.session_list,
                 Some(ClientRequest::ListSessionsRequest(args)),
-            )
+            )?
             .await
     }
 
@@ -181,7 +181,7 @@ impl Agent for ClientSideConnection {
             .request(
                 AGENT_METHOD_NAMES.session_fork,
                 Some(ClientRequest::ForkSessionRequest(args)),
-            )
+            )?
             .await
     }
 
@@ -191,7 +191,7 @@ impl Agent for ClientSideConnection {
             .request(
                 AGENT_METHOD_NAMES.session_resume,
                 Some(ClientRequest::ResumeSessionRequest(args)),
-            )
+            )?
             .await
     }
 
@@ -201,7 +201,7 @@ impl Agent for ClientSideConnection {
             .request::<Option<_>>(
                 AGENT_METHOD_NAMES.session_close,
                 Some(ClientRequest::CloseSessionRequest(args)),
-            )
+            )?
             .await
             .map(Option::unwrap_or_default)
     }
@@ -214,7 +214,7 @@ impl Agent for ClientSideConnection {
             .request(
                 AGENT_METHOD_NAMES.session_set_config_option,
                 Some(ClientRequest::SetSessionConfigOptionRequest(args)),
-            )
+            )?
             .await
     }
 
@@ -223,7 +223,7 @@ impl Agent for ClientSideConnection {
             .request(
                 format!("_{}", args.method),
                 Some(ClientRequest::ExtMethodRequest(args)),
-            )
+            )?
             .await
     }
 
@@ -441,7 +441,7 @@ impl Client for AgentSideConnection {
             .request(
                 CLIENT_METHOD_NAMES.session_request_permission,
                 Some(AgentRequest::RequestPermissionRequest(args)),
-            )
+            )?
             .await
     }
 
@@ -450,7 +450,7 @@ impl Client for AgentSideConnection {
             .request::<Option<_>>(
                 CLIENT_METHOD_NAMES.fs_write_text_file,
                 Some(AgentRequest::WriteTextFileRequest(args)),
-            )
+            )?
             .await
             .map(Option::unwrap_or_default)
     }
@@ -460,7 +460,7 @@ impl Client for AgentSideConnection {
             .request(
                 CLIENT_METHOD_NAMES.fs_read_text_file,
                 Some(AgentRequest::ReadTextFileRequest(args)),
-            )
+            )?
             .await
     }
 
@@ -469,7 +469,7 @@ impl Client for AgentSideConnection {
             .request(
                 CLIENT_METHOD_NAMES.terminal_create,
                 Some(AgentRequest::CreateTerminalRequest(args)),
-            )
+            )?
             .await
     }
 
@@ -478,7 +478,7 @@ impl Client for AgentSideConnection {
             .request(
                 CLIENT_METHOD_NAMES.terminal_output,
                 Some(AgentRequest::TerminalOutputRequest(args)),
-            )
+            )?
             .await
     }
 
@@ -490,7 +490,7 @@ impl Client for AgentSideConnection {
             .request::<Option<_>>(
                 CLIENT_METHOD_NAMES.terminal_release,
                 Some(AgentRequest::ReleaseTerminalRequest(args)),
-            )
+            )?
             .await
             .map(Option::unwrap_or_default)
     }
@@ -503,7 +503,7 @@ impl Client for AgentSideConnection {
             .request(
                 CLIENT_METHOD_NAMES.terminal_wait_for_exit,
                 Some(AgentRequest::WaitForTerminalExitRequest(args)),
-            )
+            )?
             .await
     }
 
@@ -512,7 +512,7 @@ impl Client for AgentSideConnection {
             .request::<Option<_>>(
                 CLIENT_METHOD_NAMES.terminal_kill,
                 Some(AgentRequest::KillTerminalRequest(args)),
-            )
+            )?
             .await
             .map(Option::unwrap_or_default)
     }
@@ -529,7 +529,7 @@ impl Client for AgentSideConnection {
             .request(
                 format!("_{}", args.method),
                 Some(AgentRequest::ExtMethodRequest(args)),
-            )
+            )?
             .await
     }
 

src/agent-client-protocol/src/rpc.rs

@@ -113,7 +113,7 @@ where
         &self,
         method: impl Into<Arc<str>>,
         params: Option<Remote::InRequest>,
-    ) -> impl Future<Output = Result<Out>> {
+    ) -> Result<impl Future<Output = Result<Out>>> {
         let (tx, rx) = oneshot::channel();
         let id = self.next_id.fetch_add(1, Ordering::SeqCst);
         let id = RequestId::Number(id);
@@ -139,16 +139,19 @@ where
             .is_err()
         {
             self.pending_responses.lock().unwrap().remove(&id);
+            return Err(
+                Error::internal_error().data("connection closed before request could be sent")
+            );
         }
-        async move {
+        Ok(async move {
             let result = rx
                 .await
                 .map_err(|_| Error::internal_error().data("server shut down unexpectedly"))??
                 .downcast::<Out>()
                 .map_err(|_| Error::internal_error().data("failed to deserialize response"))?;
 
             Ok(*result)
-        }
+        })
     }
 
     async fn handle_io(