Harden remote pairing, transfer and lifecycle recovery

Author: AnthonyRonning

apps/maple-agent/Cargo.lock

@@ -26,6 +26,7 @@ desktop = [
     "dep:spellbook",
     "dep:wayland-client",
     "dep:tray-icon",
+    "dep:zbus",
 ]
 # `maple-agent acp`: the Agent Client Protocol server on stdio.
 acp = ["maple-agent/acp"]
@@ -84,6 +85,7 @@ tower-http = { version = "0.6", features = ["cors"], optional = true }
 # `ui::decorations`. gpui cannot answer that question.
 [target.'cfg(target_os = "linux")'.dependencies]
 wayland-client = { version = "0.31", optional = true }
+zbus = { version = "=5.19.0", optional = true }
 
 [target.'cfg(any(target_os = "linux", target_os = "windows"))'.dependencies]
 tray-icon = { version = "=0.26.1", default-features = false, features = ["ksni"], optional = true }

apps/maple-agent/app/Cargo.toml

@@ -127,7 +127,6 @@ operations! {
     cancel_queued_message(user_id: (&str), session_id: (&str), queue_id: (&str)) -> AgentDesktopQueueSnapshot;
     begin_queued_message_edit(user_id: (&str), session_id: (&str), queue_id: (&str)) -> ();
     end_queued_message_edit(user_id: (&str), session_id: (&str), queue_id: (&str)) -> ();
-    send_message(user_id: (&str), request: (AgentSendMessageRequest)) -> String;
     read_image_attachment(user_id: (&str), session_id: (&str), attachment_id: (&str)) -> Vec<u8>;
     cancel_run(user_id: (&str), run_id: (&str)) -> ();
     available_model_ids(user_id: (&str)) -> Vec<String>;
@@ -260,6 +259,27 @@ impl AgentBackend {
             }
         }
     }
+    /// The composer retains this ID through uncertain delivery. Only receipt
+    /// queries, never another send, use it after a connection loss.
+    pub async fn send_message_identified(
+        &self,
+        user: &str,
+        request: AgentSendMessageRequest,
+        id: String,
+    ) -> Result<String, String> {
+        match &self.client {
+            BackendClient::Host(client) => {
+                client
+                    .call_identified("send_message", (user, request), Some(id))
+                    .await
+            }
+            #[cfg(test)]
+            BackendClient::InProcess(backend) => backend.send_message(user, request).await,
+        }
+    }
+    pub async fn send_outcome(&self, id: String) -> Result<crate::host::receipts::Outcome, String> {
+        self.call("send_outcome", id).await
+    }
     pub fn runtime_handle(&self) -> tokio::runtime::Handle {
         self.runtime.clone()
     }
@@ -389,6 +409,15 @@ async fn dispatch_special(
         serde_json::to_value(value).map_err(|_| "Could not encode host response".to_string())
     };
     match request.method.as_str() {
+        "send_message" => {
+            let (user, message): (String, AgentSendMessageRequest) =
+                serde_json::from_value(request.params)
+                    .map_err(|_| "Invalid host request parameters")?;
+            backend.authorize_account(&user)?;
+            serde_json::to_value(backend.send_message(&user, message).await?)
+                .map_err(|_| "Could not encode host response".into())
+        }
+
         "restore_auth" => {
             serde_json::to_value(backend.validate_persisted_auth().await).map_err(|e| e.to_string())
         }

apps/maple-agent/app/src/backend.rs

@@ -2689,6 +2689,60 @@ mod tests {
         assert!(!remote.stays_running());
     }
 
+    #[tokio::test]
+    async fn disabling_remote_does_not_wait_for_initializing_credentials() {
+        use std::{future::Future, task::Poll};
+        let root = tempfile::tempdir().unwrap();
+        let backend = Arc::new(LocalBackend::isolated_for_test(root.path()));
+        backend.persisted_auth.claim_and_persist(&auth_snapshot(
+            "fixture-user",
+            "fixture-session",
+            1,
+            "fixture-token",
+        ));
+        let (events, _) = tokio::sync::broadcast::channel(8);
+        let receipts = Arc::new(
+            crate::host::receipts::Receipts::open(&root.path().join("receipts.db")).unwrap(),
+        );
+        let remote = crate::remote::manager::Manager::new(backend.clone(), events, receipts);
+        let request = |method: &str, params| crate::host::protocol::Request {
+            id: 1,
+            method: method.into(),
+            params,
+            send_id: None,
+        };
+        remote
+            .dispatch(request("remote_enable", serde_json::json!([true])))
+            .await
+            .unwrap();
+        let validation = backend.restore_lock.lock().await;
+        let active = remote.active();
+        tokio::pin!(active);
+        std::future::poll_fn(|cx| {
+            assert!(active.as_mut().poll(cx).is_pending());
+            Poll::Ready(())
+        })
+        .await;
+        tokio::time::timeout(
+            std::time::Duration::from_secs(1),
+            remote.dispatch(request("remote_enable", serde_json::json!([false]))),
+        )
+        .await
+        .unwrap()
+        .unwrap();
+        assert!(!remote.stays_running());
+        let status = remote
+            .dispatch(request("remote_status", serde_json::json!([])))
+            .await
+            .unwrap();
+        assert_eq!(status["enabled"], false);
+        assert_eq!(status["ready"], false);
+        drop(validation);
+        // Invalid fixture credentials cannot create a node after the cancellation.
+        assert!(active.await.is_err());
+        remote.shutdown().await;
+    }
+
     #[tokio::test]
     async fn credential_validation_and_acp_attachment_do_not_block_local_operations() {
         use std::{future::Future, task::Poll};