flwr-nightly 1.23.0.dev20250930__py3-none-any.whl → 1.26.0.dev20260121__py3-none-any.whl

flwr/simulation/ray_transport/ray_client_proxy.py

@@ -17,19 +17,14 @@
 
 import traceback
 from logging import ERROR
-from typing import Optional
 
 from flwr import common
+from flwr.app.message_type import MessageType
 from flwr.client import ClientFnExt
-from flwr.client.client_app import ClientApp
 from flwr.client.run_info_store import DeprecatedRunInfoStore
+from flwr.clientapp.client_app import ClientApp
 from flwr.common import DEFAULT_TTL, Message, Metadata, RecordDict, now
-from flwr.common.constant import (
-    NUM_PARTITIONS_KEY,
-    PARTITION_ID_KEY,
-    MessageType,
-    MessageTypeLegacy,
-)
+from flwr.common.constant import NUM_PARTITIONS_KEY, PARTITION_ID_KEY, MessageTypeLegacy
 from flwr.common.logger import log
 from flwr.common.message import make_message
 from flwr.common.recorddict_compat import (
@@ -74,7 +69,7 @@ class RayActorClientProxy(ClientProxy):
             },
         )
 
-    def _submit_job(self, message: Message, timeout: Optional[float]) -> Message:
+    def _submit_job(self, message: Message, timeout: float | None) -> Message:
         """Sumbit a message to the ActorPool."""
         run_id = message.metadata.run_id
 
@@ -114,8 +109,8 @@ class RayActorClientProxy(ClientProxy):
         self,
         recorddict: RecordDict,
         message_type: str,
-        timeout: Optional[float],
-        group_id: Optional[int],
+        timeout: float | None,
+        group_id: int | None,
     ) -> Message:
         """Wrap a RecordDict inside a Message."""
         return make_message(
@@ -136,8 +131,8 @@ class RayActorClientProxy(ClientProxy):
     def get_properties(
         self,
         ins: common.GetPropertiesIns,
-        timeout: Optional[float],
-        group_id: Optional[int],
+        timeout: float | None,
+        group_id: int | None,
     ) -> common.GetPropertiesRes:
         """Return client's properties."""
         recorddict = getpropertiesins_to_recorddict(ins)
@@ -155,8 +150,8 @@ class RayActorClientProxy(ClientProxy):
     def get_parameters(
         self,
         ins: common.GetParametersIns,
-        timeout: Optional[float],
-        group_id: Optional[int],
+        timeout: float | None,
+        group_id: int | None,
     ) -> common.GetParametersRes:
         """Return the current local model parameters."""
         recorddict = getparametersins_to_recorddict(ins)
@@ -172,7 +167,7 @@ class RayActorClientProxy(ClientProxy):
         return recorddict_to_getparametersres(message_out.content, keep_input=False)
 
     def fit(
-        self, ins: common.FitIns, timeout: Optional[float], group_id: Optional[int]
+        self, ins: common.FitIns, timeout: float | None, group_id: int | None
     ) -> common.FitRes:
         """Train model parameters on the locally held dataset."""
         recorddict = fitins_to_recorddict(
@@ -190,7 +185,7 @@ class RayActorClientProxy(ClientProxy):
         return recorddict_to_fitres(message_out.content, keep_input=False)
 
     def evaluate(
-        self, ins: common.EvaluateIns, timeout: Optional[float], group_id: Optional[int]
+        self, ins: common.EvaluateIns, timeout: float | None, group_id: int | None
     ) -> common.EvaluateRes:
         """Evaluate model parameters on the locally held dataset."""
         recorddict = evaluateins_to_recorddict(
@@ -210,8 +205,8 @@ class RayActorClientProxy(ClientProxy):
     def reconnect(
         self,
         ins: common.ReconnectIns,
-        timeout: Optional[float],
-        group_id: Optional[int],
+        timeout: float | None,
+        group_id: int | None,
     ) -> common.DisconnectRes:
         """Disconnect and (optionally) reconnect later."""
         return common.DisconnectRes(reason="")  # Nothing to do here (yet)

flwr/simulation/run_simulation.py

@@ -26,11 +26,12 @@ import traceback
 from logging import DEBUG, ERROR, INFO, WARNING
 from pathlib import Path
 from queue import Empty, Queue
-from typing import Any, Optional
+from typing import Any, cast
 
+from flwr.app.user_config import UserConfig
 from flwr.cli.config_utils import load_and_validate
 from flwr.cli.utils import get_sha256_hash
-from flwr.client import ClientApp
+from flwr.clientapp import ClientApp
 from flwr.common import Context, EventType, RecordDict, event, log, now
 from flwr.common.config import get_fused_config_from_dir, parse_config_args
 from flwr.common.constant import RUN_ID_NUM_BYTES, Status
@@ -39,7 +40,7 @@ from flwr.common.logger import (
     update_console_handler,
     warn_deprecated_feature_with_example,
 )
-from flwr.common.typing import Run, RunStatus, UserConfig
+from flwr.common.typing import Run, RunStatus
 from flwr.server.grid import Grid, InMemoryGrid
 from flwr.server.run_serverapp import run as _run
 from flwr.server.server_app import ServerApp
@@ -51,6 +52,9 @@ from flwr.server.superlink.linkstate.utils import generate_rand_int_from_bytes
 from flwr.simulation.ray_transport.utils import (
     enable_tf_gpu_growth as enable_gpu_growth,
 )
+from flwr.supercore.constant import FLWR_IN_MEMORY_DB_NAME, NOOP_FEDERATION
+from flwr.supercore.object_store import ObjectStoreFactory
+from flwr.superlink.federation import NoOpFederationManager
 
 
 def _replace_keys(d: Any, match: str, target: str) -> Any:
@@ -98,12 +102,7 @@ def run_simulation_from_cli() -> None:
     _check_ray_support(args.backend)
 
     # Load JSON config
-    backend_config_dict = json.loads(args.backend_config)
-
-    if backend_config_dict:
-        # Backend config internally operates with `_` not with `-`
-        backend_config_dict = _replace_keys(backend_config_dict, match="-", target="_")
-        log(DEBUG, "backend_config_dict: %s", backend_config_dict)
+    backend_config = json.loads(args.backend_config)
 
     run_id = (
         generate_rand_int_from_bytes(RUN_ID_NUM_BYTES)
@@ -141,6 +140,7 @@ def run_simulation_from_cli() -> None:
 
     # Create run
     run = Run.create_empty(run_id)
+    run.federation = NOOP_FEDERATION
     run.override_config = override_config
 
     # Create Context
@@ -157,7 +157,7 @@ def run_simulation_from_cli() -> None:
         client_app_attr=client_app_attr,
         num_supernodes=args.num_supernodes,
         backend_name=args.backend,
-        backend_config=backend_config_dict,
+        backend_config=backend_config,
         app_dir=args.app,
         run=run,
         enable_tf_gpu_growth=args.enable_tf_gpu_growth,
@@ -175,7 +175,7 @@ def run_simulation(
     client_app: ClientApp,
     num_supernodes: int,
     backend_name: str = "ray",
-    backend_config: Optional[BackendConfig] = None,
+    backend_config: BackendConfig | None = None,
     enable_tf_gpu_growth: bool = False,
     verbose_logging: bool = False,
 ) -> None:
@@ -248,8 +248,8 @@ def run_simulation(
 
 # pylint: disable=too-many-arguments,too-many-positional-arguments
 def run_serverapp_th(
-    server_app_attr: Optional[str],
-    server_app: Optional[ServerApp],
+    server_app_attr: str | None,
+    server_app: ServerApp | None,
     server_app_context: Context,
     grid: Grid,
     app_dir: str,
@@ -266,8 +266,8 @@ def run_serverapp_th(
         exception_event: threading.Event,
         _grid: Grid,
         _server_app_dir: str,
-        _server_app_attr: Optional[str],
-        _server_app: Optional[ServerApp],
+        _server_app_attr: str | None,
+        _server_app: ServerApp | None,
         _ctx_queue: "Queue[Context]",
     ) -> None:
         """Run SeverApp, after check if GPU memory growth has to be set.
@@ -327,16 +327,18 @@ def _main_loop(
     enable_tf_gpu_growth: bool,
     run: Run,
     exit_event: EventType,
-    flwr_dir: Optional[str] = None,
-    client_app: Optional[ClientApp] = None,
-    client_app_attr: Optional[str] = None,
-    server_app: Optional[ServerApp] = None,
-    server_app_attr: Optional[str] = None,
-    server_app_context: Optional[Context] = None,
+    flwr_dir: str | None = None,
+    client_app: ClientApp | None = None,
+    client_app_attr: str | None = None,
+    server_app: ServerApp | None = None,
+    server_app_attr: str | None = None,
+    server_app_context: Context | None = None,
 ) -> Context:
     """Start ServerApp on a separate thread, then launch Simulation Engine."""
     # Initialize StateFactory
-    state_factory = LinkStateFactory(":flwr-in-memory-state:")
+    state_factory = LinkStateFactory(
+        FLWR_IN_MEMORY_DB_NAME, NoOpFederationManager(), ObjectStoreFactory()
+    )
 
     f_stop = threading.Event()
     # A Threading event to indicate if an exception was raised in the ServerApp thread
@@ -427,16 +429,16 @@ def _main_loop(
 def _run_simulation(
     num_supernodes: int,
     exit_event: EventType,
-    client_app: Optional[ClientApp] = None,
-    server_app: Optional[ServerApp] = None,
+    client_app: ClientApp | None = None,
+    server_app: ServerApp | None = None,
     backend_name: str = "ray",
-    backend_config: Optional[BackendConfig] = None,
-    client_app_attr: Optional[str] = None,
-    server_app_attr: Optional[str] = None,
-    server_app_context: Optional[Context] = None,
+    backend_config: BackendConfig | None = None,
+    client_app_attr: str | None = None,
+    server_app_attr: str | None = None,
+    server_app_context: Context | None = None,
     app_dir: str = "",
-    flwr_dir: Optional[str] = None,
-    run: Optional[Run] = None,
+    flwr_dir: str | None = None,
+    run: Run | None = None,
     enable_tf_gpu_growth: bool = False,
     verbose_logging: bool = False,
     is_app: bool = False,
@@ -444,29 +446,28 @@ def _run_simulation(
     """Launch the Simulation Engine."""
     if backend_config is None:
         backend_config = {}
+    elif backend_config:
+        # Backend config internally operates with `_` not with `-`
+        backend_config = cast(
+            BackendConfig, _replace_keys(backend_config, match="-", target="_")
+        )
+        log(DEBUG, "backend_config: %s", backend_config)
 
-    if "init_args" not in backend_config:
-        backend_config["init_args"] = {}
-
+    # Set default init_args if not passed
+    backend_config.setdefault("init_args", {})
     # Set default client_resources if not passed
-    if "client_resources" not in backend_config:
-        backend_config["client_resources"] = {"num_cpus": 2, "num_gpus": 0}
-
+    backend_config.setdefault("client_resources", {"num_cpus": 2, "num_gpus": 0})
     # Initialization of backend config to enable GPU growth globally when set
-    if "actor" not in backend_config:
-        backend_config["actor"] = {"tensorflow": 0}
+    backend_config.setdefault("actor", {"tensorflow": 0})
 
     # Set logging level
     logger = logging.getLogger("flwr")
     if verbose_logging:
         update_console_handler(level=DEBUG, timestamps=True, colored=True)
     else:
-        backend_config["init_args"]["logging_level"] = backend_config["init_args"].get(
-            "logging_level", WARNING
-        )
-        backend_config["init_args"]["log_to_driver"] = backend_config["init_args"].get(
-            "log_to_driver", True
-        )
+        init_args = backend_config["init_args"]
+        init_args.setdefault("logging_level", WARNING)
+        init_args.setdefault("log_to_driver", True)
 
     if enable_tf_gpu_growth:
         # Check that Backend config has also enabled using GPU growth
@@ -482,6 +483,7 @@ def _run_simulation(
     if run is None:
         run_id = generate_rand_int_from_bytes(RUN_ID_NUM_BYTES)
         run = Run.create_empty(run_id=run_id)
+        run.federation = NOOP_FEDERATION
 
     args = (
         num_supernodes,

flwr/simulation/simulationio_connection.py

@@ -16,7 +16,7 @@
 
 
 from logging import DEBUG, WARNING
-from typing import Optional, cast
+from typing import cast
 
 import grpc
 
@@ -43,12 +43,12 @@ class SimulationIoConnection:
     def __init__(  # pylint: disable=too-many-arguments
         self,
         simulationio_service_address: str = SIMULATIONIO_API_DEFAULT_CLIENT_ADDRESS,
-        root_certificates: Optional[bytes] = None,
+        root_certificates: bytes | None = None,
     ) -> None:
         self._addr = simulationio_service_address
         self._cert = root_certificates
-        self._grpc_stub: Optional[SimulationIoStub] = None
-        self._channel: Optional[grpc.Channel] = None
+        self._grpc_stub: SimulationIoStub | None = None
+        self._channel: grpc.Channel | None = None
         self._retry_invoker = _make_simple_grpc_retry_invoker()
 
     @property

flwr/{common → supercore}/address.py

@@ -15,17 +15,13 @@
 """Flower IP address utils."""
 
 
-import re
 import socket
 from ipaddress import ip_address
-from typing import Optional
-
-import grpc
 
 IPV6: int = 6
 
 
-def parse_address(address: str) -> Optional[tuple[str, int, Optional[bool]]]:
+def parse_address(address: str) -> tuple[str, int, bool | None] | None:
     """Parse an IP address into host, port, and version.
 
     Parameters
@@ -104,35 +100,3 @@ def is_port_in_use(address: str) -> bool:
             return True
 
         return False
-
-
-def get_ip_address_from_servicer_context(context: grpc.ServicerContext) -> str:
-    """Extract the client's IPv4 or IPv6 address from the gRPC ServicerContext.
-
-    Parameters
-    ----------
-    context : grpc.ServicerContext
-        The gRPC ServicerContext object. The context.peer() returns a string like
-        "ipv4:127.0.0.1:56789" for IPv4 and "ipv6:[2001:db8::1]:54321" for IPv6.
-
-    Returns
-    -------
-    str
-        If one of the format matches, the function will return the client's IP address,
-        otherwise, it will raise a ValueError.
-    """
-    peer: str = context.peer()
-    # Match IPv4: "ipv4:IP:port"
-    ipv4_match = re.match(r"^ipv4:(?P<ip>[^:]+):", peer)
-    if ipv4_match:
-        return ipv4_match.group("ip")
-
-    # Match IPv6: "ipv6:[IP]:port"
-    ipv6_match = re.match(r"^ipv6:\[(?P<ip>[^\]]+)\]:", peer)
-    if ipv6_match:
-        return ipv6_match.group("ip")
-
-    raise ValueError(
-        f"Unsupported peer address format: {peer} for the transport protocol. "
-        "The supported formats are ipv4:IP:port and ipv6:[IP]:port."
-    )

flwr/supercore/cli/flower_superexec.py

@@ -17,7 +17,7 @@
 
 import argparse
 from logging import INFO
-from typing import Any, Optional
+from typing import Any
 
 import yaml
 
@@ -54,7 +54,7 @@ except ImportError:
 
     def get_ee_plugin_and_stub_class(  # pylint: disable=unused-argument
         plugin_type: str,
-    ) -> Optional[tuple[type[ExecPlugin], type[object]]]:
+    ) -> tuple[type[ExecPlugin], type[object]] | None:
         """Get the EE plugin class and stub class based on the plugin type."""
         return None
 
@@ -75,7 +75,6 @@ def flower_superexec() -> None:
     # Log the first message after parsing arguments in case of `--help`
     log(INFO, "Starting Flower SuperExec")
 
-    # Trigger telemetry event
     event(EventType.RUN_SUPEREXEC_ENTER, {"plugin_type": args.plugin_type})
 
     # Load plugin config from YAML file if provided
@@ -83,7 +82,7 @@ def flower_superexec() -> None:
     if plugin_config_path := getattr(args, "plugin_config", None):
         try:
             with open(plugin_config_path, encoding="utf-8") as file:
-                yaml_config: Optional[dict[str, Any]] = yaml.safe_load(file)
+                yaml_config: dict[str, Any] | None = yaml.safe_load(file)
                 if yaml_config is None or EXEC_PLUGIN_SECTION not in yaml_config:
                     raise ValueError(f"Missing '{EXEC_PLUGIN_SECTION}' section.")
                 plugin_config = yaml_config[EXEC_PLUGIN_SECTION]

flwr/supercore/constant.py

@@ -15,5 +15,74 @@
 """Constants for Flower infrastructure."""
 
 
+from __future__ import annotations
+
+from flwr.common.constant import FLWR_DIR
+
 # Top-level key in YAML config for exec plugin settings
 EXEC_PLUGIN_SECTION = "exec_plugin"
+
+# Flower in-memory Python-based database name
+FLWR_IN_MEMORY_DB_NAME = ":flwr-in-memory:"
+
+# Constants for Hub
+APP_ID_PATTERN = r"^@[A-Za-z0-9_.-]+/[A-Za-z0-9_.-]+$"
+APP_VERSION_PATTERN = r"^\d+\.\d+\.\d+$"
+PLATFORM_API_URL = "https://api.flower.ai/v1"
+
+# Specification for app publishing
+APP_PUBLISH_INCLUDE_PATTERNS = (
+    "**/*.py",
+    "**/*.toml",
+    "**/*.md",
+)
+APP_PUBLISH_EXCLUDE_PATTERNS = FAB_EXCLUDE_PATTERNS = (
+    f"{FLWR_DIR}/**",  # Exclude the .flwr directory
+    "**/__pycache__/**",
+)
+MAX_TOTAL_BYTES = 10 * 1024 * 1024  # 10 MB
+MAX_FILE_BYTES = 1 * 1024 * 1024  # 1 MB
+MAX_FILE_COUNT = 1000
+MAX_DIR_DEPTH = 10  # relative depth (number of parts in relpath)
+UTF8 = "utf-8"
+MIME_MAP = {
+    ".py": "text/x-python; charset=utf-8",
+    ".md": "text/markdown; charset=utf-8",
+    ".toml": "application/toml; charset=utf-8",
+}
+
+# Constants for federations
+NOOP_FEDERATION = "default"
+
+# Constants for exit handling
+FORCE_EXIT_TIMEOUT_SECONDS = 5  # Used in `flwr_exit` function
+
+# Constants for message processing timing
+MESSAGE_TIME_ENTRY_MAX_AGE_SECONDS = 3600
+
+
+# System message type
+SYSTEM_MESSAGE_TYPE = "system"
+
+# SQLite PRAGMA settings for optimal performance and correctness
+SQLITE_PRAGMAS = (
+    ("journal_mode", "WAL"),  # Enable Write-Ahead Logging for better concurrency
+    ("synchronous", "NORMAL"),
+    ("foreign_keys", "ON"),
+    ("cache_size", "-64000"),  # 64MB cache
+    ("temp_store", "MEMORY"),  # In-memory temp tables
+    ("mmap_size", "268435456"),  # 256MB memory-mapped I/O
+)
+
+
+class NodeStatus:
+    """Event log writer types."""
+
+    REGISTERED = "registered"
+    ONLINE = "online"
+    OFFLINE = "offline"
+    UNREGISTERED = "unregistered"
+
+    def __new__(cls) -> NodeStatus:
+        """Prevent instantiation."""
+        raise TypeError(f"{cls.__name__} cannot be instantiated.")

flwr/supercore/corestate/corestate.py

@@ -16,14 +16,20 @@
 
 
 from abc import ABC, abstractmethod
-from typing import Optional
+
+from ..object_store import ObjectStore
 
 
 class CoreState(ABC):
     """Abstract base class for core state."""
 
+    @property
+    @abstractmethod
+    def object_store(self) -> ObjectStore:
+        """Return the ObjectStore instance used by this CoreState."""
+
     @abstractmethod
-    def create_token(self, run_id: int) -> Optional[str]:
+    def create_token(self, run_id: int) -> str | None:
         """Create a token for the given run ID.
 
         Parameters
@@ -66,7 +72,7 @@ class CoreState(ABC):
         """
 
     @abstractmethod
-    def get_run_id_by_token(self, token: str) -> Optional[int]:
+    def get_run_id_by_token(self, token: str) -> int | None:
         """Get the run ID associated with a given token.
 
         Parameters
@@ -79,3 +85,18 @@ class CoreState(ABC):
         Optional[int]
             The run ID if the token is valid, otherwise None.
         """
+
+    @abstractmethod
+    def acknowledge_app_heartbeat(self, token: str) -> bool:
+        """Acknowledge an app heartbeat with the provided token.
+
+        Parameters
+        ----------
+        token : str
+            The token associated with the app.
+
+        Returns
+        -------
+        bool
+            True if the heartbeat is acknowledged successfully, False otherwise.
+        """

flwr/supercore/corestate/in_memory_corestate.py

@@ -0,0 +1,138 @@
+# 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.
+# ==============================================================================
+"""In-memory CoreState implementation."""
+
+
+import secrets
+from dataclasses import dataclass
+from threading import Lock
+
+from flwr.common import now
+from flwr.common.constant import (
+    FLWR_APP_TOKEN_LENGTH,
+    HEARTBEAT_DEFAULT_INTERVAL,
+    HEARTBEAT_PATIENCE,
+)
+
+from ..object_store import ObjectStore
+from .corestate import CoreState
+
+
+@dataclass
+class TokenRecord:
+    """Record containing token and heartbeat information."""
+
+    token: str
+    active_until: float
+
+
+class InMemoryCoreState(CoreState):
+    """In-memory CoreState implementation."""
+
+    def __init__(self, object_store: ObjectStore) -> None:
+        self._object_store = object_store
+        # Store run ID to token mapping and token to run ID mapping
+        self.token_store: dict[int, TokenRecord] = {}
+        self.token_to_run_id: dict[str, int] = {}
+        self.lock_token_store = Lock()
+
+    @property
+    def object_store(self) -> ObjectStore:
+        """Return the ObjectStore instance used by this CoreState."""
+        return self._object_store
+
+    def create_token(self, run_id: int) -> str | None:
+        """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] = TokenRecord(
+                token=token, active_until=now().timestamp() + HEARTBEAT_DEFAULT_INTERVAL
+            )
+            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."""
+        self._cleanup_expired_tokens()
+        with self.lock_token_store:
+            record = self.token_store.get(run_id)
+            return record is not None and record.token == token
+
+    def delete_token(self, run_id: int) -> None:
+        """Delete the token for the given run ID."""
+        with self.lock_token_store:
+            record = self.token_store.pop(run_id, None)
+            if record is not None:
+                self.token_to_run_id.pop(record.token, None)
+
+    def get_run_id_by_token(self, token: str) -> int | None:
+        """Get the run ID associated with a given token."""
+        self._cleanup_expired_tokens()
+        with self.lock_token_store:
+            return self.token_to_run_id.get(token)
+
+    def acknowledge_app_heartbeat(self, token: str) -> bool:
+        """Acknowledge an app heartbeat with the provided token."""
+        # Clean up expired tokens
+        self._cleanup_expired_tokens()
+
+        with self.lock_token_store:
+            # Return False if token is not found
+            if token not in self.token_to_run_id:
+                return False
+
+            # Get the run_id and update heartbeat info
+            run_id = self.token_to_run_id[token]
+            record = self.token_store[run_id]
+            current = now().timestamp()
+            record.active_until = (
+                current + HEARTBEAT_PATIENCE * HEARTBEAT_DEFAULT_INTERVAL
+            )
+            return True
+
+    def _cleanup_expired_tokens(self) -> None:
+        """Remove expired tokens and perform additional cleanup.
+
+        This method is called before token operations to ensure integrity.
+        Subclasses can override `_on_tokens_expired` to add custom cleanup logic.
+        """
+        with self.lock_token_store:
+            current = now().timestamp()
+            expired_records: list[tuple[int, float]] = []
+            for run_id, record in list(self.token_store.items()):
+                if record.active_until < current:
+                    expired_records.append((run_id, record.active_until))
+                    # Remove from both stores
+                    del self.token_store[run_id]
+                    self.token_to_run_id.pop(record.token, None)
+
+            # Hook for subclasses
+            if expired_records:
+                self._on_tokens_expired(expired_records)
+
+    def _on_tokens_expired(self, expired_records: list[tuple[int, float]]) -> None:
+        """Handle cleanup of expired tokens.
+
+        Override in subclasses to add custom cleanup logic.
+
+        Parameters
+        ----------
+        expired_records : list[tuple[int, float]]
+            List of tuples containing (run_id, active_until timestamp)
+            for expired tokens.
+        """