feat: official multi world support (#351)

* feat(storage): multi world scoping for entities

* formatting

* formatting storeage

* update tests

* fix set nttiies

* rm bad mig

* fix up all storeage queries

* big refactor for caching Y& processors mujlti world

* update grpc & all services to multi world support

* fix: storing

* finished multi world support for queries

* debug log rmeove

* further optimize queries

* uodate historifal query

* multi world fixes

* offchain messages world address

* fmt

* world ids test

* ids gql

* fix graphql subscription

* fix last gql subscriptions

* finnalty fix all subs

* add world addresses to grpc subs

* fmt

* clippy fix

Author: Larkooo

crates/cache/src/lib.rs

@@ -3,7 +3,6 @@ use std::sync::Arc;
 
 use async_trait::async_trait;
 use dashmap::DashMap;
-use dojo_types::naming;
 use starknet::core::types::contract::AbiEntry;
 use starknet::core::types::{
     BlockId, BlockTag, ContractClass, EntryPointsByType, LegacyContractAbiEntry, StarknetError,
@@ -24,11 +23,17 @@ pub type CacheError = Box<dyn std::error::Error + Send + Sync>;
 
 #[async_trait]
 pub trait ReadOnlyCache: Send + Sync + std::fmt::Debug {
-    /// Get models by selectors. If selectors is empty, returns all models.
-    async fn models(&self, selectors: &[Felt]) -> Result<Vec<Model>, CacheError>;
+    /// Get models by selectors from specified worlds.
+    /// If world_addresses is empty, returns models from all worlds.
+    /// If selectors is empty, returns all models from the specified worlds.
+    async fn models(
+        &self,
+        world_addresses: &[Felt],
+        selectors: &[Felt],
+    ) -> Result<Vec<Model>, CacheError>;
 
-    /// Get a specific model by selector.
-    async fn model(&self, selector: Felt) -> Result<Model, CacheError>;
+    /// Get a specific model by selector for the given world.
+    async fn model(&self, world_address: Felt, selector: Felt) -> Result<Model, CacheError>;
 
     /// Check if a token is registered.
     async fn is_token_registered(&self, token_id: &TokenId) -> bool;
@@ -45,8 +50,8 @@ pub trait ReadOnlyCache: Send + Sync + std::fmt::Debug {
 
 #[async_trait]
 pub trait Cache: ReadOnlyCache + Send + Sync + std::fmt::Debug {
-    /// Register a model in the cache.
-    async fn register_model(&self, selector: Felt, model: Model);
+    /// Register a model in the cache, scoped to a world.
+    async fn register_model(&self, world_address: Felt, selector: Felt, model: Model);
 
     /// Clear all models from the cache.
     async fn clear_models(&self);
@@ -78,16 +83,20 @@ impl InMemoryCache {
 
 #[async_trait]
 impl ReadOnlyCache for InMemoryCache {
-    async fn models(&self, selectors: &[Felt]) -> Result<Vec<Model>, CacheError> {
+    async fn models(
+        &self,
+        world_addresses: &[Felt],
+        selectors: &[Felt],
+    ) -> Result<Vec<Model>, CacheError> {
         self.model_cache
-            .models(selectors)
+            .models(world_addresses, selectors)
             .await
             .map_err(|e| Box::new(e) as CacheError)
     }
 
-    async fn model(&self, selector: Felt) -> Result<Model, CacheError> {
+    async fn model(&self, world_address: Felt, selector: Felt) -> Result<Model, CacheError> {
         self.model_cache
-            .model(selector)
+            .model(world_address, selector)
             .await
             .map_err(|e| Box::new(e) as CacheError)
     }
@@ -119,8 +128,8 @@ impl ReadOnlyCache for InMemoryCache {
 
 #[async_trait]
 impl Cache for InMemoryCache {
-    async fn register_model(&self, selector: Felt, model: Model) {
-        self.model_cache.set(selector, model).await
+    async fn register_model(&self, world_address: Felt, selector: Felt, model: Model) {
+        self.model_cache.set(world_address, selector, model).await
     }
 
     async fn clear_models(&self) {
@@ -147,48 +156,86 @@ impl Cache for InMemoryCache {
 
 #[derive(Debug)]
 pub struct ModelCache {
-    model_cache: RwLock<HashMap<Felt, Model>>,
+    // Outer key: world_address, Inner key: model_selector
+    model_cache: RwLock<HashMap<Felt, HashMap<Felt, Model>>>,
 }
 
 impl ModelCache {
     pub async fn new(storage: Arc<dyn ReadOnlyStorage>) -> Result<Self, Error> {
-        let models = storage.models(&[]).await?;
+        let models = storage.models(&[], &[]).await?;
 
-        let mut model_cache = HashMap::new();
+        let mut model_cache: HashMap<Felt, HashMap<Felt, Model>> = HashMap::new();
         for model in models {
-            let selector = naming::compute_selector_from_names(&model.namespace, &model.name);
-            model_cache.insert(selector, model);
+            model_cache
+                .entry(model.world_address)
+                .or_default()
+                .insert(model.selector, model);
         }
 
         Ok(Self {
             model_cache: RwLock::new(model_cache),
         })
     }
 
-    pub async fn models(&self, selectors: &[Felt]) -> Result<Vec<Model>, Error> {
-        if selectors.is_empty() {
-            return Ok(self.model_cache.read().await.values().cloned().collect());
-        }
-
-        let mut schemas = Vec::with_capacity(selectors.len());
-        for selector in selectors {
-            schemas.push(self.model(*selector).await?);
+    pub async fn models(
+        &self,
+        world_addresses: &[Felt],
+        selectors: &[Felt],
+    ) -> Result<Vec<Model>, Error> {
+        let cache = self.model_cache.read().await;
+        let mut result = Vec::new();
+
+        if world_addresses.is_empty() {
+            // Return from all worlds
+            for world_models in cache.values() {
+                if selectors.is_empty() {
+                    result.extend(world_models.values().cloned());
+                } else {
+                    for selector in selectors {
+                        if let Some(model) = world_models.get(selector) {
+                            result.push(model.clone());
+                        }
+                    }
+                }
+            }
+        } else {
+            // Return from specific worlds
+            for world in world_addresses {
+                let Some(world_models) = cache.get(world) else {
+                    continue;
+                };
+
+                if selectors.is_empty() {
+                    result.extend(world_models.values().cloned());
+                } else {
+                    for selector in selectors {
+                        if let Some(model) = world_models.get(selector) {
+                            result.push(model.clone());
+                        }
+                    }
+                }
+            }
         }
 
-        Ok(schemas)
+        Ok(result)
     }
 
-    pub async fn model(&self, selector: Felt) -> Result<Model, Error> {
+    pub async fn model(&self, world_address: Felt, selector: Felt) -> Result<Model, Error> {
         let cache = self.model_cache.read().await;
+
         cache
-            .get(&selector)
+            .get(&world_address)
+            .and_then(|world_models| world_models.get(&selector))
             .cloned()
             .ok_or_else(|| Error::ModelNotFound(selector))
     }
 
-    pub async fn set(&self, selector: Felt, model: Model) {
+    pub async fn set(&self, world_address: Felt, selector: Felt, model: Model) {
         let mut cache = self.model_cache.write().await;
-        cache.insert(selector, model);
+        cache
+            .entry(world_address)
+            .or_insert_with(HashMap::new)
+            .insert(selector, model);
     }
 
     pub async fn clear(&self) {

crates/client/src/lib.rs

@@ -32,8 +32,8 @@ pub struct Client {
 
 impl Client {
     /// Returns a initialized [Client] with default max message size (4MB).
-    pub async fn new(torii_url: String, world: Felt) -> Result<Self, Error> {
-        let grpc_client = WorldClient::new(torii_url, world).await?;
+    pub async fn new(torii_url: String) -> Result<Self, Error> {
+        let grpc_client = WorldClient::new(torii_url).await?;
 
         Ok(Self { inner: grpc_client })
     }
@@ -42,39 +42,40 @@ impl Client {
     ///
     /// # Arguments
     /// * `torii_url` - The URL of the Torii server
-    /// * `world` - The world address
     /// * `max_message_size` - Maximum size in bytes for gRPC messages (both incoming and outgoing)
     pub async fn new_with_config(
         torii_url: String,
-        world: Felt,
         max_message_size: usize,
     ) -> Result<Self, Error> {
-        let grpc_client = WorldClient::new_with_config(torii_url, world, max_message_size).await?;
+        let grpc_client = WorldClient::new_with_config(torii_url, max_message_size).await?;
 
         Ok(Self { inner: grpc_client })
     }
 
     /// Publishes an offchain message to the world.
     /// Returns the entity id of the offchain message.
-    pub async fn publish_message(&self, message: Message) -> Result<Felt, Error> {
+    pub async fn publish_message(&self, message: Message) -> Result<String, Error> {
         let mut grpc_client = self.inner.clone();
         let entity_id = grpc_client.publish_message(message).await?;
         Ok(entity_id)
     }
 
     /// Publishes a set of offchain messages to the world.
     /// Returns the entity ids of the offchain messages.
-    pub async fn publish_message_batch(&self, messages: Vec<Message>) -> Result<Vec<Felt>, Error> {
+    pub async fn publish_message_batch(
+        &self,
+        messages: Vec<Message>,
+    ) -> Result<Vec<String>, Error> {
         let mut grpc_client = self.inner.clone();
         let entity_ids = grpc_client.publish_message_batch(messages).await?;
         Ok(entity_ids)
     }
 
     /// Returns a read lock on the World metadata that the client is connected to.
-    pub async fn metadata(&self) -> Result<World, Error> {
+    pub async fn worlds(&self, world_addresses: Vec<Felt>) -> Result<Vec<World>, Error> {
         let mut grpc_client = self.inner.clone();
-        let world = grpc_client.metadata().await?;
-        Ok(world)
+        let worlds = grpc_client.worlds(world_addresses).await?;
+        Ok(worlds)
     }
 
     /// Retrieves controllers matching contract addresses.
@@ -402,9 +403,12 @@ impl Client {
     pub async fn on_entity_updated(
         &self,
         clause: Option<Clause>,
+        world_addresses: Vec<Felt>,
     ) -> Result<EntityUpdateStreaming, Error> {
         let mut grpc_client = self.inner.clone();
-        let stream = grpc_client.subscribe_entities(clause).await?;
+        let stream = grpc_client
+            .subscribe_entities(clause, world_addresses)
+            .await?;
         Ok(stream)
     }
 
@@ -413,10 +417,11 @@ impl Client {
         &self,
         subscription_id: u64,
         clause: Option<Clause>,
+        world_addresses: Vec<Felt>,
     ) -> Result<(), Error> {
         let mut grpc_client = self.inner.clone();
         grpc_client
-            .update_entities_subscription(subscription_id, clause)
+            .update_entities_subscription(subscription_id, clause, world_addresses)
             .await?;
         Ok(())
     }
@@ -425,9 +430,12 @@ impl Client {
     pub async fn on_event_message_updated(
         &self,
         clause: Option<Clause>,
+        world_addresses: Vec<Felt>,
     ) -> Result<EntityUpdateStreaming, Error> {
         let mut grpc_client = self.inner.clone();
-        let stream = grpc_client.subscribe_event_messages(clause).await?;
+        let stream = grpc_client
+            .subscribe_event_messages(clause, world_addresses)
+            .await?;
         Ok(stream)
     }
 
@@ -436,10 +444,11 @@ impl Client {
         &self,
         subscription_id: u64,
         clause: Option<Clause>,
+        world_addresses: Vec<Felt>,
     ) -> Result<(), Error> {
         let mut grpc_client = self.inner.clone();
         grpc_client
-            .update_event_messages_subscription(subscription_id, clause)
+            .update_event_messages_subscription(subscription_id, clause, world_addresses)
             .await?;
         Ok(())
     }

crates/graphql/src/mapping.rs

@@ -461,6 +461,10 @@ pub static TOKEN_TYPE_MAPPING: LazyLock<TypeMapping> = LazyLock::new(|| {
 
 pub static PUBLISH_MESSAGE_INPUT_MAPPING: LazyLock<TypeMapping> = LazyLock::new(|| {
     IndexMap::from([
+        (
+            Name::new("worldAddress"),
+            TypeData::Simple(TypeRef::named_nn(TypeRef::STRING)),
+        ),
         (
             Name::new("signature"),
             TypeData::Simple(TypeRef::named_nn_list(TypeRef::STRING)),
@@ -474,7 +478,7 @@ pub static PUBLISH_MESSAGE_INPUT_MAPPING: LazyLock<TypeMapping> = LazyLock::new(
 
 pub static PUBLISH_MESSAGE_RESPONSE_MAPPING: LazyLock<TypeMapping> = LazyLock::new(|| {
     IndexMap::from([(
-        Name::new("entityId"),
+        Name::new("id"),
         TypeData::Simple(TypeRef::named_nn(TypeRef::STRING)),
     )])
 });

crates/graphql/src/object/entity.rs

@@ -11,6 +11,7 @@ use torii_broker::types::EntityUpdate;
 use torii_broker::MemoryBroker;
 use torii_sqlite::types::Entity;
 use torii_sqlite::utils::felt_to_sql_string;
+use torii_storage::utils::format_world_scoped_id;
 
 use super::inputs::keys_input::keys_argument;
 use super::{BasicObject, ResolvableObject, TypeMapping, ValueMapping};
@@ -83,7 +84,14 @@ impl ResolvableObject for EntityObject {
                             let id = id.clone();
                             async move {
                                 let entity_id = felt_to_sql_string(&entity.entity.hashed_keys);
-                                if id.is_none() || id == Some(entity_id.clone()) {
+                                let scoped_entity_id = format_world_scoped_id(
+                                    &entity.entity.world_address,
+                                    &entity.entity.hashed_keys,
+                                );
+                                if id.is_none()
+                                    || id == Some(entity_id.clone())
+                                    || id == Some(scoped_entity_id.clone())
+                                {
                                     let mut conn = match pool.acquire().await {
                                         Ok(conn) => conn,
                                         Err(_) => return None,
@@ -92,7 +100,7 @@ impl ResolvableObject for EntityObject {
                                     let entity = match sqlx::query_as::<_, Entity>(
                                         "SELECT * FROM entities WHERE id = ?",
                                     )
-                                    .bind(&entity_id)
+                                    .bind(&scoped_entity_id)
                                     .fetch_one(&mut *conn)
                                     .await
                                     {

crates/graphql/src/object/model.rs

@@ -9,6 +9,7 @@ use torii_broker::types::ModelUpdate;
 use torii_broker::MemoryBroker;
 use torii_sqlite::types::Model;
 use torii_sqlite::utils::felt_to_sql_string;
+use torii_storage::utils::format_world_scoped_id;
 
 use super::{resolve_many, BasicObject, ResolvableObject, TypeMapping, ValueMapping};
 use crate::constants::{
@@ -108,7 +109,14 @@ impl ResolvableObject for ModelObject {
                                 let id = id.clone();
                                 async move {
                                     let model_id = felt_to_sql_string(&model.selector);
-                                    if id.is_none() || id == Some(model_id.clone()) {
+                                    let scoped_model_id = format_world_scoped_id(
+                                        &model.world_address,
+                                        &model.selector,
+                                    );
+                                    if id.is_none()
+                                        || id == Some(model_id.clone())
+                                        || id == Some(scoped_model_id.clone())
+                                    {
                                         let mut conn = match pool.acquire().await {
                                             Ok(conn) => conn,
                                             Err(_) => return None,
@@ -117,7 +125,7 @@ impl ResolvableObject for ModelObject {
                                         let model = match sqlx::query_as::<_, Model>(
                                             "SELECT * FROM models WHERE id = ?",
                                         )
-                                        .bind(&model_id)
+                                        .bind(&scoped_model_id)
                                         .fetch_one(&mut *conn)
                                         .await
                                         {