flwr-nightly 1.19.0.dev20250601__py3-none-any.whl → 1.19.0.dev20250603__py3-none-any.whl

flwr/proto/recorddict_pb2.pyi

@@ -197,23 +197,34 @@ global___ConfigRecordValue = ConfigRecordValue
 
 class ArrayRecord(google.protobuf.message.Message):
     DESCRIPTOR: google.protobuf.descriptor.Descriptor
-    DATA_KEYS_FIELD_NUMBER: builtins.int
-    DATA_VALUES_FIELD_NUMBER: builtins.int
-    @property
-    def data_keys(self) -> google.protobuf.internal.containers.RepeatedScalarFieldContainer[typing.Text]: ...
+    class Item(google.protobuf.message.Message):
+        DESCRIPTOR: google.protobuf.descriptor.Descriptor
+        KEY_FIELD_NUMBER: builtins.int
+        VALUE_FIELD_NUMBER: builtins.int
+        key: typing.Text
+        @property
+        def value(self) -> global___Array: ...
+        def __init__(self,
+            *,
+            key: typing.Text = ...,
+            value: typing.Optional[global___Array] = ...,
+            ) -> None: ...
+        def HasField(self, field_name: typing_extensions.Literal["value",b"value"]) -> builtins.bool: ...
+        def ClearField(self, field_name: typing_extensions.Literal["key",b"key","value",b"value"]) -> None: ...
+
+    ITEMS_FIELD_NUMBER: builtins.int
     @property
-    def data_values(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___Array]: ...
+    def items(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___ArrayRecord.Item]: ...
     def __init__(self,
         *,
-        data_keys: typing.Optional[typing.Iterable[typing.Text]] = ...,
-        data_values: typing.Optional[typing.Iterable[global___Array]] = ...,
+        items: typing.Optional[typing.Iterable[global___ArrayRecord.Item]] = ...,
         ) -> None: ...
-    def ClearField(self, field_name: typing_extensions.Literal["data_keys",b"data_keys","data_values",b"data_values"]) -> None: ...
+    def ClearField(self, field_name: typing_extensions.Literal["items",b"items"]) -> None: ...
 global___ArrayRecord = ArrayRecord
 
 class MetricRecord(google.protobuf.message.Message):
     DESCRIPTOR: google.protobuf.descriptor.Descriptor
-    class DataEntry(google.protobuf.message.Message):
+    class Item(google.protobuf.message.Message):
         DESCRIPTOR: google.protobuf.descriptor.Descriptor
         KEY_FIELD_NUMBER: builtins.int
         VALUE_FIELD_NUMBER: builtins.int
@@ -228,19 +239,19 @@ class MetricRecord(google.protobuf.message.Message):
         def HasField(self, field_name: typing_extensions.Literal["value",b"value"]) -> builtins.bool: ...
         def ClearField(self, field_name: typing_extensions.Literal["key",b"key","value",b"value"]) -> None: ...
 
-    DATA_FIELD_NUMBER: builtins.int
+    ITEMS_FIELD_NUMBER: builtins.int
     @property
-    def data(self) -> google.protobuf.internal.containers.MessageMap[typing.Text, global___MetricRecordValue]: ...
+    def items(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___MetricRecord.Item]: ...
     def __init__(self,
         *,
-        data: typing.Optional[typing.Mapping[typing.Text, global___MetricRecordValue]] = ...,
+        items: typing.Optional[typing.Iterable[global___MetricRecord.Item]] = ...,
         ) -> None: ...
-    def ClearField(self, field_name: typing_extensions.Literal["data",b"data"]) -> None: ...
+    def ClearField(self, field_name: typing_extensions.Literal["items",b"items"]) -> None: ...
 global___MetricRecord = MetricRecord
 
 class ConfigRecord(google.protobuf.message.Message):
     DESCRIPTOR: google.protobuf.descriptor.Descriptor
-    class DataEntry(google.protobuf.message.Message):
+    class Item(google.protobuf.message.Message):
         DESCRIPTOR: google.protobuf.descriptor.Descriptor
         KEY_FIELD_NUMBER: builtins.int
         VALUE_FIELD_NUMBER: builtins.int
@@ -255,77 +266,48 @@ class ConfigRecord(google.protobuf.message.Message):
         def HasField(self, field_name: typing_extensions.Literal["value",b"value"]) -> builtins.bool: ...
         def ClearField(self, field_name: typing_extensions.Literal["key",b"key","value",b"value"]) -> None: ...
 
-    DATA_FIELD_NUMBER: builtins.int
+    ITEMS_FIELD_NUMBER: builtins.int
     @property
-    def data(self) -> google.protobuf.internal.containers.MessageMap[typing.Text, global___ConfigRecordValue]: ...
+    def items(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___ConfigRecord.Item]: ...
     def __init__(self,
         *,
-        data: typing.Optional[typing.Mapping[typing.Text, global___ConfigRecordValue]] = ...,
+        items: typing.Optional[typing.Iterable[global___ConfigRecord.Item]] = ...,
         ) -> None: ...
-    def ClearField(self, field_name: typing_extensions.Literal["data",b"data"]) -> None: ...
+    def ClearField(self, field_name: typing_extensions.Literal["items",b"items"]) -> None: ...
 global___ConfigRecord = ConfigRecord
 
 class RecordDict(google.protobuf.message.Message):
     DESCRIPTOR: google.protobuf.descriptor.Descriptor
-    class ArraysEntry(google.protobuf.message.Message):
+    class Item(google.protobuf.message.Message):
         DESCRIPTOR: google.protobuf.descriptor.Descriptor
         KEY_FIELD_NUMBER: builtins.int
-        VALUE_FIELD_NUMBER: builtins.int
+        ARRAY_RECORD_FIELD_NUMBER: builtins.int
+        METRIC_RECORD_FIELD_NUMBER: builtins.int
+        CONFIG_RECORD_FIELD_NUMBER: builtins.int
         key: typing.Text
         @property
-        def value(self) -> global___ArrayRecord: ...
-        def __init__(self,
-            *,
-            key: typing.Text = ...,
-            value: typing.Optional[global___ArrayRecord] = ...,
-            ) -> None: ...
-        def HasField(self, field_name: typing_extensions.Literal["value",b"value"]) -> builtins.bool: ...
-        def ClearField(self, field_name: typing_extensions.Literal["key",b"key","value",b"value"]) -> None: ...
-
-    class MetricsEntry(google.protobuf.message.Message):
-        DESCRIPTOR: google.protobuf.descriptor.Descriptor
-        KEY_FIELD_NUMBER: builtins.int
-        VALUE_FIELD_NUMBER: builtins.int
-        key: typing.Text
+        def array_record(self) -> global___ArrayRecord: ...
         @property
-        def value(self) -> global___MetricRecord: ...
-        def __init__(self,
-            *,
-            key: typing.Text = ...,
-            value: typing.Optional[global___MetricRecord] = ...,
-            ) -> None: ...
-        def HasField(self, field_name: typing_extensions.Literal["value",b"value"]) -> builtins.bool: ...
-        def ClearField(self, field_name: typing_extensions.Literal["key",b"key","value",b"value"]) -> None: ...
-
-    class ConfigsEntry(google.protobuf.message.Message):
-        DESCRIPTOR: google.protobuf.descriptor.Descriptor
-        KEY_FIELD_NUMBER: builtins.int
-        VALUE_FIELD_NUMBER: builtins.int
-        key: typing.Text
+        def metric_record(self) -> global___MetricRecord: ...
         @property
-        def value(self) -> global___ConfigRecord: ...
+        def config_record(self) -> global___ConfigRecord: ...
         def __init__(self,
             *,
             key: typing.Text = ...,
-            value: typing.Optional[global___ConfigRecord] = ...,
+            array_record: typing.Optional[global___ArrayRecord] = ...,
+            metric_record: typing.Optional[global___MetricRecord] = ...,
+            config_record: typing.Optional[global___ConfigRecord] = ...,
             ) -> None: ...
-        def HasField(self, field_name: typing_extensions.Literal["value",b"value"]) -> builtins.bool: ...
-        def ClearField(self, field_name: typing_extensions.Literal["key",b"key","value",b"value"]) -> None: ...
+        def HasField(self, field_name: typing_extensions.Literal["array_record",b"array_record","config_record",b"config_record","metric_record",b"metric_record","value",b"value"]) -> builtins.bool: ...
+        def ClearField(self, field_name: typing_extensions.Literal["array_record",b"array_record","config_record",b"config_record","key",b"key","metric_record",b"metric_record","value",b"value"]) -> None: ...
+        def WhichOneof(self, oneof_group: typing_extensions.Literal["value",b"value"]) -> typing.Optional[typing_extensions.Literal["array_record","metric_record","config_record"]]: ...
 
-    ARRAYS_FIELD_NUMBER: builtins.int
-    METRICS_FIELD_NUMBER: builtins.int
-    CONFIGS_FIELD_NUMBER: builtins.int
-    @property
-    def arrays(self) -> google.protobuf.internal.containers.MessageMap[typing.Text, global___ArrayRecord]: ...
-    @property
-    def metrics(self) -> google.protobuf.internal.containers.MessageMap[typing.Text, global___MetricRecord]: ...
+    ITEMS_FIELD_NUMBER: builtins.int
     @property
-    def configs(self) -> google.protobuf.internal.containers.MessageMap[typing.Text, global___ConfigRecord]: ...
+    def items(self) -> google.protobuf.internal.containers.RepeatedCompositeFieldContainer[global___RecordDict.Item]: ...
     def __init__(self,
         *,
-        arrays: typing.Optional[typing.Mapping[typing.Text, global___ArrayRecord]] = ...,
-        metrics: typing.Optional[typing.Mapping[typing.Text, global___MetricRecord]] = ...,
-        configs: typing.Optional[typing.Mapping[typing.Text, global___ConfigRecord]] = ...,
+        items: typing.Optional[typing.Iterable[global___RecordDict.Item]] = ...,
         ) -> None: ...
-    def ClearField(self, field_name: typing_extensions.Literal["arrays",b"arrays","configs",b"configs","metrics",b"metrics"]) -> None: ...
+    def ClearField(self, field_name: typing_extensions.Literal["items",b"items"]) -> None: ...
 global___RecordDict = RecordDict

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

@@ -15,11 +15,12 @@
 """Fleet API gRPC request-response servicer."""
 
 
-from logging import DEBUG, INFO
+from logging import DEBUG, ERROR, INFO
 
 import grpc
 from google.protobuf.json_format import MessageToDict
 
+from flwr.common.constant import Status
 from flwr.common.inflatable import check_body_len_consistency
 from flwr.common.logger import log
 from flwr.common.typing import InvalidRunStatusException
@@ -49,8 +50,9 @@ from flwr.proto.run_pb2 import GetRunRequest, GetRunResponse  # pylint: disable=
 from flwr.server.superlink.ffs.ffs_factory import FfsFactory
 from flwr.server.superlink.fleet.message_handler import message_handler
 from flwr.server.superlink.linkstate import LinkStateFactory
-from flwr.server.superlink.utils import abort_grpc_context
+from flwr.server.superlink.utils import abort_grpc_context, check_abort
 from flwr.supercore.object_store import ObjectStoreFactory
+from flwr.supercore.object_store.object_store import NoObjectInStoreError
 
 
 class FleetServicer(fleet_pb2_grpc.FleetServicer):
@@ -183,11 +185,39 @@ class FleetServicer(fleet_pb2_grpc.FleetServicer):
             request.object_id,
         )
 
+        state = self.state_factory.state()
+
+        # Abort if the run is not running
+        abort_msg = check_abort(
+            request.run_id,
+            [Status.PENDING, Status.STARTING, Status.FINISHED],
+            state,
+        )
+        if abort_msg:
+            abort_grpc_context(abort_msg, context)
+
+        if request.node.node_id not in state.get_nodes(run_id=request.run_id):
+            # Cancel insertion in ObjectStore
+            context.abort(grpc.StatusCode.FAILED_PRECONDITION, "Unexpected node ID.")
+
         if not check_body_len_consistency(request.object_content):
             # Cancel insertion in ObjectStore
-            context.abort(grpc.StatusCode.PERMISSION_DENIED, "Unexpected object length")
+            context.abort(
+                grpc.StatusCode.FAILED_PRECONDITION, "Unexpected object length"
+            )
+
+        # Init store
+        store = self.objectstore_factory.store()
+
+        # Insert in store
+        stored = False
+        try:
+            store.put(request.object_id, request.object_content)
+            stored = True
+        except (NoObjectInStoreError, ValueError) as e:
+            log(ERROR, str(e))
 
-        return PushObjectResponse()
+        return PushObjectResponse(stored=stored)
 
     def PullObject(
         self, request: PullObjectRequest, context: grpc.ServicerContext
@@ -199,4 +229,31 @@ class FleetServicer(fleet_pb2_grpc.FleetServicer):
             request.object_id,
         )
 
-        return PullObjectResponse()
+        state = self.state_factory.state()
+
+        # Abort if the run is not running
+        abort_msg = check_abort(
+            request.run_id,
+            [Status.PENDING, Status.STARTING, Status.FINISHED],
+            state,
+        )
+        if abort_msg:
+            abort_grpc_context(abort_msg, context)
+
+        if request.node.node_id not in state.get_nodes(run_id=request.run_id):
+            # Cancel insertion in ObjectStore
+            context.abort(grpc.StatusCode.FAILED_PRECONDITION, "Unexpected node ID.")
+
+        # Init store
+        store = self.objectstore_factory.store()
+
+        # Fetch from store
+        content = store.get(request.object_id)
+        if content is not None:
+            object_available = content != b""
+            return PullObjectResponse(
+                object_found=True,
+                object_available=object_available,
+                object_content=content,
+            )
+        return PullObjectResponse(object_found=False, object_available=False)

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

@@ -114,7 +114,7 @@ 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()
+    store: ObjectStore = cast(ObjectStoreFactory, app.state.OBJECTSTORE_FACTORY).store()
 
     # Handle message
     return message_handler.pull_messages(request=request, state=state, store=store)
@@ -125,7 +125,7 @@ async def push_message(request: PushMessagesRequest) -> PushMessagesResponse:
     """Pull PushMessages."""
     # Get state from app
     state: LinkState = cast(LinkStateFactory, app.state.STATE_FACTORY).state()
-    store: ObjectStore = cast(ObjectStoreFactory, app.state.STATE_FACTORY).store()
+    store: ObjectStore = cast(ObjectStoreFactory, app.state.OBJECTSTORE_FACTORY).store()
 
     # Handle message
     return message_handler.push_messages(request=request, state=state, store=store)

flwr/server/superlink/serverappio/serverappio_servicer.py

@@ -409,11 +409,39 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
         """Push an object to the ObjectStore."""
         log(DEBUG, "ServerAppIoServicer.PushObject")
 
+        # Init state
+        state: LinkState = self.state_factory.state()
+
+        # Abort if the run is not running
+        abort_if(
+            request.run_id,
+            [Status.PENDING, Status.STARTING, Status.FINISHED],
+            state,
+            context,
+        )
+
+        if request.node.node_id != SUPERLINK_NODE_ID:
+            # Cancel insertion in ObjectStore
+            context.abort(grpc.StatusCode.FAILED_PRECONDITION, "Unexpected node ID.")
+
         if not check_body_len_consistency(request.object_content):
             # Cancel insertion in ObjectStore
-            context.abort(grpc.StatusCode.PERMISSION_DENIED, "Unexpected object length")
+            context.abort(
+                grpc.StatusCode.FAILED_PRECONDITION, "Unexpected object length."
+            )
+
+        # Init store
+        store = self.objectstore_factory.store()
+
+        # Insert in store
+        stored = False
+        try:
+            store.put(request.object_id, request.object_content)
+            stored = True
+        except (NoObjectInStoreError, ValueError) as e:
+            log(ERROR, str(e))
 
-        return PushObjectResponse()
+        return PushObjectResponse(stored=stored)
 
     def PullObject(
         self, request: PullObjectRequest, context: grpc.ServicerContext
@@ -421,7 +449,34 @@ class ServerAppIoServicer(serverappio_pb2_grpc.ServerAppIoServicer):
         """Pull an object from the ObjectStore."""
         log(DEBUG, "ServerAppIoServicer.PullObject")
 
-        return PullObjectResponse()
+        # Init state
+        state: LinkState = self.state_factory.state()
+
+        # Abort if the run is not running
+        abort_if(
+            request.run_id,
+            [Status.PENDING, Status.STARTING, Status.FINISHED],
+            state,
+            context,
+        )
+
+        if request.node.node_id != SUPERLINK_NODE_ID:
+            # Cancel insertion in ObjectStore
+            context.abort(grpc.StatusCode.FAILED_PRECONDITION, "Unexpected node ID.")
+
+        # Init store
+        store = self.objectstore_factory.store()
+
+        # Fetch from store
+        content = store.get(request.object_id)
+        if content is not None:
+            object_available = content != b""
+            return PullObjectResponse(
+                object_found=True,
+                object_available=object_available,
+                object_content=content,
+            )
+        return PullObjectResponse(object_found=False, object_available=False)
 
 
 def _raise_if(validation_error: bool, request_name: str, detail: str) -> None:

flwr/supernode/cli/__init__.py

@@ -16,5 +16,9 @@
 
 
 from .flower_supernode import flower_supernode
+from .flwr_clientapp import flwr_clientapp
 
-__all__ = ["flower_supernode"]
+__all__ = [
+    "flower_supernode",
+    "flwr_clientapp",
+]

flwr/supernode/cli/flower_supernode.py

@@ -42,8 +42,7 @@ from flwr.common.constant import (
 from flwr.common.exit import ExitCode, flwr_exit
 from flwr.common.exit_handlers import register_exit_handlers
 from flwr.common.logger import log
-
-from ..start_client_internal import start_client_internal
+from flwr.supernode.start_client_internal import start_client_internal
 
 
 def flower_supernode() -> None:

flwr/supernode/cli/flwr_clientapp.py

@@ -0,0 +1,73 @@
+# Copyright 2025 Flower Labs GmbH. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""`flwr-clientapp` command."""
+
+
+import argparse
+from logging import DEBUG, INFO
+
+from flwr.common.args import add_args_flwr_app_common
+from flwr.common.constant import CLIENTAPPIO_API_DEFAULT_CLIENT_ADDRESS
+from flwr.common.exit import ExitCode, flwr_exit
+from flwr.common.logger import log
+from flwr.supernode.runtime.run_clientapp import run_clientapp
+
+
+def flwr_clientapp() -> None:
+    """Run process-isolated Flower ClientApp."""
+    args = _parse_args_run_flwr_clientapp().parse_args()
+    if not args.insecure:
+        flwr_exit(
+            ExitCode.COMMON_TLS_NOT_SUPPORTED,
+            "flwr-clientapp does not support TLS yet.",
+        )
+
+    log(INFO, "Start `flwr-clientapp` process")
+    log(
+        DEBUG,
+        "`flwr-clientapp` will attempt to connect to SuperNode's "
+        "ClientAppIo API at %s with token %s",
+        args.clientappio_api_address,
+        args.token,
+    )
+    run_clientapp(
+        clientappio_api_address=args.clientappio_api_address,
+        run_once=(args.token is not None),
+        token=args.token,
+        flwr_dir=args.flwr_dir,
+        certificates=None,
+    )
+
+
+def _parse_args_run_flwr_clientapp() -> argparse.ArgumentParser:
+    """Parse flwr-clientapp command line arguments."""
+    parser = argparse.ArgumentParser(
+        description="Run a Flower ClientApp",
+    )
+    parser.add_argument(
+        "--clientappio-api-address",
+        default=CLIENTAPPIO_API_DEFAULT_CLIENT_ADDRESS,
+        type=str,
+        help="Address of SuperNode's ClientAppIo API (IPv4, IPv6, or a domain name)."
+        f"By default, it is set to {CLIENTAPPIO_API_DEFAULT_CLIENT_ADDRESS}.",
+    )
+    parser.add_argument(
+        "--token",
+        type=int,
+        required=False,
+        help="Unique token generated by SuperNode for each ClientApp execution",
+    )
+    add_args_flwr_app_common(parser=parser)
+    return parser

flwr/supernode/nodestate/in_memory_nodestate.py

@@ -15,17 +15,40 @@
 """In-memory NodeState implementation."""
 
 
+from collections.abc import Sequence
+from dataclasses import dataclass
+from threading import Lock
 from typing import Optional
 
+from flwr.common import Context, Message
+from flwr.common.typing import Run
+
 from .nodestate import NodeState
 
 
+@dataclass
+class MessageEntry:
+    """Data class to represent a message entry."""
+
+    message: Message
+    is_retrieved: bool = False
+
+
 class InMemoryNodeState(NodeState):
     """In-memory NodeState implementation."""
 
     def __init__(self) -> None:
         # Store node_id
         self.node_id: Optional[int] = None
+        # Store Object ID to MessageEntry mapping
+        self.msg_store: dict[str, MessageEntry] = {}
+        self.lock_msg_store = Lock()
+        # Store run ID to Run mapping
+        self.run_store: dict[int, Run] = {}
+        self.lock_run_store = Lock()
+        # Store run ID to Context mapping
+        self.ctx_store: dict[int, Context] = {}
+        self.lock_ctx_store = Lock()
 
     def set_node_id(self, node_id: Optional[int]) -> None:
         """Set the node ID."""
@@ -36,3 +59,92 @@ class InMemoryNodeState(NodeState):
         if self.node_id is None:
             raise ValueError("Node ID not set")
         return self.node_id
+
+    def store_message(self, message: Message) -> Optional[str]:
+        """Store a message."""
+        with self.lock_msg_store:
+            msg_id = message.metadata.message_id
+            if msg_id == "" or msg_id in self.msg_store:
+                return None
+            self.msg_store[msg_id] = MessageEntry(message=message)
+            return msg_id
+
+    def get_messages(
+        self,
+        *,
+        run_ids: Optional[Sequence[int]] = None,
+        is_reply: Optional[bool] = None,
+        limit: Optional[int] = None,
+    ) -> Sequence[Message]:
+        """Retrieve messages based on the specified filters."""
+        selected_messages: list[Message] = []
+
+        with self.lock_msg_store:
+            # Iterate through all messages in the store
+            for object_id in list(self.msg_store.keys()):
+                entry = self.msg_store[object_id]
+                message = entry.message
+
+                # Skip messages that have already been retrieved
+                if entry.is_retrieved:
+                    continue
+
+                # Skip messages whose run_id doesn't match the filter
+                if run_ids is not None:
+                    if message.metadata.run_id not in run_ids:
+                        continue
+
+                # If is_reply filter is set, filter for reply/non-reply messages
+                if is_reply is not None:
+                    is_reply_message = message.metadata.reply_to_message_id != ""
+                    # XOR logic to filter mismatched types (reply vs non-reply)
+                    if is_reply ^ is_reply_message:
+                        continue
+
+                # Add the message to the result set
+                selected_messages.append(message)
+
+                # Mark the message as retrieved
+                entry.is_retrieved = True
+
+                # Stop if the number of collected messages reaches the limit
+                if limit is not None and len(selected_messages) >= limit:
+                    break
+
+        return selected_messages
+
+    def delete_messages(
+        self,
+        *,
+        message_ids: Optional[Sequence[str]] = None,
+    ) -> None:
+        """Delete messages based on the specified filters."""
+        with self.lock_msg_store:
+            if message_ids is None:
+                # If no message IDs are provided, clear the entire store
+                self.msg_store.clear()
+                return
+
+            # Remove specified messages from the store
+            for msg_id in message_ids:
+                self.msg_store.pop(msg_id, None)
+
+    def store_run(self, run: Run) -> None:
+        """Store a run."""
+        with self.lock_run_store:
+            self.run_store[run.run_id] = run
+
+    def get_run(self, run_id: int) -> Optional[Run]:
+        """Retrieve a run by its ID."""
+        with self.lock_run_store:
+            return self.run_store.get(run_id)
+
+    def store_context(self, context: Context) -> None:
+        """Store a context."""
+        with self.lock_ctx_store:
+            self.ctx_store[context.run_id] = context
+
+    def get_context(self, run_id: int) -> Optional[Context]:
+        """Retrieve a context by its run ID."""
+        with self.lock_ctx_store:
+            return self.ctx_store.get(run_id)