feat(torii-client): token subscription & update subscription via id

crates/torii/client/src/client/mod.rs

@@ -12,6 +12,7 @@
 use tokio::sync::RwLock as AsyncRwLock;
 use torii_grpc::client::{
     EntityUpdateStreaming, EventUpdateStreaming, IndexerUpdateStreaming, TokenBalanceStreaming,
+    TokenUpdateStreaming,
 };
 use torii_grpc::proto::world::{
     RetrieveControllersResponse, RetrieveEntitiesResponse, RetrieveEventsResponse,
@@ -114,6 +115,16 @@
         Ok(tokens.into_iter().map(TryInto::try_into).collect::<Result<Vec<Token>, _>>()?)
     }
 
+    /// A direct stream to grpc subscribe tokens
+    pub async fn on_token_updated(
+        &self,
+        contract_addresses: Vec<Felt>,
+    ) -> Result<TokenUpdateStreaming, Error> {
+        let mut grpc_client = self.inner.write().await;
+        let stream = grpc_client.subscribe_tokens(contract_addresses).await?;
+        Ok(stream)
+    }
+
     /// Retrieves token balances for account addresses and contract addresses.
     pub async fn token_balances(
         &self,

crates/torii/grpc/proto/world.proto

@@ -43,6 +43,9 @@ service World {
     // Subscribe to token updates.
     rpc SubscribeTokens (RetrieveTokensRequest) returns (stream SubscribeTokensResponse);
 
+    // Update token subscription
+    rpc UpdateTokensSubscription (UpdateTokenSubscriptionRequest) returns (google.protobuf.Empty);
+
     // Retrieve entities
     rpc RetrieveEventMessages (RetrieveEventMessagesRequest) returns (RetrieveEntitiesResponse);
 
@@ -101,7 +104,18 @@ message RetrieveTokensResponse {
 
 // A response containing token updates
 message SubscribeTokensResponse {
-    types.Token token = 1;
+    // The subscription ID
+    uint64 subscription_id = 1;
+    // The token
+    types.Token token = 2;
+}
+
+// A request to update a token subscription
+message UpdateTokenSubscriptionRequest {
+    // The subscription ID
+    uint64 subscription_id = 1;
+    // The list of contract addresses to subscribe to
+    repeated bytes contract_addresses = 2;
 }
 
 // A request to retrieve token balances

crates/torii/grpc/src/client.rs

@@ -2,6 +2,7 @@
 use std::num::ParseIntError;
 use std::time::Duration;
 
+use crypto_bigint::U256;
 use futures_util::stream::MapOk;
 use futures_util::{Stream, StreamExt, TryStreamExt};
 use starknet::core::types::{Felt, FromStrError, StateDiff, StateUpdate};
@@ -17,12 +18,13 @@
     SubscribeEntityResponse, SubscribeEventMessagesRequest, SubscribeEventsRequest,
     SubscribeEventsResponse, SubscribeIndexerRequest, SubscribeIndexerResponse,
     SubscribeModelsRequest, SubscribeModelsResponse, SubscribeTokenBalancesResponse,
-    UpdateEntitiesSubscriptionRequest, UpdateEventMessagesSubscriptionRequest,
-    UpdateTokenBalancesSubscriptionRequest, WorldMetadataRequest,
+    SubscribeTokensResponse, UpdateEntitiesSubscriptionRequest,
+    UpdateEventMessagesSubscriptionRequest, UpdateTokenBalancesSubscriptionRequest,
+    UpdateTokenSubscriptionRequest, WorldMetadataRequest,
 };
 use crate::types::schema::{Entity, SchemaError};
 use crate::types::{
-    EntityKeysClause, Event, EventQuery, IndexerUpdate, ModelKeysClause, Query, TokenBalance,
+    EntityKeysClause, Event, EventQuery, IndexerUpdate, ModelKeysClause, Query, Token, TokenBalance,
 };
 
 #[derive(Debug, thiserror::Error)]
@@ -127,6 +129,55 @@
             .map(|res| res.into_inner())
     }
 
+    pub async fn subscribe_tokens(
+        &mut self,
+        contract_addresses: Vec<Felt>,
+    ) -> Result<TokenUpdateStreaming, Error> {
+        let request = RetrieveTokensRequest {
+            contract_addresses: contract_addresses
+                .into_iter()
+                .map(|c| c.to_bytes_be().to_vec())
+                .collect(),
+        };
+        let stream = self
+            .inner
+            .subscribe_tokens(request)
+            .await
+            .map_err(Error::Grpc)
+            .map(|res| res.into_inner())?;
+        Ok(TokenUpdateStreaming(stream.map_ok(Box::new(|res| {
+            (
+                res.subscription_id,
+                match res.token {
+                    Some(token) => token.try_into().expect("must able to serialize"),
+                    None => Token {
+                        id: "".to_string(),
+                        contract_address: Felt::ZERO,
+                        name: "".to_string(),
+                        symbol: "".to_string(),
+                        decimals: 0,
+                        metadata: "".to_string(),
+                    },
+                },
+            )
+        }))))
+    }
+
+    pub async fn update_tokens_subscription(
+        &mut self,
+        subscription_id: u64,
+        contract_addresses: Vec<Felt>,
+    ) -> Result<(), Error> {
+        let contract_addresses =
+            contract_addresses.into_iter().map(|c| c.to_bytes_be().to_vec()).collect();
+        let request = UpdateTokenSubscriptionRequest { subscription_id, contract_addresses };
+        self.inner
+            .update_tokens_subscription(request)
+            .await
+            .map_err(Error::Grpc)
+            .map(|res| res.into_inner())
+    }
+
     pub async fn retrieve_token_balances(
         &mut self,
         account_addresses: Vec<Felt>,
@@ -338,7 +389,18 @@
             .map_err(Error::Grpc)
             .map(|res| res.into_inner())?;
         Ok(TokenBalanceStreaming(stream.map_ok(Box::new(|res| {
-            (res.subscription_id, res.balance.unwrap().try_into().expect("must able to serialize"))
+            (
+                res.subscription_id,
+                match res.balance {
+                    Some(balance) => balance.try_into().expect("must able to serialize"),
+                    None => TokenBalance {
+                        balance: U256::ZERO,
+                        account_address: Felt::ZERO,
+                        contract_address: Felt::ZERO,
+                        token_id: "".to_string(),
+                    },
+                },
+            )
         }))))
     }
 
@@ -368,6 +430,24 @@
     }
 }
 
+type TokenMappedStream = MapOk<
+    tonic::Streaming<SubscribeTokensResponse>,
+    Box<dyn Fn(SubscribeTokensResponse) -> (SubscriptionId, Token) + Send>,
+>;
+
+#[derive(Debug)]
+pub struct TokenUpdateStreaming(TokenMappedStream);
+
+impl Stream for TokenUpdateStreaming {
+    type Item = <TokenMappedStream as Stream>::Item;
+    fn poll_next(
+        mut self: std::pin::Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+    ) -> std::task::Poll<Option<Self::Item>> {
+        self.0.poll_next_unpin(cx)
+    }
+}
+
 type TokenBalanceMappedStream = MapOk<
     tonic::Streaming<SubscribeTokenBalancesResponse>,
     Box<dyn Fn(SubscribeTokenBalancesResponse) -> (SubscriptionId, TokenBalance) + Send>,

crates/torii/grpc/src/server/mod.rs

@@ -63,7 +63,8 @@
     SubscribeEntityResponse, SubscribeEventMessagesRequest, SubscribeEventsResponse,
     SubscribeIndexerRequest, SubscribeIndexerResponse, SubscribeTokenBalancesResponse,
     SubscribeTokensResponse, UpdateEventMessagesSubscriptionRequest,
-    UpdateTokenBalancesSubscriptionRequest, WorldMetadataRequest, WorldMetadataResponse,
+    UpdateTokenBalancesSubscriptionRequest, UpdateTokenSubscriptionRequest, WorldMetadataRequest,
+    WorldMetadataResponse,
 };
 use crate::proto::{self};
 use crate::types::schema::SchemaError;
@@ -1406,6 +1407,20 @@
         Ok(Response::new(Box::pin(ReceiverStream::new(rx)) as Self::SubscribeTokensStream))
     }
 
+    async fn update_tokens_subscription(
+        &self,
+        request: Request<UpdateTokenSubscriptionRequest>,
+    ) -> ServiceResult<()> {
+        let UpdateTokenSubscriptionRequest { subscription_id, contract_addresses } =
+            request.into_inner();
+        let contract_addresses = contract_addresses
+            .iter()
+            .map(|address| Felt::from_bytes_be_slice(address))
+            .collect::<Vec<_>>();
+        self.token_manager.update_subscriber(subscription_id, contract_addresses).await;
+        Ok(Response::new(()))
+    }
+
     async fn retrieve_token_balances(
         &self,
         request: Request<RetrieveTokenBalancesRequest>,

crates/torii/grpc/src/server/subscriptions/token.rs

@@ -14,13 +14,13 @@
 use tokio::sync::RwLock;
 use torii_sqlite::error::{Error, ParseError};
 use torii_sqlite::simple_broker::SimpleBroker;
-use torii_sqlite::types::Token;
+use torii_sqlite::types::OptimisticToken;
 use tracing::{error, trace};
 
 use crate::proto;
 use crate::proto::world::SubscribeTokensResponse;
 
-pub(crate) const LOG_TARGET: &str = "torii::grpc::server::subscriptions::balance";
+pub(crate) const LOG_TARGET: &str = "torii::grpc::server::subscriptions::token";
 
 #[derive(Debug)]
 pub struct TokenSubscriber {
@@ -45,7 +45,7 @@
         let (sender, receiver) = channel(1);
 
         // Send initial empty response
-        let _ = sender.send(Ok(SubscribeTokensResponse { token: None })).await;
+        let _ = sender.send(Ok(SubscribeTokensResponse { subscription_id, token: None })).await;
 
         self.subscribers.write().await.insert(
             subscription_id,
@@ -85,33 +85,38 @@
 #[must_use = "Service does nothing unless polled"]
 #[allow(missing_debug_implementations)]
 pub struct Service {
-    simple_broker: Pin<Box<dyn Stream<Item = Token> + Send>>,
-    balance_sender: UnboundedSender<Token>,
+    simple_broker: Pin<Box<dyn Stream<Item = OptimisticToken> + Send>>,
+    token_sender: UnboundedSender<OptimisticToken>,
 }
 
 impl Service {
     pub fn new(subs_manager: Arc<TokenManager>) -> Self {
-        let (balance_sender, balance_receiver) = unbounded_channel();
-        let service =
-            Self { simple_broker: Box::pin(SimpleBroker::<Token>::subscribe()), balance_sender };
+        let (token_sender, token_receiver) = unbounded_channel();
+        let service = Self {
+            simple_broker: Box::pin(SimpleBroker::<OptimisticToken>::subscribe()),
+            token_sender,
+        };
 
-        tokio::spawn(Self::publish_updates(subs_manager, balance_receiver));
+        tokio::spawn(Self::publish_updates(subs_manager, token_receiver));
 
         service
     }
 
     async fn publish_updates(
         subs: Arc<TokenManager>,
-        mut balance_receiver: UnboundedReceiver<Token>,
+        mut token_receiver: UnboundedReceiver<OptimisticToken>,
     ) {
-        while let Some(balance) = balance_receiver.recv().await {
-            if let Err(e) = Self::process_balance_update(&subs, &balance).await {
-                error!(target = LOG_TARGET, error = %e, "Processing balance update.");
+        while let Some(token) = token_receiver.recv().await {
+            if let Err(e) = Self::process_token_update(&subs, &token).await {
+                error!(target = LOG_TARGET, error = %e, "Processing token update.");
             }
         }
     }
 
-    async fn process_balance_update(subs: &Arc<TokenManager>, token: &Token) -> Result<(), Error> {
+    async fn process_token_update(
+        subs: &Arc<TokenManager>,
+        token: &OptimisticToken,
+    ) -> Result<(), Error> {
         let mut closed_stream = Vec::new();
 
         for (idx, sub) in subs.subscribers.read().await.iter() {
@@ -126,6 +131,7 @@
             }
 
             let resp = SubscribeTokensResponse {
+                subscription_id: *idx,
                 token: Some(proto::types::Token {
                     token_id: token.id.clone(),
                     contract_address: token.contract_address.clone(),
@@ -142,7 +148,7 @@
         }
 
         for id in closed_stream {
-            trace!(target = LOG_TARGET, id = %id, "Closing balance stream.");
+            trace!(target = LOG_TARGET, id = %id, "Closing token stream.");
             subs.remove_subscriber(id).await
         }
 
@@ -156,9 +162,9 @@
     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
         let this = self.get_mut();
 
-        while let Poll::Ready(Some(balance)) = this.simple_broker.poll_next_unpin(cx) {
-            if let Err(e) = this.balance_sender.send(balance) {
-                error!(target = LOG_TARGET, error = %e, "Sending balance update to processor.");
+        while let Poll::Ready(Some(token)) = this.simple_broker.poll_next_unpin(cx) {
+            if let Err(e) = this.token_sender.send(token) {
+                error!(target = LOG_TARGET, error = %e, "Sending token update to processor.");
             }
         }
 

crates/torii/grpc/src/server/subscriptions/token_balance.rs

@@ -14,7 +14,7 @@
 use tokio::sync::RwLock;
 use torii_sqlite::error::{Error, ParseError};
 use torii_sqlite::simple_broker::SimpleBroker;
-use torii_sqlite::types::TokenBalance;
+use torii_sqlite::types::OptimisticTokenBalance;
 use tracing::{error, trace};
 
 use crate::proto;
@@ -98,15 +98,15 @@
 #[must_use = "Service does nothing unless polled"]
 #[allow(missing_debug_implementations)]
 pub struct Service {
-    simple_broker: Pin<Box<dyn Stream<Item = TokenBalance> + Send>>,
-    balance_sender: UnboundedSender<TokenBalance>,
+    simple_broker: Pin<Box<dyn Stream<Item = OptimisticTokenBalance> + Send>>,
+    balance_sender: UnboundedSender<OptimisticTokenBalance>,
 }
 
 impl Service {
     pub fn new(subs_manager: Arc<TokenBalanceManager>) -> Self {
         let (balance_sender, balance_receiver) = unbounded_channel();
         let service = Self {
-            simple_broker: Box::pin(SimpleBroker::<TokenBalance>::subscribe()),
+            simple_broker: Box::pin(SimpleBroker::<OptimisticTokenBalance>::subscribe()),
             balance_sender,
         };
 
@@ -117,7 +117,7 @@
 
     async fn publish_updates(
         subs: Arc<TokenBalanceManager>,
-        mut balance_receiver: UnboundedReceiver<TokenBalance>,
+        mut balance_receiver: UnboundedReceiver<OptimisticTokenBalance>,
     ) {
         while let Some(balance) = balance_receiver.recv().await {
             if let Err(e) = Self::process_balance_update(&subs, &balance).await {
@@ -128,7 +128,7 @@
 
     async fn process_balance_update(
         subs: &Arc<TokenBalanceManager>,
-        balance: &TokenBalance,
+        balance: &OptimisticTokenBalance,
     ) -> Result<(), Error> {
         let mut closed_stream = Vec::new();
 

crates/torii/grpc/src/types/mod.rs

@@ -38,6 +38,7 @@
 
 #[derive(Debug, Serialize, Deserialize, PartialEq, Hash, Eq, Clone)]
 pub struct Token {
+    pub id: String,
     pub contract_address: Felt,
     pub name: String,
     pub symbol: String,
@@ -49,6 +50,7 @@
     type Error = SchemaError;
     fn try_from(value: proto::types::Token) -> Result<Self, Self::Error> {
         Ok(Self {
+            id: value.token_id,
             contract_address: Felt::from_str(&value.contract_address)?,
             name: value.name,
             symbol: value.symbol,