feat: enhance bot functionality with message sending capability and state management

Author: krypton-byte

src/client.rs

@@ -1,4 +1,4 @@
-use std::sync::{Arc, Once};
+use std::sync::{Arc, Mutex, Once};
 use std::future::Future;
 use std::pin::Pin;
 use pyo3::{Bound, PyAny, pyclass, pymethods};
@@ -8,16 +8,20 @@ use pyo3_async_runtimes::tokio::{future_into_py_with_locals, get_current_locals,
 use tokio::runtime;
 use tokio::time::{Duration, interval};
 use wacore::types::events::Event;
+use whatsapp_rust::Client;
 use whatsapp_rust::bot::Bot;
 use whatsapp_rust::store::Backend;
 use whatsapp_rust_tokio_transport::TokioWebSocketTransportFactory;
 use whatsapp_rust_ureq_http_client::UreqHttpClient;
+use waproto::whatsapp::Message as WhatsappMessage;
+use prost::Message;
 use tokio::signal;
 use tracing::{debug, error, info, warn};
 use tracing_subscriber::EnvFilter;
 
 use crate::backend::{SqliteBackend, BackendBase};
-use crate::events::{Dispatcher, Message, PairingQrCode};
+use crate::events::{Dispatcher, Message as WAMessage, PairingQrCode};
+use crate::types::JID;
 
 static LOG_INIT: Once = Once::new();
 
@@ -36,13 +40,15 @@ fn init_logging() {
 pub struct Tryx {
     backend: Arc<dyn Backend>,
     handlers: Py<Dispatcher>,
+    bot: Arc<Mutex<Option<Arc<Client>>>>,
 }
 
 impl Tryx {
     async fn run_bot(
         backend: Arc<dyn Backend>,
         handlers: Py<Dispatcher>,
         locals: Option<TaskLocals>,
+        bot_state: Arc<Mutex<Option<Arc<Client>>>>,
     ) -> PyResult<()> {
         info!("building WhatsApp bot");
         let mut bot = Bot::builder()
@@ -113,7 +119,7 @@ impl Tryx {
                                 debug!(handler_index = idx, message_id = %info.id, "calling message Python callback");
                                 let locals = locals.clone();
                                 let py_future = Python::attach(|py| -> PyResult<_> {
-                                    let payload = Py::new(py, Message::new(msg.clone(), info.clone()))?;
+                                    let payload = Py::new(py, WAMessage::new(msg.clone(), info.clone()))?;
                                     let awaitable = callback.bind(py).call1((py.None(), payload))?;
                                     let fut: Pin<Box<dyn Future<Output = PyResult<Py<PyAny>>> + Send>> = match &locals {
                                         Some(locals) => {
@@ -157,8 +163,15 @@ impl Tryx {
                 PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string())
             })?;
 
-        info!("bot built successfully, starting run loop");
+        let client = bot.client();
+        {
+            let mut state = bot_state
+                .lock()
+                .map_err(|_| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>("Failed to lock bot state"))?;
+            *state = Some(client);
+        }
 
+        info!("bot built successfully, starting run loop");
         bot.run()
             .await
             .map_err(|e| {
@@ -195,6 +208,7 @@ impl Tryx {
             Ok(Tryx {
                 backend: Arc::new(store),
                 handlers: Py::new(py, Dispatcher::empty())?,
+                bot: Arc::new(Mutex::new(None)),
             })
         } else {
             error!("unsupported backend type passed to Tryx");
@@ -219,10 +233,12 @@ impl Tryx {
         info!("starting bot in async mode via Tryx.run");
         let backend = self.backend.clone();
         let handlers = self.handlers.clone_ref(py);
+        let bot_state = self.bot.clone();
         let locals = get_current_locals(py)?;
         future_into_py_with_locals(py, locals.clone(), async move {
-            Self::run_bot(backend, handlers, Some(locals)).await
+            Self::run_bot(backend, handlers, Some(locals), bot_state).await
         })
+    
     }
 
     /// Starts the bot and blocks until it exits.
@@ -234,7 +250,7 @@ impl Tryx {
         info!("starting bot in blocking mode via Tryx.run_blocking");
         let backend = self.backend.clone();
         let handlers = Python::attach(|py| self.handlers.clone_ref(py));
-
+        let bot_state = self.bot.clone();
         py.detach(move || {
             let rt = runtime::Runtime::new()
                 .map_err(|e| {
@@ -243,7 +259,7 @@ impl Tryx {
                 })?;
 
             rt.block_on(async {
-                let mut bot_task = tokio::spawn(Self::run_bot(backend, handlers, None));
+                let mut bot_task = tokio::spawn(Self::run_bot(backend, handlers, None, bot_state));
                 let mut signal_tick = interval(Duration::from_millis(200));
 
                 loop {
@@ -303,4 +319,38 @@ impl Tryx {
             })
         })
     }
-}
+
+    fn send_message<'py>(&self, py: Python<'py>, to: Py<JID>, message: Py<PyAny>) -> PyResult<Bound<'py, PyAny>> {
+        let client = {
+            let state = self
+                .bot
+                .lock()
+                .map_err(|_| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>("Failed to lock bot state"))?;
+            state
+                .clone()
+                .ok_or_else(|| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>("Bot is not running"))?
+        };
+
+        let jid = to.bind(py).borrow().as_whatsapp_jid();
+
+        // Python protobuf object -> bytes -> Rust proto
+        let serialized: Vec<u8> = message
+            .call_method0(py, "SerializeToString")?
+            .extract(py)?;
+
+        let whatsapp_message = WhatsappMessage::decode(serialized.as_slice()).map_err(|e| {
+            PyErr::new::<pyo3::exceptions::PyValueError, _>(
+                format!("Failed to decode WhatsAppMessage proto: {}", e),
+            )
+        })?;
+
+        let locals = get_current_locals(py)?;
+        future_into_py_with_locals(py, locals, async move {
+            let _message_id = client
+                .send_message(jid, whatsapp_message)
+                .await
+                .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))?;
+            Ok(())
+        })
+    }
+}

src/main.rs

@@ -26,10 +26,12 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
                 _ => {}
             }
         })
-        .build()
-        .await?;
+        .build();
+        let mut bot2 = bot.await?;
+        let g = bot2.client();
+        g.send_message(to, message)
 
     // Start the bot
-    bot.run().await?.await?;
+    bot2.run().await?.await?;
     Ok(())
 }

src/types.rs

@@ -11,6 +11,13 @@ use prost::Message;
 pub struct JID {
     inner: Arc<WhatsAppJID>,
 }
+
+impl JID {
+    pub fn as_whatsapp_jid(&self) -> WhatsAppJID {
+        (*self.inner).clone()
+    }
+}
+
 #[pymethods]
 impl JID {
     #[new]