feat: subscription token balalance updates

dojo.h

@@ -898,19 +898,19 @@ struct ResultCArrayToken client_tokens(struct ToriiClient *client,
  *
  * # Parameters
  * * `client` - Pointer to ToriiClient instance
- * * `account_addresses` - Array of account addresses
- * * `account_addresses_len` - Length of account addresses array
  * * `contract_addresses` - Array of contract addresses
  * * `contract_addresses_len` - Length of contract addresses array
+ * * `account_addresses` - Array of account addresses
+ * * `account_addresses_len` - Length of account addresses array
  *
  * # Returns
  * Result containing array of TokenBalance information or error
  */
 struct ResultCArrayTokenBalance client_token_balances(struct ToriiClient *client,
-                                                      const struct FieldElement *account_addresses,
-                                                      uintptr_t account_addresses_len,
                                                       const struct FieldElement *contract_addresses,
-                                                      uintptr_t contract_addresses_len);
+                                                      uintptr_t contract_addresses_len,
+                                                      const struct FieldElement *account_addresses,
+                                                      uintptr_t account_addresses_len);
 
 /**
  * Subscribes to indexer updates
@@ -927,6 +927,48 @@ struct ResultSubscription on_indexer_update(struct ToriiClient *client,
                                             const struct FieldElement *contract_address,
                                             void (*callback)(struct IndexerUpdate));
 
+/**
+ * Subscribes to token balance updates
+ *
+ * # Parameters
+ * * `client` - Pointer to ToriiClient instance
+ * * `contract_addresses` - Array of contract addresses to filter (empty for all)
+ * * `contract_addresses_len` - Length of contract addresses array
+ * * `account_addresses` - Array of account addresses to filter (empty for all)
+ * * `account_addresses_len` - Length of account addresses array
+ * * `callback` - Function called when updates occur
+ *
+ * # Returns
+ * Result containing pointer to Subscription or error
+ */
+struct ResultSubscription client_on_token_balance_update(struct ToriiClient *client,
+                                                         const struct FieldElement *contract_addresses,
+                                                         uintptr_t contract_addresses_len,
+                                                         const struct FieldElement *account_addresses,
+                                                         uintptr_t account_addresses_len,
+                                                         void (*callback)(struct TokenBalance));
+
+/**
+ * Updates an existing token balance subscription
+ *
+ * # Parameters
+ * * `client` - Pointer to ToriiClient instance
+ * * `subscription` - Pointer to existing Subscription
+ * * `contract_addresses` - Array of contract addresses to filter (empty for all)
+ * * `contract_addresses_len` - Length of contract addresses array
+ * * `account_addresses` - Array of account addresses to filter (empty for all)
+ * * `account_addresses_len` - Length of account addresses array
+ *
+ * # Returns
+ * Result containing success boolean or error
+ */
+struct Resultbool client_update_token_balance_subscription(struct ToriiClient *client,
+                                                           struct Subscription *subscription,
+                                                           const struct FieldElement *contract_addresses,
+                                                           uintptr_t contract_addresses_len,
+                                                           const struct FieldElement *account_addresses,
+                                                           uintptr_t account_addresses_len);
+
 /**
  * Serializes a string into a byte array
  *

dojo.hpp

@@ -1085,18 +1085,18 @@ Result<CArray<Token>> client_tokens(ToriiClient *client,
 ///
 /// # Parameters
 /// * `client` - Pointer to ToriiClient instance
-/// * `account_addresses` - Array of account addresses
-/// * `account_addresses_len` - Length of account addresses array
 /// * `contract_addresses` - Array of contract addresses
 /// * `contract_addresses_len` - Length of contract addresses array
+/// * `account_addresses` - Array of account addresses
+/// * `account_addresses_len` - Length of account addresses array
 ///
 /// # Returns
 /// Result containing array of TokenBalance information or error
 Result<CArray<TokenBalance>> client_token_balances(ToriiClient *client,
-                                                   const FieldElement *account_addresses,
-                                                   uintptr_t account_addresses_len,
                                                    const FieldElement *contract_addresses,
-                                                   uintptr_t contract_addresses_len);
+                                                   uintptr_t contract_addresses_len,
+                                                   const FieldElement *account_addresses,
+                                                   uintptr_t account_addresses_len);
 
 /// Subscribes to indexer updates
 ///
@@ -1111,6 +1111,44 @@ Result<Subscription*> on_indexer_update(ToriiClient *client,
                                         const FieldElement *contract_address,
                                         void (*callback)(IndexerUpdate));
 
+/// Subscribes to token balance updates
+///
+/// # Parameters
+/// * `client` - Pointer to ToriiClient instance
+/// * `contract_addresses` - Array of contract addresses to filter (empty for all)
+/// * `contract_addresses_len` - Length of contract addresses array
+/// * `account_addresses` - Array of account addresses to filter (empty for all)
+/// * `account_addresses_len` - Length of account addresses array
+/// * `callback` - Function called when updates occur
+///
+/// # Returns
+/// Result containing pointer to Subscription or error
+Result<Subscription*> client_on_token_balance_update(ToriiClient *client,
+                                                     const FieldElement *contract_addresses,
+                                                     uintptr_t contract_addresses_len,
+                                                     const FieldElement *account_addresses,
+                                                     uintptr_t account_addresses_len,
+                                                     void (*callback)(TokenBalance));
+
+/// Updates an existing token balance subscription
+///
+/// # Parameters
+/// * `client` - Pointer to ToriiClient instance
+/// * `subscription` - Pointer to existing Subscription
+/// * `contract_addresses` - Array of contract addresses to filter (empty for all)
+/// * `contract_addresses_len` - Length of contract addresses array
+/// * `account_addresses` - Array of account addresses to filter (empty for all)
+/// * `account_addresses_len` - Length of account addresses array
+///
+/// # Returns
+/// Result containing success boolean or error
+Result<bool> client_update_token_balance_subscription(ToriiClient *client,
+                                                      Subscription *subscription,
+                                                      const FieldElement *contract_addresses,
+                                                      uintptr_t contract_addresses_len,
+                                                      const FieldElement *account_addresses,
+                                                      uintptr_t account_addresses_len);
+
 /// Serializes a string into a byte array
 ///
 /// # Parameters

dojo.pyx

@@ -612,18 +612,18 @@ cdef extern from *:
   #
   # # Parameters
   # * `client` - Pointer to ToriiClient instance
-  # * `account_addresses` - Array of account addresses
-  # * `account_addresses_len` - Length of account addresses array
   # * `contract_addresses` - Array of contract addresses
   # * `contract_addresses_len` - Length of contract addresses array
+  # * `account_addresses` - Array of account addresses
+  # * `account_addresses_len` - Length of account addresses array
   #
   # # Returns
   # Result containing array of TokenBalance information or error
   ResultCArrayTokenBalance client_token_balances(ToriiClient *client,
-                                                 const FieldElement *account_addresses,
-                                                 uintptr_t account_addresses_len,
                                                  const FieldElement *contract_addresses,
-                                                 uintptr_t contract_addresses_len);
+                                                 uintptr_t contract_addresses_len,
+                                                 const FieldElement *account_addresses,
+                                                 uintptr_t account_addresses_len);
 
   # Subscribes to indexer updates
   #
@@ -638,6 +638,44 @@ cdef extern from *:
                                        const FieldElement *contract_address,
                                        void (*callback)(IndexerUpdate));
 
+  # Subscribes to token balance updates
+  #
+  # # Parameters
+  # * `client` - Pointer to ToriiClient instance
+  # * `contract_addresses` - Array of contract addresses to filter (empty for all)
+  # * `contract_addresses_len` - Length of contract addresses array
+  # * `account_addresses` - Array of account addresses to filter (empty for all)
+  # * `account_addresses_len` - Length of account addresses array
+  # * `callback` - Function called when updates occur
+  #
+  # # Returns
+  # Result containing pointer to Subscription or error
+  ResultSubscription client_on_token_balance_update(ToriiClient *client,
+                                                    const FieldElement *contract_addresses,
+                                                    uintptr_t contract_addresses_len,
+                                                    const FieldElement *account_addresses,
+                                                    uintptr_t account_addresses_len,
+                                                    void (*callback)(TokenBalance));
+
+  # Updates an existing token balance subscription
+  #
+  # # Parameters
+  # * `client` - Pointer to ToriiClient instance
+  # * `subscription` - Pointer to existing Subscription
+  # * `contract_addresses` - Array of contract addresses to filter (empty for all)
+  # * `contract_addresses_len` - Length of contract addresses array
+  # * `account_addresses` - Array of account addresses to filter (empty for all)
+  # * `account_addresses_len` - Length of account addresses array
+  #
+  # # Returns
+  # Result containing success boolean or error
+  Resultbool client_update_token_balance_subscription(ToriiClient *client,
+                                                      Subscription *subscription,
+                                                      const FieldElement *contract_addresses,
+                                                      uintptr_t contract_addresses_len,
+                                                      const FieldElement *account_addresses,
+                                                      uintptr_t account_addresses_len);
+
   # Serializes a string into a byte array
   #
   # # Parameters

src/c/mod.rs

@@ -466,20 +466,20 @@ pub unsafe extern "C" fn client_tokens(
 ///
 /// # Parameters
 /// * `client` - Pointer to ToriiClient instance
-/// * `account_addresses` - Array of account addresses
-/// * `account_addresses_len` - Length of account addresses array
 /// * `contract_addresses` - Array of contract addresses
 /// * `contract_addresses_len` - Length of contract addresses array
+/// * `account_addresses` - Array of account addresses
+/// * `account_addresses_len` - Length of account addresses array
 ///
 /// # Returns
 /// Result containing array of TokenBalance information or error
 #[no_mangle]
 pub unsafe extern "C" fn client_token_balances(
     client: *mut ToriiClient,
-    account_addresses: *const types::FieldElement,
-    account_addresses_len: usize,
     contract_addresses: *const types::FieldElement,
     contract_addresses_len: usize,
+    account_addresses: *const types::FieldElement,
+    account_addresses_len: usize,
 ) -> Result<CArray<TokenBalance>> {
     let account_addresses =
         unsafe { std::slice::from_raw_parts(account_addresses, account_addresses_len) };
@@ -561,6 +561,138 @@ pub unsafe extern "C" fn on_indexer_update(
     Result::Ok(Box::into_raw(Box::new(subscription)))
 }
 
+/// Subscribes to token balance updates
+///
+/// # Parameters
+/// * `client` - Pointer to ToriiClient instance
+/// * `contract_addresses` - Array of contract addresses to filter (empty for all)
+/// * `contract_addresses_len` - Length of contract addresses array
+/// * `account_addresses` - Array of account addresses to filter (empty for all)
+/// * `account_addresses_len` - Length of account addresses array
+/// * `callback` - Function called when updates occur
+///
+/// # Returns
+/// Result containing pointer to Subscription or error
+#[no_mangle]
+pub unsafe extern "C" fn client_on_token_balance_update(
+    client: *mut ToriiClient,
+    contract_addresses: *const types::FieldElement,
+    contract_addresses_len: usize,
+    account_addresses: *const types::FieldElement,
+    account_addresses_len: usize,
+    callback: unsafe extern "C" fn(TokenBalance),
+) -> Result<*mut Subscription> {
+    let client = Arc::new(unsafe { &*client });
+
+    // Convert account addresses array to Vec<Felt> if not empty
+    let account_addresses = if account_addresses.is_null() || account_addresses_len == 0 {
+        Vec::new()
+    } else {
+        let addresses =
+            unsafe { std::slice::from_raw_parts(account_addresses, account_addresses_len) };
+        addresses.iter().map(|f| (&f.clone()).into()).collect::<Vec<Felt>>()
+    };
+
+    // Convert contract addresses array to Vec<Felt> if not empty
+    let contract_addresses = if contract_addresses.is_null() || contract_addresses_len == 0 {
+        Vec::new()
+    } else {
+        let addresses =
+            unsafe { std::slice::from_raw_parts(contract_addresses, contract_addresses_len) };
+        addresses.iter().map(|f| (&f.clone()).into()).collect::<Vec<Felt>>()
+    };
+
+    let subscription_id = Arc::new(AtomicU64::new(0));
+    let (trigger, tripwire) = Tripwire::new();
+
+    let subscription = Subscription { id: Arc::clone(&subscription_id), trigger };
+
+    // Spawn a new thread to handle the stream and reconnections
+    let client_clone = client.clone();
+    client.runtime.spawn(async move {
+        let mut backoff = Duration::from_secs(1);
+        let max_backoff = Duration::from_secs(60);
+
+        loop {
+            let rcv = client_clone
+                .inner
+                .on_token_balance_updated(contract_addresses.clone(), account_addresses.clone())
+                .await;
+
+            if let Ok(rcv) = rcv {
+                backoff = Duration::from_secs(1); // Reset backoff on successful connection
+
+                let mut rcv = rcv.take_until_if(tripwire.clone());
+
+                while let Some(Ok((id, balance))) = rcv.next().await {
+                    subscription_id.store(id, Ordering::SeqCst);
+                    let balance: TokenBalance = (&balance).into();
+                    callback(balance);
+                }
+            }
+
+            // If we've reached this point, the stream has ended (possibly due to disconnection)
+            // We'll try to reconnect after a delay, unless the tripwire has been triggered
+            if tripwire.clone().now_or_never().unwrap_or_default() {
+                break; // Exit the loop if the subscription has been cancelled
+            }
+            sleep(backoff).await;
+            backoff = std::cmp::min(backoff * 2, max_backoff);
+        }
+    });
+
+    Result::Ok(Box::into_raw(Box::new(subscription)))
+}
+
+/// Updates an existing token balance subscription
+///
+/// # Parameters
+/// * `client` - Pointer to ToriiClient instance
+/// * `subscription` - Pointer to existing Subscription
+/// * `contract_addresses` - Array of contract addresses to filter (empty for all)
+/// * `contract_addresses_len` - Length of contract addresses array
+/// * `account_addresses` - Array of account addresses to filter (empty for all)
+/// * `account_addresses_len` - Length of account addresses array
+///
+/// # Returns
+/// Result containing success boolean or error
+#[no_mangle]
+pub unsafe extern "C" fn client_update_token_balance_subscription(
+    client: *mut ToriiClient,
+    subscription: *mut Subscription,
+    contract_addresses: *const types::FieldElement,
+    contract_addresses_len: usize,
+    account_addresses: *const types::FieldElement,
+    account_addresses_len: usize,
+) -> Result<bool> {
+    // Convert account addresses array to Vec<Felt> if not empty
+    let account_addresses = if account_addresses.is_null() || account_addresses_len == 0 {
+        Vec::new()
+    } else {
+        let addresses =
+            unsafe { std::slice::from_raw_parts(account_addresses, account_addresses_len) };
+        addresses.iter().map(|f| (&f.clone()).into()).collect::<Vec<Felt>>()
+    };
+
+    // Convert contract addresses array to Vec<Felt> if not empty
+    let contract_addresses = if contract_addresses.is_null() || contract_addresses_len == 0 {
+        Vec::new()
+    } else {
+        let addresses =
+            unsafe { std::slice::from_raw_parts(contract_addresses, contract_addresses_len) };
+        addresses.iter().map(|f| (&f.clone()).into()).collect::<Vec<Felt>>()
+    };
+
+    match (*client).runtime.block_on((*client).inner.update_token_balance_subscription(
+        (*subscription).id.load(Ordering::SeqCst),
+        contract_addresses,
+        account_addresses,
+    )) {
+        Ok(_) => Result::Ok(true),
+        Err(e) => Result::Err(e.into()),
+    }
+}
+
 /// Serializes a string into a byte array
 ///
 /// # Parameters
@@ -1258,5 +1390,3 @@ pub unsafe extern "C" fn string_free(string: *mut c_char) {
         let _: String = CString::from_raw(string).into_string().unwrap();
     }
 }
-
-// TODO: free keys clause