flwr-nightly 1.19.0.dev20250526__py3-none-any.whl → 1.19.0.dev20250528__py3-none-any.whl

flwr/server/grid/inmemory_grid.py

@@ -18,7 +18,7 @@
 import time
 from collections.abc import Iterable
 from typing import Optional, cast
-from uuid import UUID
+from uuid import uuid4
 
 from flwr.common import Message, RecordDict
 from flwr.common.constant import SUPERLINK_NODE_ID
@@ -56,7 +56,7 @@ class InMemoryGrid(Grid):
     def _check_message(self, message: Message) -> None:
         # Check if the message is valid
         if not (
-            message.metadata.message_id == ""
+            message.metadata.message_id != ""
             and message.metadata.reply_to_message_id == ""
             and message.metadata.ttl > 0
             and message.metadata.delivered_at == ""
@@ -111,6 +111,7 @@ class InMemoryGrid(Grid):
             # Populate metadata
             msg.metadata.__dict__["_run_id"] = cast(Run, self._run).run_id
             msg.metadata.__dict__["_src_node_id"] = self.node.node_id
+            msg.metadata.__dict__["_message_id"] = str(uuid4())
             # Check message
             self._check_message(msg)
             # Store in state
@@ -126,12 +127,12 @@ class InMemoryGrid(Grid):
         This method is used to collect messages from the SuperLink that correspond to a
         set of given message IDs.
         """
-        msg_ids = {UUID(msg_id) for msg_id in message_ids}
+        msg_ids = set(message_ids)
         # Pull Messages
         message_res_list = self.state.get_message_res(message_ids=msg_ids)
         # Get IDs of Messages these replies are for
         message_ins_ids_to_delete = {
-            UUID(msg_res.metadata.reply_to_message_id) for msg_res in message_res_list
+            msg_res.metadata.reply_to_message_id for msg_res in message_res_list
         }
         # Delete
         self.state.delete_messages(message_ins_ids=message_ins_ids_to_delete)

flwr/server/superlink/fleet/message_handler/message_handler.py

@@ -16,7 +16,6 @@
 
 
 from typing import Optional
-from uuid import UUID
 
 from flwr.common import Message
 from flwr.common.constant import Status
@@ -122,7 +121,7 @@ def push_messages(
         raise InvalidRunStatusException(abort_msg)
 
     # Store Message in State
-    message_id: Optional[UUID] = state.store_message_res(message=msg)
+    message_id: Optional[str] = state.store_message_res(message=msg)
 
     # Build response
     response = PushMessagesResponse(

flwr/server/superlink/fleet/vce/vce_api.py

@@ -25,6 +25,7 @@ from pathlib import Path
 from queue import Empty, Queue
 from time import sleep
 from typing import Callable, Optional
+from uuid import uuid4
 
 from flwr.app.error import Error
 from flwr.client.client_app import ClientApp, ClientAppException, LoadClientAppError
@@ -134,6 +135,8 @@ def worker(
 
         finally:
             if out_mssg:
+                # Assign a message_id
+                out_mssg.metadata.__dict__["_message_id"] = str(uuid4())
                 # Store reply Messages in state
                 messageres_queue.put(out_mssg)
 

flwr/server/superlink/linkstate/in_memory_linkstate.py

@@ -21,7 +21,6 @@ from bisect import bisect_right
 from dataclasses import dataclass, field
 from logging import ERROR, WARNING
 from typing import Optional
-from uuid import UUID, uuid4
 
 from flwr.common import Context, Message, log, now
 from flwr.common.constant import (
@@ -76,15 +75,15 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         self.run_ids: dict[int, RunRecord] = {}
         self.contexts: dict[int, Context] = {}
         self.federation_options: dict[int, ConfigRecord] = {}
-        self.message_ins_store: dict[UUID, Message] = {}
-        self.message_res_store: dict[UUID, Message] = {}
-        self.message_ins_id_to_message_res_id: dict[UUID, UUID] = {}
+        self.message_ins_store: dict[str, Message] = {}
+        self.message_res_store: dict[str, Message] = {}
+        self.message_ins_id_to_message_res_id: dict[str, str] = {}
 
         self.node_public_keys: set[bytes] = set()
 
         self.lock = threading.RLock()
 
-    def store_message_ins(self, message: Message) -> Optional[UUID]:
+    def store_message_ins(self, message: Message) -> Optional[str]:
         """Store one Message."""
         # Validate message
         errors = validate_message(message, is_reply_message=False)
@@ -112,12 +111,7 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
             )
             return None
 
-        # Create message_id
-        message_id = uuid4()
-
-        # Store Message
-        # pylint: disable-next=W0212
-        message.metadata._message_id = str(message_id)  # type: ignore
+        message_id = message.metadata.message_id
         with self.lock:
             self.message_ins_store[message_id] = message
 
@@ -153,7 +147,7 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         return message_ins_list
 
     # pylint: disable=R0911
-    def store_message_res(self, message: Message) -> Optional[UUID]:
+    def store_message_res(self, message: Message) -> Optional[str]:
         """Store one Message."""
         # Validate message
         errors = validate_message(message, is_reply_message=True)
@@ -165,7 +159,7 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         with self.lock:
             # Check if the Message it is replying to exists and is valid
             msg_ins_id = res_metadata.reply_to_message_id
-            msg_ins = self.message_ins_store.get(UUID(msg_ins_id))
+            msg_ins = self.message_ins_store.get(msg_ins_id)
 
             # Ensure that dst_node_id of original Message matches the src_node_id of
             # reply Message.
@@ -220,22 +214,17 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
             log(ERROR, "`metadata.run_id` is invalid")
             return None
 
-        # Create message_id
-        message_id = uuid4()
-
-        # Store Message
-        # pylint: disable-next=W0212
-        message.metadata._message_id = str(message_id)  # type: ignore
+        message_id = message.metadata.message_id
         with self.lock:
             self.message_res_store[message_id] = message
-            self.message_ins_id_to_message_res_id[UUID(msg_ins_id)] = message_id
+            self.message_ins_id_to_message_res_id[msg_ins_id] = message_id
 
         # Return the new message_id
         return message_id
 
-    def get_message_res(self, message_ids: set[UUID]) -> list[Message]:
+    def get_message_res(self, message_ids: set[str]) -> list[Message]:
         """Get reply Messages for the given Message IDs."""
-        ret: dict[UUID, Message] = {}
+        ret: dict[str, Message] = {}
 
         with self.lock:
             current = time.time()
@@ -287,7 +276,7 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
 
         return list(ret.values())
 
-    def delete_messages(self, message_ins_ids: set[UUID]) -> None:
+    def delete_messages(self, message_ins_ids: set[str]) -> None:
         """Delete a Message and its reply based on provided Message IDs."""
         if not message_ins_ids:
             return
@@ -304,9 +293,9 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
                     )
                     del self.message_res_store[message_res_id]
 
-    def get_message_ids_from_run_id(self, run_id: int) -> set[UUID]:
+    def get_message_ids_from_run_id(self, run_id: int) -> set[str]:
         """Get all instruction Message IDs for the given run_id."""
-        message_id_list: set[UUID] = set()
+        message_id_list: set[str] = set()
         with self.lock:
             for message_id, message in self.message_ins_store.items():
                 if message.metadata.run_id == run_id:

flwr/server/superlink/linkstate/linkstate.py

@@ -17,7 +17,6 @@
 
 import abc
 from typing import Optional
-from uuid import UUID
 
 from flwr.common import Context, Message
 from flwr.common.record import ConfigRecord
@@ -28,13 +27,13 @@ class LinkState(abc.ABC):  # pylint: disable=R0904
     """Abstract LinkState."""
 
     @abc.abstractmethod
-    def store_message_ins(self, message: Message) -> Optional[UUID]:
+    def store_message_ins(self, message: Message) -> Optional[str]:
         """Store one Message.
 
         Usually, the ServerAppIo API calls this to schedule instructions.
 
         Stores the value of the `message` in the link state and, if successful,
-        returns the `message_id` (UUID) of the `message`. If, for any reason,
+        returns the `message_id` (str) of the `message`. If, for any reason,
         storing the `message` fails, `None` is returned.
 
         Constraints
@@ -61,12 +60,12 @@ class LinkState(abc.ABC):  # pylint: disable=R0904
         """
 
     @abc.abstractmethod
-    def store_message_res(self, message: Message) -> Optional[UUID]:
+    def store_message_res(self, message: Message) -> Optional[str]:
         """Store one Message.
 
         Usually, the Fleet API calls this for Nodes returning results.
 
-        Stores the Message and, if successful, returns the `message_id` (UUID) of
+        Stores the Message and, if successful, returns the `message_id` (str) of
         the `message`. If storing the `message` fails, `None` is returned.
 
         Constraints
@@ -78,7 +77,7 @@ class LinkState(abc.ABC):  # pylint: disable=R0904
         """
 
     @abc.abstractmethod
-    def get_message_res(self, message_ids: set[UUID]) -> list[Message]:
+    def get_message_res(self, message_ids: set[str]) -> list[Message]:
         """Get reply Messages for the given Message IDs.
 
         This method is typically called by the ServerAppIo API to obtain
@@ -94,7 +93,7 @@ class LinkState(abc.ABC):  # pylint: disable=R0904
 
         Parameters
         ----------
-        message_ids : set[UUID]
+        message_ids : set[str]
             A set of Message IDs used to retrieve reply Messages responding to them.
 
         Returns
@@ -113,18 +112,18 @@ class LinkState(abc.ABC):  # pylint: disable=R0904
         """Calculate the number of reply Messages in store."""
 
     @abc.abstractmethod
-    def delete_messages(self, message_ins_ids: set[UUID]) -> None:
+    def delete_messages(self, message_ins_ids: set[str]) -> None:
         """Delete a Message and its reply based on provided Message IDs.
 
         Parameters
         ----------
-        message_ins_ids : set[UUID]
+        message_ins_ids : set[str]
             A set of Message IDs. For each ID in the set, the corresponding
             Message and its associated reply Message will be deleted.
         """
 
     @abc.abstractmethod
-    def get_message_ids_from_run_id(self, run_id: int) -> set[UUID]:
+    def get_message_ids_from_run_id(self, run_id: int) -> set[str]:
         """Get all instruction Message IDs for the given run_id."""
 
     @abc.abstractmethod

flwr/server/superlink/linkstate/sqlite_linkstate.py

@@ -24,7 +24,6 @@ import time
 from collections.abc import Sequence
 from logging import DEBUG, ERROR, WARNING
 from typing import Any, Optional, Union, cast
-from uuid import UUID, uuid4
 
 from flwr.common import Context, Message, Metadata, log, now
 from flwr.common.constant import (
@@ -251,19 +250,15 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
 
         return result
 
-    def store_message_ins(self, message: Message) -> Optional[UUID]:
+    def store_message_ins(self, message: Message) -> Optional[str]:
         """Store one Message."""
         # Validate message
         errors = validate_message(message=message, is_reply_message=False)
         if any(errors):
             log(ERROR, errors)
             return None
-        # Create message_id
-        message_id = uuid4()
 
         # Store Message
-        # pylint: disable-next=W0212
-        message.metadata._message_id = str(message_id)  # type: ignore
         data = (message_to_dict(message),)
 
         # Convert values from uint64 to sint64 for SQLite
@@ -303,7 +298,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         # This may need to be changed in the future version with more integrity checks.
         self.query(query, data)
 
-        return message_id
+        return message.metadata.message_id
 
     def get_message_ins(self, node_id: int, limit: Optional[int]) -> list[Message]:
         """Get all Messages that have not been delivered yet."""
@@ -366,7 +361,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
 
         return result
 
-    def store_message_res(self, message: Message) -> Optional[UUID]:
+    def store_message_res(self, message: Message) -> Optional[str]:
         """Store one Message."""
         # Validate message
         errors = validate_message(message=message, is_reply_message=True)
@@ -418,12 +413,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
             )
             return None
 
-        # Create message_id
-        message_id = uuid4()
-
         # Store Message
-        # pylint: disable-next=W0212
-        message.metadata._message_id = str(message_id)  # type: ignore
         data = (message_to_dict(message),)
 
         # Convert values from uint64 to sint64 for SQLite
@@ -442,12 +432,12 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
             log(ERROR, "`run` is invalid")
             return None
 
-        return message_id
+        return message.metadata.message_id
 
-    def get_message_res(self, message_ids: set[UUID]) -> list[Message]:
+    def get_message_res(self, message_ids: set[str]) -> list[Message]:
         """Get reply Messages for the given Message IDs."""
         # pylint: disable-msg=too-many-locals
-        ret: dict[UUID, Message] = {}
+        ret: dict[str, Message] = {}
 
         # Verify Message IDs
         current = time.time()
@@ -457,12 +447,12 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
             WHERE message_id IN ({",".join(["?"] * len(message_ids))});
         """
         rows = self.query(query, tuple(str(message_id) for message_id in message_ids))
-        found_message_ins_dict: dict[UUID, Message] = {}
+        found_message_ins_dict: dict[str, Message] = {}
         for row in rows:
             convert_sint64_values_in_dict_to_uint64(
                 row, ["run_id", "src_node_id", "dst_node_id"]
             )
-            found_message_ins_dict[UUID(row["message_id"])] = dict_to_message(row)
+            found_message_ins_dict[row["message_id"]] = dict_to_message(row)
 
         ret = verify_message_ids(
             inquired_message_ids=message_ids,
@@ -551,7 +541,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         result: dict[str, int] = rows[0]
         return result["num"]
 
-    def delete_messages(self, message_ins_ids: set[UUID]) -> None:
+    def delete_messages(self, message_ins_ids: set[str]) -> None:
         """Delete a Message and its reply based on provided Message IDs."""
         if not message_ins_ids:
             return
@@ -577,7 +567,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
             self.conn.execute(query_1, data)
             self.conn.execute(query_2, data)
 
-    def get_message_ids_from_run_id(self, run_id: int) -> set[UUID]:
+    def get_message_ids_from_run_id(self, run_id: int) -> set[str]:
         """Get all instruction Message IDs for the given run_id."""
         if self.conn is None:
             raise AttributeError("LinkState not initialized")
@@ -594,7 +584,7 @@ class SqliteLinkState(LinkState):  # pylint: disable=R0904
         with self.conn:
             rows = self.conn.execute(query, data).fetchall()
 
-        return {UUID(row["message_id"]) for row in rows}
+        return {row["message_id"] for row in rows}
 
     def create_node(self, heartbeat_interval: float) -> int:
         """Create, store in the link state, and return `node_id`."""

flwr/server/superlink/linkstate/utils.py

@@ -17,7 +17,7 @@
 
 from os import urandom
 from typing import Optional
-from uuid import UUID, uuid4
+from uuid import uuid4
 
 from flwr.common import ConfigRecord, Context, Error, Message, Metadata, now, serde
 from flwr.common.constant import (
@@ -273,7 +273,7 @@ def create_message_error_unavailable_res_message(
     )
 
 
-def create_message_error_unavailable_ins_message(reply_to_message_id: UUID) -> Message:
+def create_message_error_unavailable_ins_message(reply_to_message_id: str) -> Message:
     """Error to indicate that the enquired Message had expired before reply arrived or
     that it isn't found."""
     metadata = Metadata(
@@ -281,7 +281,7 @@ def create_message_error_unavailable_ins_message(reply_to_message_id: UUID) -> M
         message_id=str(uuid4()),
         src_node_id=SUPERLINK_NODE_ID,
         dst_node_id=SUPERLINK_NODE_ID,
-        reply_to_message_id=str(reply_to_message_id),
+        reply_to_message_id=reply_to_message_id,
         group_id="",  # Unknown
         message_type=MessageType.SYSTEM,
         created_at=now().timestamp(),
@@ -303,18 +303,18 @@ def message_ttl_has_expired(message_metadata: Metadata, current_time: float) ->
 
 
 def verify_message_ids(
-    inquired_message_ids: set[UUID],
-    found_message_ins_dict: dict[UUID, Message],
+    inquired_message_ids: set[str],
+    found_message_ins_dict: dict[str, Message],
     current_time: Optional[float] = None,
     update_set: bool = True,
-) -> dict[UUID, Message]:
+) -> dict[str, Message]:
     """Verify found Messages and generate error Messages for invalid ones.
 
     Parameters
     ----------
-    inquired_message_ids : set[UUID]
+    inquired_message_ids : set[str]
         Set of Message IDs for which to generate error Message if invalid.
-    found_message_ins_dict : dict[UUID, Message]
+    found_message_ins_dict : dict[str, Message]
         Dictionary containing all found Message indexed by their IDs.
     current_time : Optional[float] (default: None)
         The current time to check for expiration. If set to `None`, the current time
@@ -325,7 +325,7 @@ def verify_message_ids(
 
     Returns
     -------
-    dict[UUID, Message]
+    dict[str, Message]
         A dictionary of error Message indexed by the corresponding ID of the message
         they are a reply of.
     """
@@ -345,19 +345,19 @@ def verify_message_ids(
 
 
 def verify_found_message_replies(
-    inquired_message_ids: set[UUID],
-    found_message_ins_dict: dict[UUID, Message],
+    inquired_message_ids: set[str],
+    found_message_ins_dict: dict[str, Message],
     found_message_res_list: list[Message],
     current_time: Optional[float] = None,
     update_set: bool = True,
-) -> dict[UUID, Message]:
+) -> dict[str, Message]:
     """Verify found Message replies and generate error Message for invalid ones.
 
     Parameters
     ----------
-    inquired_message_ids : set[UUID]
+    inquired_message_ids : set[str]
         Set of Message IDs for which to generate error Message if invalid.
-    found_message_ins_dict : dict[UUID, Message]
+    found_message_ins_dict : dict[str, Message]
         Dictionary containing all found instruction Messages indexed by their IDs.
     found_message_res_list : dict[Message, Message]
         List of found Message to be verified.
@@ -370,13 +370,13 @@ def verify_found_message_replies(
 
     Returns
     -------
-    dict[UUID, Message]
+    dict[str, Message]
         A dictionary of Message indexed by the corresponding Message ID.
     """
-    ret_dict: dict[UUID, Message] = {}
+    ret_dict: dict[str, Message] = {}
     current = current_time if current_time else now().timestamp()
     for message_res in found_message_res_list:
-        message_ins_id = UUID(message_res.metadata.reply_to_message_id)
+        message_ins_id = message_res.metadata.reply_to_message_id
         if update_set:
             inquired_message_ids.remove(message_ins_id)
         # Check if the reply Message has expired
@@ -390,21 +390,21 @@ def verify_found_message_replies(
 
 
 def check_node_availability_for_in_message(
-    inquired_in_message_ids: set[UUID],
-    found_in_message_dict: dict[UUID, Message],
+    inquired_in_message_ids: set[str],
+    found_in_message_dict: dict[str, Message],
     node_id_to_online_until: dict[int, float],
     current_time: Optional[float] = None,
     update_set: bool = True,
-) -> dict[UUID, Message]:
+) -> dict[str, Message]:
     """Check node availability for given Message and generate error reply Message if
     unavailable. A Message error indicating node unavailability will be generated for
     each given Message whose destination node is offline or non-existent.
 
     Parameters
     ----------
-    inquired_in_message_ids : set[UUID]
+    inquired_in_message_ids : set[str]
         Set of Message IDs for which to check destination node availability.
-    found_in_message_dict : dict[UUID, Message]
+    found_in_message_dict : dict[str, Message]
         Dictionary containing all found Message indexed by their IDs.
     node_id_to_online_until : dict[int, float]
         Dictionary mapping node IDs to their online-until timestamps.
@@ -417,7 +417,7 @@ def check_node_availability_for_in_message(
 
     Returns
     -------
-    dict[UUID, Message]
+    dict[str, Message]
         A dictionary of error Message indexed by the corresponding Message ID.
     """
     ret_dict = {}

flwr/server/superlink/serverappio/serverappio_servicer.py

@@ -18,7 +18,6 @@
 import threading
 from logging import DEBUG, INFO
 from typing import Optional
-from uuid import UUID
 
 import grpc
 
@@ -140,7 +139,7 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
             request_name="PushMessages",
             detail="`messages_list` must not be empty",
         )
-        message_ids: list[Optional[UUID]] = []
+        message_ids: list[Optional[str]] = []
         while request.messages_list:
             message_proto = request.messages_list.pop(0)
             message = message_from_proto(message_proto=message_proto)
@@ -156,7 +155,7 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
                 detail="`Message.metadata` has mismatched `run_id`",
             )
             # Store
-            message_id: Optional[UUID] = state.store_message_ins(message=message)
+            message_id: Optional[str] = state.store_message_ins(message=message)
             message_ids.append(message_id)
 
         return PushInsMessagesResponse(
@@ -182,17 +181,14 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
             context,
         )
 
-        # Convert each message_id str to UUID
-        message_ids: set[UUID] = {
-            UUID(message_id) for message_id in request.message_ids
-        }
-
         # Read from state
-        messages_res: list[Message] = state.get_message_res(message_ids=message_ids)
+        messages_res: list[Message] = state.get_message_res(
+            message_ids=set(request.message_ids)
+        )
 
         # Delete the instruction Messages and their replies if found
         message_ins_ids_to_delete = {
-            UUID(msg_res.metadata.reply_to_message_id) for msg_res in messages_res
+            msg_res.metadata.reply_to_message_id for msg_res in messages_res
         }
 
         state.delete_messages(message_ins_ids=message_ins_ids_to_delete)

flwr/server/utils/validator.py

@@ -27,8 +27,8 @@ def validate_message(message: Message, is_reply_message: bool) -> list[str]:
     validation_errors = []
     metadata = message.metadata
 
-    if metadata.message_id != "":
-        validation_errors.append("non-empty `metadata.message_id`")
+    if metadata.message_id == "":
+        validation_errors.append("empty `metadata.message_id`")
 
     # Created/delivered/TTL/Pushed
     if (

flwr/supercore/object_store/in_memory_object_store.py

@@ -28,12 +28,27 @@ class InMemoryObjectStore(ObjectStore):
     def __init__(self, verify: bool = True) -> None:
         self.verify = verify
         self.store: dict[str, bytes] = {}
+        # Mapping the Object ID of a message to the list of children object IDs
+        self.msg_children_objects_mapping: dict[str, list[str]] = {}
+
+    def preregister(self, object_ids: list[str]) -> list[str]:
+        """Identify and preregister missing objects."""
+        new_objects = []
+        for obj_id in object_ids:
+            # Verify object ID format (must be a valid sha256 hash)
+            if not is_valid_sha256_hash(obj_id):
+                raise ValueError(f"Invalid object ID format: {obj_id}")
+            if obj_id not in self.store:
+                self.store[obj_id] = b""
+                new_objects.append(obj_id)
+
+        return new_objects
 
     def put(self, object_id: str, object_content: bytes) -> None:
         """Put an object into the store."""
-        # Verify object ID format (must be a valid sha256 hash)
-        if not is_valid_sha256_hash(object_id):
-            raise ValueError(f"Invalid object ID format: {object_id}")
+        # Only allow adding the object if it has been preregistered
+        if object_id not in self.store:
+            raise KeyError(f"Object with id {object_id} was not preregistered.")
 
         # Verify object_id and object_content match
         if self.verify:
@@ -42,11 +57,22 @@ class InMemoryObjectStore(ObjectStore):
                 raise ValueError(f"Object ID {object_id} does not match content hash")
 
         # Return if object is already present in the store
-        if object_id in self.store:
+        if self.store[object_id] != b"":
             return
 
         self.store[object_id] = object_content
 
+    def set_message_descendant_ids(
+        self, msg_object_id: str, descendant_ids: list[str]
+    ) -> None:
+        """Store the mapping from a ``Message`` object ID to the object IDs of its
+        descendants."""
+        self.msg_children_objects_mapping[msg_object_id] = descendant_ids
+
+    def get_message_descendant_ids(self, msg_object_id: str) -> list[str]:
+        """Retrieve the object IDs of all descendants of a given Message."""
+        return self.msg_children_objects_mapping[msg_object_id]
+
     def get(self, object_id: str) -> Optional[bytes]:
         """Get an object from the store."""
         return self.store.get(object_id)