flwr 1.22.0__py3-none-any.whl → 1.24.0__py3-none-any.whl

flwr/server/superlink/linkstate/in_memory_linkstate.py

@@ -15,19 +15,16 @@
 """In-memory LinkState implementation."""
 
 
-import secrets
 import threading
-import time
 from bisect import bisect_right
 from collections import defaultdict
+from collections.abc import Sequence
 from dataclasses import dataclass, field
+from datetime import datetime, timezone
 from logging import ERROR, WARNING
-from typing import Optional
 
 from flwr.common import Context, Message, log, now
 from flwr.common.constant import (
-    FLWR_APP_TOKEN_LENGTH,
-    HEARTBEAT_MAX_INTERVAL,
     HEARTBEAT_PATIENCE,
     MESSAGE_TTL_TOLERANCE,
     NODE_ID_NUM_BYTES,
@@ -39,8 +36,13 @@ from flwr.common.constant import (
 )
 from flwr.common.record import ConfigRecord
 from flwr.common.typing import Run, RunStatus, UserConfig
+from flwr.proto.node_pb2 import NodeInfo  # pylint: disable=E0611
 from flwr.server.superlink.linkstate.linkstate import LinkState
 from flwr.server.utils import validate_message
+from flwr.supercore.constant import NodeStatus
+from flwr.supercore.corestate.in_memory_corestate import InMemoryCoreState
+from flwr.supercore.object_store.object_store import ObjectStore
+from flwr.superlink.federation import FederationManager
 
 from .utils import (
     check_node_availability_for_in_message,
@@ -57,22 +59,23 @@ class RunRecord:  # pylint: disable=R0902
     """The record of a specific run, including its status and timestamps."""
 
     run: Run
-    active_until: float = 0.0
-    heartbeat_interval: float = 0.0
     logs: list[tuple[float, str]] = field(default_factory=list)
     log_lock: threading.Lock = field(default_factory=threading.Lock)
     lock: threading.RLock = field(default_factory=threading.RLock)
 
 
-class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
+class InMemoryLinkState(LinkState, InMemoryCoreState):  # pylint: disable=R0902,R0904
     """In-memory LinkState implementation."""
 
-    def __init__(self) -> None:
+    def __init__(
+        self, federation_manager: FederationManager, object_store: ObjectStore
+    ) -> None:
+        super().__init__(object_store)
 
-        # Map node_id to (online_until, heartbeat_interval)
-        self.node_ids: dict[int, tuple[float, float]] = {}
-        self.public_key_to_node_id: dict[bytes, int] = {}
-        self.node_id_to_public_key: dict[int, bytes] = {}
+        # Map node_id to NodeInfo
+        self.nodes: dict[int, NodeInfo] = {}
+        self.node_public_key_to_node_id: dict[bytes, int] = {}
+        self.owner_to_node_ids: dict[str, set[int]] = {}  # Quick lookup
 
         # Map run_id to RunRecord
         self.run_ids: dict[int, RunRecord] = {}
@@ -82,19 +85,21 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         self.message_res_store: dict[str, Message] = {}
         self.message_ins_id_to_message_res_id: dict[str, str] = {}
 
-        # Store run ID to token mapping and token to run ID mapping
-        self.token_store: dict[int, str] = {}
-        self.token_to_run_id: dict[str, int] = {}
-        self.lock_token_store = threading.Lock()
-
         # Map flwr_aid to run_ids for O(1) reverse index lookup
         self.flwr_aid_to_run_ids: dict[str, set[int]] = defaultdict(set)
 
         self.node_public_keys: set[bytes] = set()
 
         self.lock = threading.RLock()
+        federation_manager.linkstate = self
+        self._federation_manager = federation_manager
+
+    @property
+    def federation_manager(self) -> FederationManager:
+        """Get the FederationManager instance."""
+        return self._federation_manager
 
-    def store_message_ins(self, message: Message) -> Optional[str]:
+    def store_message_ins(self, message: Message) -> str | None:
         """Store one Message."""
         # Validate message
         errors = validate_message(message, is_reply_message=False)
@@ -105,6 +110,7 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         if message.metadata.run_id not in self.run_ids:
             log(ERROR, "Invalid run ID for Message: %s", message.metadata.run_id)
             return None
+        federation = self.run_ids[message.metadata.run_id].run.federation
         # Validate source node ID
         if message.metadata.src_node_id != SUPERLINK_NODE_ID:
             log(
@@ -114,7 +120,15 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
             )
             return None
         # Validate destination node ID
-        if message.metadata.dst_node_id not in self.node_ids:
+        dst_node = self.nodes.get(message.metadata.dst_node_id)
+        if (
+            # Node must exist
+            dst_node is None
+            # Node must be online or offline
+            or dst_node.status not in (NodeStatus.ONLINE, NodeStatus.OFFLINE)
+            # Node must belong to the same federation
+            or not self.federation_manager.has_node(dst_node.node_id, federation)
+        ):
             log(
                 ERROR,
                 "Invalid destination node ID for Message: %s",
@@ -129,21 +143,50 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         # Return the new message_id
         return message_id
 
-    def get_message_ins(self, node_id: int, limit: Optional[int]) -> list[Message]:
+    def _check_stored_messages(self, message_ids: set[str]) -> None:
+        """Check and delete the message if it's invalid."""
+        with self.lock:
+            invalid_msg_ids: set[str] = set()
+            current = now().timestamp()
+            for msg_id in message_ids:
+                if not (message := self.message_ins_store.get(msg_id)):
+                    continue
+
+                # Check if the message has expired
+                available_until = message.metadata.created_at + message.metadata.ttl
+                if available_until <= current:
+                    invalid_msg_ids.add(msg_id)
+                    continue
+
+                # Check if the destination node and the source node are still in the
+                # same federation
+                src_node_id = message.metadata.src_node_id
+                dst_node_id = message.metadata.dst_node_id
+                filtered = self.federation_manager.filter_nodes(
+                    {src_node_id, dst_node_id},
+                    self.run_ids[message.metadata.run_id].run.federation,
+                )
+                if len(filtered) != 2:  # Not both nodes are in the federation
+                    invalid_msg_ids.add(msg_id)
+
+            # Delete all invalid messages
+            self.delete_messages(invalid_msg_ids)
+
+    def get_message_ins(self, node_id: int, limit: int | None) -> list[Message]:
         """Get all Messages that have not been delivered yet."""
         if limit is not None and limit < 1:
             raise AssertionError("`limit` must be >= 1")
 
         # Find Message for node_id that were not delivered yet
         message_ins_list: list[Message] = []
-        current_time = time.time()
         with self.lock:
-            for _, msg_ins in self.message_ins_store.items():
+            for msg_id in list(self.message_ins_store.keys()):
+                self._check_stored_messages({msg_id})
+
                 if (
-                    msg_ins.metadata.dst_node_id == node_id
+                    (msg_ins := self.message_ins_store.get(msg_id))
+                    and msg_ins.metadata.dst_node_id == node_id
                     and msg_ins.metadata.delivered_at == ""
-                    and msg_ins.metadata.created_at + msg_ins.metadata.ttl
-                    > current_time
                 ):
                     message_ins_list.append(msg_ins)
                 if limit and len(message_ins_list) == limit:
@@ -158,7 +201,7 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         return message_ins_list
 
     # pylint: disable=R0911
-    def store_message_res(self, message: Message) -> Optional[str]:
+    def store_message_res(self, message: Message) -> str | None:
         """Store one Message."""
         # Validate message
         errors = validate_message(message, is_reply_message=True)
@@ -170,6 +213,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
+            self._check_stored_messages({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
@@ -189,22 +233,13 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
                 )
                 return None
 
-            ins_metadata = msg_ins.metadata
-            if ins_metadata.created_at + ins_metadata.ttl <= time.time():
-                log(
-                    ERROR,
-                    "Failed to store Message: the message it is replying to "
-                    "(with ID %s) has expired",
-                    msg_ins_id,
-                )
-                return None
-
             # Fail if the Message TTL exceeds the
             # expiration time of the Message it replies to.
             # Condition: ins_metadata.created_at + ins_metadata.ttl ≥
             #            res_metadata.created_at + res_metadata.ttl
             # A small tolerance is introduced to account
             # for floating-point precision issues.
+            ins_metadata = msg_ins.metadata
             max_allowed_ttl = (
                 ins_metadata.created_at + ins_metadata.ttl - res_metadata.created_at
             )
@@ -238,7 +273,8 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         ret: dict[str, Message] = {}
 
         with self.lock:
-            current = time.time()
+            self._check_stored_messages(message_ids)
+            current = now().timestamp()
 
             # Verify Message IDs
             ret = verify_message_ids(
@@ -256,9 +292,10 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
                 inquired_in_message_ids=message_ids,
                 found_in_message_dict=self.message_ins_store,
                 node_id_to_online_until={
-                    node_id: self.node_ids[node_id][0]
+                    node_id: self.nodes[node_id].online_until
                     for node_id in dst_node_ids
-                    if node_id in self.node_ids
+                    if node_id in self.nodes
+                    and self.nodes[node_id].status != NodeStatus.UNREGISTERED
                 },
                 current_time=current,
             )
@@ -330,7 +367,13 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         """
         return len(self.message_res_store)
 
-    def create_node(self, heartbeat_interval: float) -> int:
+    def create_node(
+        self,
+        owner_aid: str,
+        owner_name: str,
+        public_key: bytes,
+        heartbeat_interval: float,
+    ) -> int:
         """Create, store in the link state, and return `node_id`."""
         # Sample a random int64 as node_id
         node_id = generate_rand_int_from_bytes(
@@ -338,28 +381,89 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         )
 
         with self.lock:
-            if node_id in self.node_ids:
+            if node_id in self.nodes:
                 log(ERROR, "Unexpected node registration failure.")
                 return 0
+            if public_key in self.node_public_key_to_node_id:
+                raise ValueError("Public key already in use")
 
-            # Mark the node online until time.time() + heartbeat_interval
-            self.node_ids[node_id] = (
-                time.time() + heartbeat_interval,
-                heartbeat_interval,
+            # The node is not activated upon creation
+            self.nodes[node_id] = NodeInfo(
+                node_id=node_id,
+                owner_aid=owner_aid,
+                owner_name=owner_name,
+                status=NodeStatus.REGISTERED,
+                registered_at=now().isoformat(),
+                last_activated_at=None,
+                last_deactivated_at=None,
+                unregistered_at=None,
+                online_until=None,
+                heartbeat_interval=heartbeat_interval,
+                public_key=public_key,
             )
+            self.node_public_key_to_node_id[public_key] = node_id
+            self.owner_to_node_ids.setdefault(owner_aid, set()).add(node_id)
             return node_id
 
-    def delete_node(self, node_id: int) -> None:
+    def delete_node(self, owner_aid: str, node_id: int) -> None:
         """Delete a node."""
         with self.lock:
-            if node_id not in self.node_ids:
-                raise ValueError(f"Node {node_id} not found")
+            if (
+                not (node := self.nodes.get(node_id))
+                or node.status == NodeStatus.UNREGISTERED
+                or owner_aid != self.nodes[node_id].owner_aid
+            ):
+                raise ValueError(
+                    f"Node ID {node_id} already unregistered, not found or "
+                    "the request was unauthorized."
+                )
+
+            node.status = NodeStatus.UNREGISTERED
+            current = now()
+            node.unregistered_at = current.isoformat()
+            # Set online_until to current timestamp on deletion, if it is in the future
+            node.online_until = min(node.online_until, current.timestamp())
+
+    def activate_node(self, node_id: int, heartbeat_interval: float) -> bool:
+        """Activate the node with the specified `node_id`."""
+        with self.lock:
+            self._check_and_tag_offline_nodes(node_ids=[node_id])
+
+            # Check if the node exists
+            if not (node := self.nodes.get(node_id)):
+                return False
+
+            # Only activate if the node is currently registered or offline
+            current_dt = now()
+            if node.status in (NodeStatus.REGISTERED, NodeStatus.OFFLINE):
+                node.status = NodeStatus.ONLINE
+                node.last_activated_at = current_dt.isoformat()
+                node.online_until = (
+                    current_dt.timestamp() + HEARTBEAT_PATIENCE * heartbeat_interval
+                )
+                node.heartbeat_interval = heartbeat_interval
+                return True
+            return False
+
+    def deactivate_node(self, node_id: int) -> bool:
+        """Deactivate the node with the specified `node_id`."""
+        with self.lock:
+            self._check_and_tag_offline_nodes(node_ids=[node_id])
+
+            # Check if the node exists
+            if not (node := self.nodes.get(node_id)):
+                return False
 
-            # Remove node ID <> public key mappings
-            if pk := self.node_id_to_public_key.pop(node_id, None):
-                del self.public_key_to_node_id[pk]
+            # Only deactivate if the node is currently online
+            current_dt = now()
+            if node.status == NodeStatus.ONLINE:
+                node.status = NodeStatus.OFFLINE
+                node.last_deactivated_at = current_dt.isoformat()
 
-            del self.node_ids[node_id]
+                # Set online_until to current timestamp
+                node.online_until = current_dt.timestamp()
+                return True
+            return False
 
     def get_nodes(self, run_id: int) -> set[int]:
         """Return all available nodes.
@@ -372,48 +476,83 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         with self.lock:
             if run_id not in self.run_ids:
                 return set()
-            current_time = time.time()
-            return {
-                node_id
-                for node_id, (online_until, _) in self.node_ids.items()
-                if online_until > current_time
+            federation = self.run_ids[run_id].run.federation
+            node_ids = {
+                node.node_id
+                for node in self.get_node_info(statuses=[NodeStatus.ONLINE])
             }
+            return self.federation_manager.filter_nodes(node_ids, federation)
 
-    def set_node_public_key(self, node_id: int, public_key: bytes) -> None:
-        """Set `public_key` for the specified `node_id`."""
+    def get_node_info(
+        self,
+        *,
+        node_ids: Sequence[int] | None = None,
+        owner_aids: Sequence[str] | None = None,
+        statuses: Sequence[str] | None = None,
+    ) -> Sequence[NodeInfo]:
+        """Retrieve information about nodes based on the specified filters."""
         with self.lock:
-            if node_id not in self.node_ids:
-                raise ValueError(f"Node {node_id} not found")
-
-            if public_key in self.public_key_to_node_id:
-                raise ValueError("Public key already in use")
-
-            self.public_key_to_node_id[public_key] = node_id
-            self.node_id_to_public_key[node_id] = public_key
-
-    def get_node_public_key(self, node_id: int) -> Optional[bytes]:
+            self._check_and_tag_offline_nodes()
+            result = []
+            for node_id in self.nodes.keys() if node_ids is None else node_ids:
+                if (node := self.nodes.get(node_id)) is None:
+                    continue
+                if owner_aids is not None and node.owner_aid not in owner_aids:
+                    continue
+                if statuses is not None and node.status not in statuses:
+                    continue
+                result.append(node)
+            return result
+
+    def _check_and_tag_offline_nodes(self, node_ids: list[int] | None = None) -> None:
+        with self.lock:
+            # Set all nodes of "online" status to "offline" if they've offline
+            current_ts = now().timestamp()
+            for node_id in node_ids or self.nodes.keys():
+                if (node := self.nodes.get(node_id)) is None:
+                    continue
+                if node.status == NodeStatus.ONLINE:
+                    if node.online_until <= current_ts:
+                        node.status = NodeStatus.OFFLINE
+                        node.last_deactivated_at = datetime.fromtimestamp(
+                            node.online_until, tz=timezone.utc
+                        ).isoformat()
+
+    def get_node_public_key(self, node_id: int) -> bytes:
         """Get `public_key` for the specified `node_id`."""
         with self.lock:
-            if node_id not in self.node_ids:
-                raise ValueError(f"Node {node_id} not found")
+            if (
+                node := self.nodes.get(node_id)
+            ) is None or node.status == NodeStatus.UNREGISTERED:
+                raise ValueError(f"Node ID {node_id} not found")
+            return node.public_key
+
+    def get_node_id_by_public_key(self, public_key: bytes) -> int | None:
+        """Get `node_id` for the specified `public_key` if it exists and is not
+        deleted."""
+        with self.lock:
+            node_id = self.node_public_key_to_node_id.get(public_key)
 
-            return self.node_id_to_public_key.get(node_id)
+            if node_id is None:
+                return None
 
-    def get_node_id(self, node_public_key: bytes) -> Optional[int]:
-        """Retrieve stored `node_id` filtered by `node_public_keys`."""
-        return self.public_key_to_node_id.get(node_public_key)
+            node_info = self.nodes[node_id]
+            if node_info.status == NodeStatus.UNREGISTERED:
+                return None
+            return node_id
 
     # pylint: disable=too-many-arguments,too-many-positional-arguments
     def create_run(
         self,
-        fab_id: Optional[str],
-        fab_version: Optional[str],
-        fab_hash: Optional[str],
+        fab_id: str | None,
+        fab_version: str | None,
+        fab_hash: str | None,
         override_config: UserConfig,
+        federation: str,
         federation_options: ConfigRecord,
-        flwr_aid: Optional[str],
+        flwr_aid: str | None,
     ) -> int:
-        """Create a new run for the specified `fab_hash`."""
+        """Create a new run."""
         # Sample a random int64 as run_id
         with self.lock:
             run_id = generate_rand_int_from_bytes(RUN_ID_NUM_BYTES)
@@ -436,6 +575,7 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
                             details="",
                         ),
                         flwr_aid=flwr_aid if flwr_aid else "",
+                        federation=federation,
                     ),
                 )
                 self.run_ids[run_id] = run_record
@@ -449,27 +589,7 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         log(ERROR, "Unexpected run creation failure.")
         return 0
 
-    def clear_supernode_auth_keys(self) -> None:
-        """Clear stored `node_public_keys` in the link state if any."""
-        with self.lock:
-            self.node_public_keys.clear()
-
-    def store_node_public_keys(self, public_keys: set[bytes]) -> None:
-        """Store a set of `node_public_keys` in the link state."""
-        with self.lock:
-            self.node_public_keys.update(public_keys)
-
-    def store_node_public_key(self, public_key: bytes) -> None:
-        """Store a `node_public_key` in the link state."""
-        with self.lock:
-            self.node_public_keys.add(public_key)
-
-    def get_node_public_keys(self) -> set[bytes]:
-        """Retrieve all currently stored `node_public_keys` as a set."""
-        with self.lock:
-            return self.node_public_keys.copy()
-
-    def get_run_ids(self, flwr_aid: Optional[str]) -> set[int]:
+    def get_run_ids(self, flwr_aid: str | None) -> set[int]:
         """Retrieve all run IDs if `flwr_aid` is not specified.
 
         Otherwise, retrieve all run IDs for the specified `flwr_aid`.
@@ -480,30 +600,10 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
                 return set(self.flwr_aid_to_run_ids.get(flwr_aid, ()))
             return set(self.run_ids.keys())
 
-    def _check_and_tag_inactive_run(self, run_ids: set[int]) -> None:
-        """Check if any runs are no longer active.
-
-        Marks runs with status 'starting' or 'running' as failed
-        if they have not sent a heartbeat before `active_until`.
-        """
-        current = now()
-        for record in (self.run_ids.get(run_id) for run_id in run_ids):
-            if record is None:
-                continue
-            with record.lock:
-                if record.run.status.status in (Status.STARTING, Status.RUNNING):
-                    if record.active_until < current.timestamp():
-                        record.run.status = RunStatus(
-                            status=Status.FINISHED,
-                            sub_status=SubStatus.FAILED,
-                            details=RUN_FAILURE_DETAILS_NO_HEARTBEAT,
-                        )
-                        record.run.finished_at = now().isoformat()
-
-    def get_run(self, run_id: int) -> Optional[Run]:
+    def get_run(self, run_id: int) -> Run | None:
         """Retrieve information about the run with the specified `run_id`."""
-        # Check if runs are still active
-        self._check_and_tag_inactive_run(run_ids={run_id})
+        # Clean up expired tokens; this will flag inactive runs as needed
+        self._cleanup_expired_tokens()
 
         with self.lock:
             if run_id not in self.run_ids:
@@ -513,8 +613,8 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
 
     def get_run_status(self, run_ids: set[int]) -> dict[int, RunStatus]:
         """Retrieve the statuses for the specified runs."""
-        # Check if runs are still active
-        self._check_and_tag_inactive_run(run_ids=run_ids)
+        # Clean up expired tokens; this will flag inactive runs as needed
+        self._cleanup_expired_tokens()
 
         with self.lock:
             return {
@@ -525,8 +625,8 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
 
     def update_run_status(self, run_id: int, new_status: RunStatus) -> bool:
         """Update the status of the run with the specified `run_id`."""
-        # Check if runs are still active
-        self._check_and_tag_inactive_run(run_ids={run_id})
+        # Clean up expired tokens; this will flag inactive runs as needed
+        self._cleanup_expired_tokens()
 
         with self.lock:
             # Check if the run_id exists
@@ -556,17 +656,9 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
                 )
                 return False
 
-            # Initialize heartbeat_interval and active_until
-            # when switching to starting or running
+            # Update the run status
             current = now()
             run_record = self.run_ids[run_id]
-            if new_status.status in (Status.STARTING, Status.RUNNING):
-                run_record.heartbeat_interval = HEARTBEAT_MAX_INTERVAL
-                run_record.active_until = (
-                    current.timestamp() + run_record.heartbeat_interval
-                )
-
-            # Update the run status
             if new_status.status == Status.STARTING:
                 run_record.run.starting_at = current.isoformat()
             elif new_status.status == Status.RUNNING:
@@ -576,7 +668,7 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
             run_record.run.status = new_status
             return True
 
-    def get_pending_run_id(self) -> Optional[int]:
+    def get_pending_run_id(self) -> int | None:
         """Get the `run_id` of a run with `Status.PENDING` status, if any."""
         pending_run_id = None
 
@@ -589,7 +681,7 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
 
         return pending_run_id
 
-    def get_federation_options(self, run_id: int) -> Optional[ConfigRecord]:
+    def get_federation_options(self, run_id: int) -> ConfigRecord | None:
         """Retrieve the federation options for the specified `run_id`."""
         with self.lock:
             if run_id not in self.run_ids:
@@ -608,52 +700,46 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
         the node is marked as offline.
         """
         with self.lock:
-            if node_id in self.node_ids:
-                self.node_ids[node_id] = (
-                    time.time() + HEARTBEAT_PATIENCE * heartbeat_interval,
-                    heartbeat_interval,
+            if (
+                node := self.nodes.get(node_id)
+            ) and node.status != NodeStatus.UNREGISTERED:
+                current_dt = now()
+
+                # Set timestamp if the status changes
+                if node.status != NodeStatus.ONLINE:  # offline or registered
+                    node.status = NodeStatus.ONLINE
+                    node.last_activated_at = current_dt.isoformat()
+
+                # Refresh `online_until` and `heartbeat_interval`
+                node.online_until = (
+                    current_dt.timestamp() + HEARTBEAT_PATIENCE * heartbeat_interval
                 )
+                node.heartbeat_interval = heartbeat_interval
                 return True
-        return False
+            return False
 
-    def acknowledge_app_heartbeat(self, run_id: int, heartbeat_interval: float) -> bool:
-        """Acknowledge a heartbeat received from a ServerApp for a given run.
+    def _on_tokens_expired(self, expired_records: list[tuple[int, float]]) -> None:
+        """Transition runs with expired tokens to failed status.
 
-        A run with status `"running"` is considered alive as long as it sends heartbeats
-        within the tolerated interval: HEARTBEAT_PATIENCE × heartbeat_interval.
-        HEARTBEAT_PATIENCE = N allows for N-1 missed heartbeat before the run is
-        marked as `"completed:failed"`.
+        Parameters
+        ----------
+        expired_records : list[tuple[int, float]]
+            List of tuples containing (run_id, active_until timestamp)
+            for expired tokens.
         """
-        with self.lock:
-            # Search for the run
-            record = self.run_ids.get(run_id)
-
-            # Check if the run_id exists
-            if record is None:
-                log(ERROR, "`run_id` is invalid")
-                return False
-
-        with record.lock:
-            # Check if runs are still active
-            self._check_and_tag_inactive_run(run_ids={run_id})
-
-            # Check if the run is of status "running"/"starting"
-            current_status = record.run.status
-            if current_status.status not in (Status.RUNNING, Status.STARTING):
-                log(
-                    ERROR,
-                    'Cannot acknowledge heartbeat for run with status "%s"',
-                    current_status.status,
+        for run_id, active_until in expired_records:
+            if not (run_record := self.run_ids.get(run_id)):
+                continue
+            with run_record.lock:
+                run_record.run.status = RunStatus(
+                    status=Status.FINISHED,
+                    sub_status=SubStatus.FAILED,
+                    details=RUN_FAILURE_DETAILS_NO_HEARTBEAT,
                 )
-                return False
+                active_until_dt = datetime.fromtimestamp(active_until, tz=timezone.utc)
+                run_record.run.finished_at = active_until_dt.isoformat()
 
-            # Update the `active_until` and `heartbeat_interval` for the given run
-            current = now().timestamp()
-            record.active_until = current + HEARTBEAT_PATIENCE * heartbeat_interval
-            record.heartbeat_interval = heartbeat_interval
-            return True
-
-    def get_serverapp_context(self, run_id: int) -> Optional[Context]:
+    def get_serverapp_context(self, run_id: int) -> Context | None:
         """Get the context for the specified `run_id`."""
         return self.contexts.get(run_id)
 
@@ -672,7 +758,7 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
             run.logs.append((now().timestamp(), log_message))
 
     def get_serverapp_log(
-        self, run_id: int, after_timestamp: Optional[float]
+        self, run_id: int, after_timestamp: float | None
     ) -> tuple[str, float]:
         """Get the serverapp logs for the specified `run_id`."""
         if run_id not in self.run_ids:
@@ -685,30 +771,3 @@ class InMemoryLinkState(LinkState):  # pylint: disable=R0902,R0904
             index = bisect_right(run.logs, (after_timestamp, ""))
             latest_timestamp = run.logs[-1][0] if index < len(run.logs) else 0.0
             return "".join(log for _, log in run.logs[index:]), latest_timestamp
-
-    def create_token(self, run_id: int) -> Optional[str]:
-        """Create a token for the given run ID."""
-        token = secrets.token_hex(FLWR_APP_TOKEN_LENGTH)  # Generate a random token
-        with self.lock_token_store:
-            if run_id in self.token_store:
-                return None  # Token already created for this run ID
-            self.token_store[run_id] = token
-            self.token_to_run_id[token] = run_id
-        return token
-
-    def verify_token(self, run_id: int, token: str) -> bool:
-        """Verify a token for the given run ID."""
-        with self.lock_token_store:
-            return self.token_store.get(run_id) == token
-
-    def delete_token(self, run_id: int) -> None:
-        """Delete the token for the given run ID."""
-        with self.lock_token_store:
-            token = self.token_store.pop(run_id, None)
-            if token is not None:
-                self.token_to_run_id.pop(token, None)
-
-    def get_run_id_by_token(self, token: str) -> Optional[int]:
-        """Get the run ID associated with a given token."""
-        with self.lock_token_store:
-            return self.token_to_run_id.get(token)