flwr-nightly 1.19.0.dev20250529__py3-none-any.whl → 1.19.0.dev20250530__py3-none-any.whl

flwr/cli/utils.py

@@ -291,9 +291,9 @@ def init_channel(
 def flwr_cli_grpc_exc_handler() -> Iterator[None]:
     """Context manager to handle specific gRPC errors.
 
-    It catches grpc.RpcError exceptions with UNAUTHENTICATED and UNIMPLEMENTED statuses,
-    informs the user, and exits the application. All other exceptions will be allowed to
-    escape.
+    It catches grpc.RpcError exceptions with UNAUTHENTICATED, UNIMPLEMENTED, and
+    PERMISSION_DENIED statuses, informs the user, and exits the application. All other
+    exceptions will be allowed to escape.
     """
     try:
         yield
@@ -313,4 +313,12 @@ def flwr_cli_grpc_exc_handler() -> Iterator[None]:
                 bold=True,
             )
             raise typer.Exit(code=1) from None
+        if e.code() == grpc.StatusCode.PERMISSION_DENIED:
+            typer.secho(
+                "❌ Authorization failed. Please contact your administrator"
+                " to check your permissions.",
+                fg=typer.colors.RED,
+                bold=True,
+            )
+            raise typer.Exit(code=1) from None
         raise

flwr/common/auth_plugin/auth_plugin.py

@@ -63,7 +63,7 @@ class ExecAuthPlugin(ABC):
     @abstractmethod
     def refresh_tokens(
         self, metadata: Sequence[tuple[str, Union[str, bytes]]]
-    ) -> Optional[Sequence[tuple[str, Union[str, bytes]]]]:
+    ) -> tuple[Optional[Sequence[tuple[str, Union[str, bytes]]]], Optional[UserInfo]]:
         """Refresh authentication tokens in the provided metadata."""
 
 

flwr/common/exit_handlers.py

@@ -30,6 +30,7 @@ SIGNAL_TO_EXIT_CODE: dict[int, int] = {
     signal.SIGINT: ExitCode.GRACEFUL_EXIT_SIGINT,
     signal.SIGTERM: ExitCode.GRACEFUL_EXIT_SIGTERM,
 }
+registered_exit_handlers: list[Callable[[], None]] = []
 
 # SIGQUIT is not available on Windows
 if hasattr(signal, "SIGQUIT"):
@@ -41,6 +42,7 @@ def register_exit_handlers(
     exit_message: Optional[str] = None,
     grpc_servers: Optional[list[Server]] = None,
     bckg_threads: Optional[list[Thread]] = None,
+    exit_handlers: Optional[list[Callable[[], None]]] = None,
 ) -> None:
     """Register exit handlers for `SIGINT`, `SIGTERM` and `SIGQUIT` signals.
 
@@ -56,8 +58,12 @@ def register_exit_handlers(
     bckg_threads: Optional[List[Thread]] (default: None)
         An optional list of threads that need to be gracefully
         terminated before exiting.
+    exit_handlers: Optional[List[Callable[[], None]]] (default: None)
+        An optional list of exit handlers to be called before exiting.
+        Additional exit handlers can be added using `add_exit_handler`.
     """
     default_handlers: dict[int, Callable[[int, FrameType], None]] = {}
+    registered_exit_handlers.extend(exit_handlers or [])
 
     def graceful_exit_handler(signalnum: int, _frame: FrameType) -> None:
         """Exit handler to be registered with `signal.signal`.
@@ -68,6 +74,9 @@ def register_exit_handlers(
         # Reset to default handler
         signal.signal(signalnum, default_handlers[signalnum])  # type: ignore
 
+        for handler in registered_exit_handlers:
+            handler()
+
         if grpc_servers is not None:
             for grpc_server in grpc_servers:
                 grpc_server.stop(grace=1)
@@ -87,3 +96,24 @@ def register_exit_handlers(
     for sig in SIGNAL_TO_EXIT_CODE:
         default_handler = signal.signal(sig, graceful_exit_handler)  # type: ignore
         default_handlers[sig] = default_handler  # type: ignore
+
+
+def add_exit_handler(exit_handler: Callable[[], None]) -> None:
+    """Add an exit handler to be called on graceful exit.
+
+    This function allows you to register additional exit handlers
+    that will be executed when the application exits gracefully,
+    if `register_exit_handlers` was called.
+
+    Parameters
+    ----------
+    exit_handler : Callable[[], None]
+        A callable that takes no arguments and performs cleanup or
+        other actions before the application exits.
+
+    Notes
+    -----
+    This method is not thread-safe, and it allows you to add the
+    same exit handler multiple times.
+    """
+    registered_exit_handlers.append(exit_handler)

flwr/common/inflatable_grpc_utils.py

@@ -15,7 +15,7 @@
 """InflatableObject utils."""
 
 
-from typing import Union
+from typing import Optional, Union
 
 from flwr.proto.fleet_pb2_grpc import FleetStub  # pylint: disable=E0611
 from flwr.proto.message_pb2 import (  # pylint: disable=E0611
@@ -24,6 +24,7 @@ from flwr.proto.message_pb2 import (  # pylint: disable=E0611
     PushObjectRequest,
     PushObjectResponse,
 )
+from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
 from flwr.proto.serverappio_pb2_grpc import ServerAppIoStub  # pylint: disable=E0611
 
 from .inflatable import (
@@ -46,40 +47,51 @@ inflatable_class_registry: dict[str, type[InflatableObject]] = {
 
 
 def push_object_to_servicer(
-    obj: InflatableObject, stub: Union[FleetStub, ServerAppIoStub]
+    obj: InflatableObject,
+    stub: Union[FleetStub, ServerAppIoStub],
+    node: Node,
+    object_ids_to_push: Optional[set[str]] = None,
 ) -> set[str]:
     """Recursively deflate an object and push it to the servicer.
 
-    Objects with the same ID are not pushed twice. It returns the set of pushed object
+    Objects with the same ID are not pushed twice. If `object_ids_to_push` is set,
+    only objects with those IDs are pushed. It returns the set of pushed object
     IDs.
     """
     pushed_object_ids: set[str] = set()
     # Push children if it has any
     if children := obj.children:
         for child in children.values():
-            pushed_object_ids |= push_object_to_servicer(child, stub)
+            pushed_object_ids |= push_object_to_servicer(
+                child, stub, node, object_ids_to_push
+            )
 
     # Deflate object and push
     object_content = obj.deflate()
     object_id = get_object_id(object_content)
-    _: PushObjectResponse = stub.PushObject(
-        PushObjectRequest(
-            object_id=object_id,
-            object_content=object_content,
+    # Push always if no object set is specified, or if the object is in the set
+    if object_ids_to_push is None or object_id in object_ids_to_push:
+        _: PushObjectResponse = stub.PushObject(
+            PushObjectRequest(
+                node=node,
+                object_id=object_id,
+                object_content=object_content,
+            )
         )
-    )
-    pushed_object_ids.add(object_id)
+        pushed_object_ids.add(object_id)
 
     return pushed_object_ids
 
 
 def pull_object_from_servicer(
-    object_id: str, stub: Union[FleetStub, ServerAppIoStub]
+    object_id: str,
+    stub: Union[FleetStub, ServerAppIoStub],
+    node: Node,
 ) -> InflatableObject:
     """Recursively inflate an object by pulling it from the servicer."""
     # Pull object
     object_proto: PullObjectResponse = stub.PullObject(
-        PullObjectRequest(object_id=object_id)
+        PullObjectRequest(node=node, object_id=object_id)
     )
     object_content = object_proto.object_content
 
@@ -93,7 +105,9 @@ def pull_object_from_servicer(
     # Pull all children objects
     children: dict[str, InflatableObject] = {}
     for child_object_id in children_obj_ids:
-        children[child_object_id] = pull_object_from_servicer(child_object_id, stub)
+        children[child_object_id] = pull_object_from_servicer(
+            child_object_id, stub, node
+        )
 
     # Inflate object passing its children
     return cls_type.inflate(object_content, children=children)

flwr/common/record/array.py

@@ -62,8 +62,8 @@ class Array(InflatableObject):
         A string representing the data type of the serialized object (e.g. `"float32"`).
         Only required if you are not passing in a ndarray or a tensor.
 
-    shape : Optional[list[int]] (default: None)
-        A list representing the shape of the unserialized array-like object. Only
+    shape : Optional[tuple[int, ...]] (default: None)
+        A tuple representing the shape of the unserialized array-like object. Only
         required if you are not passing in a ndarray or a tensor.
 
     stype : Optional[str] (default: None)
@@ -107,24 +107,13 @@ class Array(InflatableObject):
     """
 
     dtype: str
+    shape: tuple[int, ...]
     stype: str
     data: bytes
 
-    @property
-    def shape(self) -> list[int]:
-        """Get the shape of the array."""
-        self.is_dirty = True  # Mark as dirty when shape is accessed
-        return cast(list[int], self.__dict__["_shape"])
-
-    @shape.setter
-    def shape(self, value: list[int]) -> None:
-        """Set the shape of the array."""
-        self.is_dirty = True  # Mark as dirty when shape is set
-        self.__dict__["_shape"] = value
-
     @overload
     def __init__(  # noqa: E704
-        self, dtype: str, shape: list[int], stype: str, data: bytes
+        self, dtype: str, shape: tuple[int, ...], stype: str, data: bytes
     ) -> None: ...
 
     @overload
@@ -137,7 +126,7 @@ class Array(InflatableObject):
         self,
         *args: Any,
         dtype: str | None = None,
-        shape: list[int] | None = None,
+        shape: tuple[int, ...] | None = None,
         stype: str | None = None,
         data: bytes | None = None,
         ndarray: NDArray | None = None,
@@ -145,7 +134,7 @@ class Array(InflatableObject):
     ) -> None:
         # Determine the initialization method and validate input arguments.
         # Support three initialization formats:
-        # 1. Array(dtype: str, shape: list[int], stype: str, data: bytes)
+        # 1. Array(dtype: str, shape: tuple[int, ...], stype: str, data: bytes)
         # 2. Array(ndarray: NDArray)
         # 3. Array(torch_tensor: torch.Tensor)
 
@@ -192,7 +181,7 @@ class Array(InflatableObject):
             if (
                 len(all_args) == 4  # pylint: disable=too-many-boolean-expressions
                 and isinstance(all_args[0], str)
-                and isinstance(all_args[1], list)
+                and isinstance(all_args[1], tuple)
                 and all(isinstance(i, int) for i in all_args[1])
                 and isinstance(all_args[2], str)
                 and isinstance(all_args[3], bytes)
@@ -232,7 +221,7 @@ class Array(InflatableObject):
         data = buffer.getvalue()
         return Array(
             dtype=str(ndarray.dtype),
-            shape=list(ndarray.shape),
+            shape=tuple(ndarray.shape),
             stype=SType.NUMPY,
             data=data,
         )
@@ -302,7 +291,7 @@ class Array(InflatableObject):
         proto_array = ArrayProto.FromString(obj_body)
         return cls(
             dtype=proto_array.dtype,
-            shape=list(proto_array.shape),
+            shape=tuple(proto_array.shape),
             stype=proto_array.stype,
             data=proto_array.data,
         )
@@ -328,7 +317,7 @@ class Array(InflatableObject):
 
     def __setattr__(self, name: str, value: Any) -> None:
         """Set attribute with special handling for dirty state."""
-        if name in ("dtype", "stype", "data"):
+        if name in ("dtype", "shape", "stype", "data"):
             # Mark as dirty if any of the main attributes are set
             self.is_dirty = True
         super().__setattr__(name, value)

flwr/common/record/arrayrecord.py

@@ -252,7 +252,7 @@ class ArrayRecord(TypedDict[str, Array], InflatableObject):
         record = ArrayRecord()
         for k, v in array_dict.items():
             record[k] = Array(
-                dtype=v.dtype, shape=list(v.shape), stype=v.stype, data=v.data
+                dtype=v.dtype, shape=tuple(v.shape), stype=v.stype, data=v.data
             )
         if not keep_input:
             array_dict.clear()

flwr/common/recorddict_compat.py

@@ -111,12 +111,12 @@ def parameters_to_arrayrecord(parameters: Parameters, keep_input: bool) -> Array
         else:
             tensor = parameters.tensors.pop(0)
         ordered_dict[str(idx)] = Array(
-            data=tensor, dtype="", stype=tensor_type, shape=[]
+            data=tensor, dtype="", stype=tensor_type, shape=()
         )
 
     if num_arrays == 0:
         ordered_dict[EMPTY_TENSOR_KEY] = Array(
-            data=b"", dtype="", stype=tensor_type, shape=[]
+            data=b"", dtype="", stype=tensor_type, shape=()
         )
     return ArrayRecord(ordered_dict, keep_input=keep_input)
 

flwr/common/serde.py

@@ -390,7 +390,7 @@ def array_from_proto(array_proto: ProtoArray) -> Array:
     """Deserialize Array from ProtoBuf."""
     return Array(
         dtype=array_proto.dtype,
-        shape=list(array_proto.shape),
+        shape=tuple(array_proto.shape),
         stype=array_proto.stype,
         data=array_proto.data,
     )

flwr/server/app.py

@@ -27,7 +27,7 @@ from collections.abc import Sequence
 from logging import DEBUG, INFO, WARN
 from pathlib import Path
 from time import sleep
-from typing import Any, Callable, Optional, Union, cast
+from typing import Any, Callable, Optional, TypeVar
 
 import grpc
 import yaml
@@ -85,6 +85,7 @@ from .superlink.simulation.simulationio_grpc import run_simulationio_api_grpc
 
 DATABASE = ":flwr-in-memory-state:"
 BASE_DIR = get_flwr_dir() / "superlink" / "ffs"
+P = TypeVar("P", ExecAuthPlugin, ExecAuthzPlugin)
 
 
 try:
@@ -151,15 +152,13 @@ def run_superlink() -> None:
     verify_tls_cert = not getattr(args, "disable_oidc_tls_cert_verification", None)
 
     auth_plugin: Optional[ExecAuthPlugin] = None
-    authz_plugin: Optional[ExecAuthzPlugin] = None  # pylint: disable=unused-variable
+    authz_plugin: Optional[ExecAuthzPlugin] = None
     event_log_plugin: Optional[EventLogWriterPlugin] = None
     # Load the auth plugin if the args.user_auth_config is provided
     if cfg_path := getattr(args, "user_auth_config", None):
-        # pylint: disable=unused-variable
-        auth_plugin, authz_plugin = _try_obtain_exec_auth_plugins(  # noqa: F841
+        auth_plugin, authz_plugin = _try_obtain_exec_auth_plugins(
             Path(cfg_path), verify_tls_cert
         )
-        # pylint: enable=unused-variable
         # Enable event logging if the args.enable_event_log is True
         if args.enable_event_log:
             event_log_plugin = _try_obtain_exec_event_log_writer_plugin()
@@ -185,6 +184,7 @@ def run_superlink() -> None:
             [args.executor_config] if args.executor_config else args.executor_config
         ),
         auth_plugin=auth_plugin,
+        authz_plugin=authz_plugin,
         event_log_plugin=event_log_plugin,
     )
     grpc_servers = [exec_server]
@@ -490,15 +490,13 @@ def _try_obtain_exec_auth_plugins(
         config: dict[str, Any] = yaml.safe_load(file)
 
     def _load_plugin(
-        section: str,
-        yaml_key: str,
-        loader: Callable[[], dict[str, type[Union[ExecAuthPlugin, ExecAuthzPlugin]]]],
-    ) -> Union[ExecAuthPlugin, ExecAuthzPlugin]:
+        section: str, yaml_key: str, loader: Callable[[], dict[str, type[P]]]
+    ) -> P:
         section_cfg = config.get(section, {})
         auth_plugin_name = section_cfg.get(yaml_key, "")
         try:
-            plugins = loader()
-            plugin_cls = plugins[auth_plugin_name]
+            plugins: dict[str, type[P]] = loader()
+            plugin_cls: type[P] = plugins[auth_plugin_name]
             return plugin_cls(
                 user_auth_config_path=config_path, verify_tls_cert=verify_tls_cert
             )
@@ -513,23 +511,17 @@ def _try_obtain_exec_auth_plugins(
             sys.exit(f"No {section} plugins are currently supported.")
 
     # Load authentication plugin
-    auth_plugin = cast(
-        ExecAuthPlugin,
-        _load_plugin(
-            section="authentication",
-            yaml_key=AUTH_TYPE_YAML_KEY,
-            loader=get_exec_auth_plugins,
-        ),
+    auth_plugin = _load_plugin(
+        section="authentication",
+        yaml_key=AUTH_TYPE_YAML_KEY,
+        loader=get_exec_auth_plugins,
     )
 
     # Load authorization plugin
-    authz_plugin = cast(
-        ExecAuthzPlugin,
-        _load_plugin(
-            section="authorization",
-            yaml_key=AUTHZ_TYPE_YAML_KEY,
-            loader=get_exec_authz_plugins,
-        ),
+    authz_plugin = _load_plugin(
+        section="authorization",
+        yaml_key=AUTHZ_TYPE_YAML_KEY,
+        loader=get_exec_authz_plugins,
     )
 
     return auth_plugin, authz_plugin

flwr/server/superlink/fleet/grpc_rere/fleet_servicer.py

@@ -114,6 +114,7 @@ class FleetServicer(fleet_pb2_grpc.FleetServicer):
         return message_handler.pull_messages(
             request=request,
             state=self.state_factory.state(),
+            store=self.objectstore_factory.store(),
         )
 
     def PushMessages(

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

@@ -14,10 +14,10 @@
 # ==============================================================================
 """Fleet API message handlers."""
 
-
+from logging import ERROR
 from typing import Optional
 
-from flwr.common import Message
+from flwr.common import Message, log
 from flwr.common.constant import Status
 from flwr.common.serde import (
     fab_to_proto,
@@ -42,6 +42,7 @@ from flwr.proto.heartbeat_pb2 import (  # pylint: disable=E0611
     SendNodeHeartbeatRequest,
     SendNodeHeartbeatResponse,
 )
+from flwr.proto.message_pb2 import ObjectIDs  # pylint: disable=E0611
 from flwr.proto.node_pb2 import Node  # pylint: disable=E0611
 from flwr.proto.run_pb2 import (  # pylint: disable=E0611
     GetRunRequest,
@@ -51,7 +52,7 @@ from flwr.proto.run_pb2 import (  # pylint: disable=E0611
 from flwr.server.superlink.ffs.ffs import Ffs
 from flwr.server.superlink.linkstate import LinkState
 from flwr.server.superlink.utils import check_abort
-from flwr.supercore.object_store import ObjectStore
+from flwr.supercore.object_store import NoObjectInStoreError, ObjectStore
 
 from ...utils import store_mapping_and_register_objects
 
@@ -89,7 +90,9 @@ def send_node_heartbeat(
 
 
 def pull_messages(
-    request: PullMessagesRequest, state: LinkState
+    request: PullMessagesRequest,
+    state: LinkState,
+    store: ObjectStore,
 ) -> PullMessagesResponse:
     """Pull Messages handler."""
     # Get node_id if client node is not anonymous
@@ -101,10 +104,25 @@ def pull_messages(
 
     # Convert to Messages
     msg_proto = []
+    objects_to_pull: dict[str, ObjectIDs] = {}
     for msg in message_list:
-        msg_proto.append(message_to_proto(msg))
-
-    return PullMessagesResponse(messages_list=msg_proto)
+        try:
+            msg_proto.append(message_to_proto(msg))
+
+            msg_object_id = msg.metadata.message_id
+            descendants = store.get_message_descendant_ids(msg_object_id)
+            # Include the object_id of the message itself
+            objects_to_pull[msg_object_id] = ObjectIDs(
+                object_ids=descendants + [msg_object_id]
+            )
+        except NoObjectInStoreError as e:
+            log(ERROR, e.message)
+            # Delete message ins from state
+            state.delete_messages(message_ins_ids={msg_object_id})
+
+    return PullMessagesResponse(
+        messages_list=msg_proto, objects_to_pull=objects_to_pull
+    )
 
 
 def push_messages(

flwr/server/superlink/fleet/rest_rere/rest_api.py

@@ -114,9 +114,10 @@ async def pull_message(request: PullMessagesRequest) -> PullMessagesResponse:
     """Pull PullMessages."""
     # Get state from app
     state: LinkState = cast(LinkStateFactory, app.state.STATE_FACTORY).state()
+    store: ObjectStore = cast(ObjectStoreFactory, app.state.STATE_FACTORY).store()
 
     # Handle message
-    return message_handler.pull_messages(request=request, state=state)
+    return message_handler.pull_messages(request=request, state=state, store=store)
 
 
 @rest_request_response(PushMessagesRequest)

flwr/server/superlink/linkstate/utils.py

@@ -17,7 +17,6 @@
 
 from os import urandom
 from typing import Optional
-from uuid import uuid4
 
 from flwr.common import ConfigRecord, Context, Error, Message, Metadata, now, serde
 from flwr.common.constant import (
@@ -246,7 +245,7 @@ def create_message_error_unavailable_res_message(
     ttl = max(ins_metadata.ttl - (current_time - ins_metadata.created_at), 0)
     metadata = Metadata(
         run_id=ins_metadata.run_id,
-        message_id=str(uuid4()),
+        message_id="",
         src_node_id=SUPERLINK_NODE_ID,
         dst_node_id=SUPERLINK_NODE_ID,
         reply_to_message_id=ins_metadata.message_id,
@@ -256,7 +255,7 @@ def create_message_error_unavailable_res_message(
         ttl=ttl,
     )
 
-    return make_message(
+    msg = make_message(
         metadata=metadata,
         error=Error(
             code=(
@@ -271,6 +270,8 @@ def create_message_error_unavailable_res_message(
             ),
         ),
     )
+    msg.metadata.__dict__["_message_id"] = msg.object_id
+    return msg
 
 
 def create_message_error_unavailable_ins_message(reply_to_message_id: str) -> Message:
@@ -278,7 +279,7 @@ def create_message_error_unavailable_ins_message(reply_to_message_id: str) -> Me
     that it isn't found."""
     metadata = Metadata(
         run_id=0,  # Unknown
-        message_id=str(uuid4()),
+        message_id="",
         src_node_id=SUPERLINK_NODE_ID,
         dst_node_id=SUPERLINK_NODE_ID,
         reply_to_message_id=reply_to_message_id,
@@ -288,13 +289,15 @@ def create_message_error_unavailable_ins_message(reply_to_message_id: str) -> Me
         ttl=0,
     )
 
-    return make_message(
+    msg = make_message(
         metadata=metadata,
         error=Error(
             code=ErrorCode.MESSAGE_UNAVAILABLE,
             reason=MESSAGE_UNAVAILABLE_ERROR_REASON,
         ),
     )
+    msg.metadata.__dict__["_message_id"] = msg.object_id
+    return msg
 
 
 def message_ttl_has_expired(message_metadata: Metadata, current_time: float) -> bool:

flwr/server/superlink/serverappio/serverappio_servicer.py

@@ -16,14 +16,14 @@
 
 
 import threading
-from logging import DEBUG, INFO
+from logging import DEBUG, ERROR, INFO
 from typing import Optional
 
 import grpc
 
 from flwr.common import Message
 from flwr.common.constant import SUPERLINK_NODE_ID, Status
-from flwr.common.inflatable import check_body_len_consistency
+from flwr.common.inflatable import check_body_len_consistency, get_desdendant_object_ids
 from flwr.common.logger import log
 from flwr.common.serde import (
     context_from_proto,
@@ -47,6 +47,7 @@ from flwr.proto.log_pb2 import (  # pylint: disable=E0611
     PushLogsResponse,
 )
 from flwr.proto.message_pb2 import (  # pylint: disable=E0611
+    ObjectIDs,
     PullObjectRequest,
     PullObjectResponse,
     PushObjectRequest,
@@ -78,7 +79,7 @@ from flwr.server.superlink.ffs.ffs_factory import FfsFactory
 from flwr.server.superlink.linkstate import LinkState, LinkStateFactory
 from flwr.server.superlink.utils import abort_if
 from flwr.server.utils.validator import validate_message
-from flwr.supercore.object_store import ObjectStoreFactory
+from flwr.supercore.object_store import NoObjectInStoreError, ObjectStoreFactory
 
 from ..utils import store_mapping_and_register_objects
 
@@ -182,6 +183,9 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
         # Init state
         state: LinkState = self.state_factory.state()
 
+        # Init store
+        store = self.objectstore_factory.store()
+
         # Abort if the run is not running
         abort_if(
             request.run_id,
@@ -195,6 +199,18 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
             message_ids=set(request.message_ids)
         )
 
+        # Register messages generated by LinkState in the Store for consistency
+        for msg_res in messages_res:
+            if msg_res.metadata.src_node_id == SUPERLINK_NODE_ID:
+                descendants = list(get_desdendant_object_ids(msg_res))
+                message_obj_id = msg_res.metadata.message_id
+                # Store mapping
+                store.set_message_descendant_ids(
+                    msg_object_id=message_obj_id, descendant_ids=descendants
+                )
+                # Preregister
+                store.preregister(descendants + [message_obj_id])
+
         # Delete the instruction Messages and their replies if found
         message_ins_ids_to_delete = {
             msg_res.metadata.reply_to_message_id for msg_res in messages_res
@@ -204,6 +220,7 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
 
         # Convert Messages to proto
         messages_list = []
+        objects_to_pull: dict[str, ObjectIDs] = {}
         while messages_res:
             msg = messages_res.pop(0)
 
@@ -216,7 +233,21 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
                 )
             messages_list.append(message_to_proto(msg))
 
-        return PullResMessagesResponse(messages_list=messages_list)
+            try:
+                msg_object_id = msg.metadata.message_id
+                descendants = store.get_message_descendant_ids(msg_object_id)
+                # Include the object_id of the message itself
+                objects_to_pull[msg_object_id] = ObjectIDs(
+                    object_ids=descendants + [msg_object_id]
+                )
+            except NoObjectInStoreError as e:
+                log(ERROR, e.message)
+                # Delete message ins from state
+                state.delete_messages(message_ins_ids={msg_object_id})
+
+        return PullResMessagesResponse(
+            messages_list=messages_list, objects_to_pull=objects_to_pull
+        )
 
     def GetRun(
         self, request: GetRunRequest, context: grpc.ServicerContext

flwr/supercore/object_store/__init__.py

@@ -14,10 +14,11 @@
 # ==============================================================================
 """Flower ObjectStore."""
 
-from .object_store import ObjectStore
+from .object_store import NoObjectInStoreError, ObjectStore
 from .object_store_factory import ObjectStoreFactory
 
 __all__ = [
+    "NoObjectInStoreError",
     "ObjectStore",
     "ObjectStoreFactory",
 ]