flwr 1.25.0__py3-none-any.whl → 1.26.0__py3-none-any.whl

flwr/cli/utils.py

@@ -18,20 +18,25 @@
 import hashlib
 import json
 import re
+import sys
 from collections.abc import Callable, Iterable, Iterator
 from contextlib import contextmanager
+from io import StringIO
 from pathlib import Path
 from typing import Any, cast
 
+import click
 import grpc
 import pathspec
 import typer
+from rich.console import Console
 
+from flwr.cli.typing import SuperLinkConnection
 from flwr.common.constant import (
     ACCESS_TOKEN_KEY,
     AUTHN_TYPE_JSON_KEY,
-    CREDENTIALS_DIR,
-    FLWR_DIR,
+    FEDERATION_NOT_FOUND_MESSAGE,
+    FEDERATION_NOT_SPECIFIED_MESSAGE,
     NO_ACCOUNT_AUTH_MESSAGE,
     NO_ARTIFACT_PROVIDER_MESSAGE,
     NODE_NOT_FOUND_MESSAGE,
@@ -41,16 +46,67 @@ from flwr.common.constant import (
     REFRESH_TOKEN_KEY,
     RUN_ID_NOT_FOUND_MESSAGE,
     AuthnType,
+    CliOutputFormat,
 )
 from flwr.common.grpc import (
     GRPC_MAX_MESSAGE_LENGTH,
     create_channel,
     on_channel_state_change,
 )
+from flwr.common.logger import print_json_error, redirect_output, restore_output
+from flwr.supercore.credential_store import get_credential_store
 
 from .auth_plugin import CliAuthPlugin, get_cli_plugin_class
 from .cli_account_auth_interceptor import CliAccountAuthInterceptor
-from .config_utils import validate_certificate_in_federation_config
+from .config_utils import load_certificate_in_connection
+from .constant import AUTHN_TYPE_STORE_KEY
+
+
+def print_json_to_stdout(data: str | Any) -> None:
+    """Print JSON data to stdout, bypassing any output redirection.
+
+    Use this function within the `cli_output_handler` context manager to print JSON
+    output directly to the terminal, even when stdout is being captured.
+    """
+    if isinstance(data, str):
+        Console(file=sys.__stdout__).print_json(data)
+    else:
+        Console(file=sys.__stdout__).print_json(data=data)
+
+
+@contextmanager  # docsig: ignore=SIG503
+def cli_output_handler(
+    output_format: str = CliOutputFormat.DEFAULT,
+) -> Iterator[bool]:
+    """Context manager for handling CLI output in different formats.
+
+    This context manager provides consistent output handling for CLI commands by:
+    - Redirecting stdout/stderr when JSON format is requested
+    - Catching and handling exceptions appropriately based on the output format
+
+    Use the `print_json_to_stdout()` utility function to print JSON output that bypasses
+    output redirection.
+    """
+    is_json = output_format == CliOutputFormat.JSON
+    captured_output = StringIO()
+
+    if is_json:
+        redirect_output(captured_output)
+
+    try:
+        yield is_json
+    except Exception as err:  # pylint: disable=broad-except
+        if is_json:
+            restore_output()
+            print_json_error(captured_output.getvalue(), err)
+        else:
+            if isinstance(err, typer.Exit):
+                raise  # Allow typer.Exit to escape normally
+            raise click.ClickException(str(err)) from None
+    finally:
+        if is_json:
+            restore_output()
+        captured_output.close()
 
 
 def prompt_text(
@@ -155,46 +211,6 @@ def is_valid_project_name(name: str) -> bool:
     return True
 
 
-def sanitize_project_name(name: str) -> str:
-    """Sanitize the given string to make it a valid Python project name.
-
-    This function replaces spaces, dots, slashes, and underscores with dashes, removes
-    any characters not allowed in Python project names, makes the string lowercase, and
-    ensures it starts with a valid character.
-
-    Parameters
-    ----------
-    name : str
-        The project name to sanitize.
-
-    Returns
-    -------
-    str
-        The sanitized project name that is valid for Python projects.
-    """
-    # Replace whitespace with '_'
-    name_with_hyphens = re.sub(r"[ ./_]", "-", name)
-
-    # Allowed characters in a module name: letters, digits, underscore
-    allowed_chars = set(
-        "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789-"
-    )
-
-    # Make the string lowercase
-    sanitized_name = name_with_hyphens.lower()
-
-    # Remove any characters not allowed in Python module names
-    sanitized_name = "".join(c for c in sanitized_name if c in allowed_chars)
-
-    # Ensure the first character is a letter or underscore
-    while sanitized_name and (
-        sanitized_name[0].isdigit() or sanitized_name[0] not in allowed_chars
-    ):
-        sanitized_name = sanitized_name[1:]
-
-    return sanitized_name
-
-
 def get_sha256_hash(file_path_or_int: Path | int) -> str:
     """Calculate the SHA-256 hash of a file or integer.
 
@@ -221,121 +237,27 @@ def get_sha256_hash(file_path_or_int: Path | int) -> str:
     return sha256.hexdigest()
 
 
-def get_account_auth_config_path(root_dir: Path, federation: str) -> Path:
-    """Return the path to the account auth config file.
-
-    Additionally, a `.gitignore` file will be created in the Flower directory to
-    include the `.credentials` folder to be excluded from git. If the `.gitignore`
-    file already exists, a warning will be displayed if the `.credentials` entry is
-    not found.
-    """
-    # Locate the credentials directory
-    abs_flwr_dir = root_dir.absolute() / FLWR_DIR
-    credentials_dir = abs_flwr_dir / CREDENTIALS_DIR
-    credentials_dir.mkdir(parents=True, exist_ok=True)
-
-    # Determine the absolute path of the Flower directory for .gitignore
-    gitignore_path = abs_flwr_dir / ".gitignore"
-    credential_entry = CREDENTIALS_DIR
-
-    try:
-        if gitignore_path.exists():
-            with open(gitignore_path, encoding="utf-8") as gitignore_file:
-                lines = gitignore_file.read().splitlines()
-
-            # Warn if .credentials is not already in .gitignore
-            if credential_entry not in lines:
-                typer.secho(
-                    f"`.gitignore` exists, but `{credential_entry}` entry not found. "
-                    "Consider adding it to your `.gitignore` to exclude Flower "
-                    "credentials from git.",
-                    fg=typer.colors.YELLOW,
-                    bold=True,
-                )
-        else:
-            typer.secho(
-                f"Creating a new `.gitignore` with `{credential_entry}` entry...",
-                fg=typer.colors.BLUE,
-            )
-            # Create a new .gitignore with .credentials
-            with open(gitignore_path, "w", encoding="utf-8") as gitignore_file:
-                gitignore_file.write(f"{credential_entry}\n")
-    except Exception as err:
-        typer.secho(
-            "❌ An error occurred while handling `.gitignore.` "
-            f"Please check the permissions of `{gitignore_path}` and try again.",
-            fg=typer.colors.RED,
-            bold=True,
-            err=True,
-        )
-        raise typer.Exit(code=1) from err
-
-    return credentials_dir / f"{federation}.json"
-
-
-def account_auth_enabled(federation_config: dict[str, Any]) -> bool:
-    """Check if account authentication is enabled in the federation config.
-
-    Parameters
-    ----------
-    federation_config : dict[str, Any]
-        The federation configuration dictionary.
-
-    Returns
-    -------
-    bool
-        True if account authentication is enabled, False otherwise.
-    """
-    enabled: bool = federation_config.get("enable-user-auth", False)
-    enabled |= federation_config.get("enable-account-auth", False)
-    if "enable-user-auth" in federation_config:
-        typer.secho(
-            "`enable-user-auth` is deprecated and will be removed in a future "
-            "release. Please use `enable-account-auth` instead.",
-            fg=typer.colors.YELLOW,
-            bold=True,
-        )
-    return enabled
-
-
-def retrieve_authn_type(config_path: Path) -> str:
-    """Retrieve the auth type from the config file or return NOOP if not found.
+def get_authn_type(host: str) -> str:
+    """Retrieve the authentication type for the given host from the credential store.
 
-    Parameters
-    ----------
-    config_path : Path
-        Path to the authentication configuration file.
-
-    Returns
-    -------
-    str
-        The authentication type string, or AuthnType.NOOP if not found.
+    `AuthnType.NOOP` is returned if no authentication type is found.
     """
-    try:
-        with config_path.open("r", encoding="utf-8") as file:
-            json_file = json.load(file)
-        authn_type: str = json_file[AUTHN_TYPE_JSON_KEY]
-        return authn_type
-    except (FileNotFoundError, KeyError):
+    store = get_credential_store()
+    authn_type = store.get(AUTHN_TYPE_STORE_KEY % host)
+    if authn_type is None:
         return AuthnType.NOOP
+    return authn_type.decode("utf-8")
 
 
-def load_cli_auth_plugin(
-    root_dir: Path,
-    federation: str,
-    federation_config: dict[str, Any],
-    authn_type: str | None = None,
+def load_cli_auth_plugin_from_connection(
+    host: str, authn_type: str | None = None
 ) -> CliAuthPlugin:
-    """Load the CLI-side account auth plugin for the given authn type.
+    """Load the CLI-side account auth plugin for the given connection.
 
     Parameters
     ----------
-    root_dir : Path
-        Root directory of the Flower project.
-    federation : str
-        Name of the federation.
-    federation_config : dict[str, Any]
-        Federation configuration dictionary.
+    host : str
+        The SuperLink Control API address.
     authn_type : str | None
         Authentication type. If None, will be determined from config.
 
@@ -346,41 +268,48 @@ def load_cli_auth_plugin(
 
     Raises
     ------
-    typer.Exit
+    click.ClickException
         If the authentication type is unknown.
     """
-    # Find the path to the account auth config file
-    config_path = get_account_auth_config_path(root_dir, federation)
-
     # Determine the auth type if not provided
     # Only `flwr login` command can provide `authn_type` explicitly, as it can query the
     # SuperLink for the auth type.
     if authn_type is None:
-        authn_type = AuthnType.NOOP
-        if account_auth_enabled(federation_config):
-            authn_type = retrieve_authn_type(config_path)
+        authn_type = get_authn_type(host)
 
     # Retrieve auth plugin class and instantiate it
     try:
         auth_plugin_class = get_cli_plugin_class(authn_type)
-        return auth_plugin_class(config_path)
+        return auth_plugin_class(host)
     except ValueError:
-        typer.echo(f"❌ Unknown account authentication type: {authn_type}")
-        raise typer.Exit(code=1) from None
+        raise click.ClickException(
+            f"Unknown account authentication type: {authn_type}"
+        ) from None
+
+
+def require_superlink_address(connection: SuperLinkConnection) -> str:
+    """Return the SuperLink address or exit if it is not configured."""
+    if connection.address is None:
+        cmd = click.get_current_context().command.name
+        raise click.ClickException(
+            f"`flwr {cmd}` currently works with a SuperLink. Ensure that the "
+            "correct SuperLink (Control API) address is provided SuperLink connection "
+            "you are using. Check your Flower configuration file. You may use `flwr "
+            "config list` to see its location in the file system."
+        )
+    return connection.address
 
 
-def init_channel(
-    app: Path, federation_config: dict[str, Any], auth_plugin: CliAuthPlugin
+def init_channel_from_connection(
+    connection: SuperLinkConnection, auth_plugin: CliAuthPlugin | None = None
 ) -> grpc.Channel:
     """Initialize gRPC channel to the Control API.
 
     Parameters
     ----------
-    app : Path
-        Path to the Flower app directory.
-    federation_config : dict[str, Any]
-        Federation configuration dictionary containing address and TLS settings.
-    auth_plugin : CliAuthPlugin
+    connection : SuperLinkConnection
+        SuperLink connection configuration.
+    auth_plugin : CliAuthPlugin | None (default: None)
         Authentication plugin instance for handling credentials.
 
     Returns
@@ -388,17 +317,20 @@ def init_channel(
     grpc.Channel
         Configured gRPC channel with authentication interceptors.
     """
-    insecure, root_certificates_bytes = validate_certificate_in_federation_config(
-        app, federation_config
-    )
+    address = require_superlink_address(connection)
 
+    root_certificates_bytes = load_certificate_in_connection(connection)
+
+    # Load authentication plugin
+    if auth_plugin is None:
+        auth_plugin = load_cli_auth_plugin_from_connection(address)
     # Load tokens
     auth_plugin.load_tokens()
 
     # Create the gRPC channel
     channel = create_channel(
-        server_address=federation_config["address"],
-        insecure=insecure,
+        server_address=address,
+        insecure=connection.insecure,
         root_certificates=root_certificates_bytes,
         max_message_length=GRPC_MAX_MESSAGE_LENGTH,
         interceptors=[CliAccountAuthInterceptor(auth_plugin)],
@@ -423,7 +355,7 @@ def flwr_cli_grpc_exc_handler() -> Iterator[None]:  # pylint: disable=too-many-b
 
     Raises
     ------
-    typer.Exit
+    click.ClickException
         On handled gRPC error statuses with appropriate exit code.
     grpc.RpcError
         For unhandled gRPC error statuses.
@@ -432,115 +364,74 @@ def flwr_cli_grpc_exc_handler() -> Iterator[None]:  # pylint: disable=too-many-b
         yield
     except grpc.RpcError as e:
         if e.code() == grpc.StatusCode.UNAUTHENTICATED:
-            typer.secho(
-                "❌ Authentication failed. Please run `flwr login`"
-                " to authenticate and try again.",
-                fg=typer.colors.RED,
-                bold=True,
-                err=True,
-            )
-            raise typer.Exit(code=1) from None
+            raise click.ClickException(
+                "Authentication failed. Please run `flwr login`"
+                " to authenticate and try again."
+            ) from None
         if e.code() == grpc.StatusCode.UNIMPLEMENTED:
             if e.details() == NO_ACCOUNT_AUTH_MESSAGE:  # pylint: disable=E1101
-                typer.secho(
-                    "❌ Account authentication is not enabled on this SuperLink.",
-                    fg=typer.colors.RED,
-                    bold=True,
-                    err=True,
-                )
-            elif e.details() == NO_ARTIFACT_PROVIDER_MESSAGE:  # pylint: disable=E1101
-                typer.secho(
-                    "❌ The SuperLink does not support `flwr pull` command.",
-                    fg=typer.colors.RED,
-                    bold=True,
-                    err=True,
-                )
-            else:
-                typer.secho(
-                    "❌ The SuperLink cannot process this request. Please verify that "
-                    "you set the address to its Control API endpoint correctly in your "
-                    "`pyproject.toml`, and ensure that the Flower versions used by "
-                    "the CLI and SuperLink are compatible.",
-                    fg=typer.colors.RED,
-                    bold=True,
-                    err=True,
-                )
-            raise typer.Exit(code=1) from None
+                raise click.ClickException(
+                    "Account authentication is not enabled on this SuperLink."
+                ) from None
+            if e.details() == NO_ARTIFACT_PROVIDER_MESSAGE:  # pylint: disable=E1101
+                raise click.ClickException(
+                    "The SuperLink does not support `flwr pull` command."
+                ) from None
+            raise click.ClickException(
+                "The SuperLink cannot process this request. Please verify that "
+                "you set the address to its Control API endpoint correctly in your "
+                "SuperLink connection in your Flower Configuration file. You may use "
+                "`flwr config list` to see its location in the file system. "
+                "Additonally, ensure that the Flower versions used by the CLI and "
+                "SuperLink are compatible."
+            ) from None
         if e.code() == grpc.StatusCode.PERMISSION_DENIED:
-            typer.secho(
-                "❌ Permission denied.",
-                fg=typer.colors.RED,
-                bold=True,
-                err=True,
-            )
             # pylint: disable-next=E1101
-            typer.secho(e.details(), fg=typer.colors.RED, bold=True)
-            raise typer.Exit(code=1) from None
+            raise click.ClickException(f"Permission denied.\n{e.details()}") from None
         if e.code() == grpc.StatusCode.UNAVAILABLE:
-            typer.secho(
+            raise click.ClickException(
                 "Connection to the SuperLink is unavailable. Please check your network "
-                "connection and 'address' in the federation configuration.",
-                fg=typer.colors.RED,
-                bold=True,
-                err=True,
-            )
-            raise typer.Exit(code=1) from None
+                "connection and 'address' in the SuperLink connection configuration."
+            ) from None
         if e.code() == grpc.StatusCode.NOT_FOUND:
             if e.details() == RUN_ID_NOT_FOUND_MESSAGE:  # pylint: disable=E1101
-                typer.secho(
-                    "❌ Run ID not found.",
-                    fg=typer.colors.RED,
-                    bold=True,
-                    err=True,
-                )
-                raise typer.Exit(code=1) from None
+                raise click.ClickException("Run ID not found.") from None
             if e.details() == NODE_NOT_FOUND_MESSAGE:  # pylint: disable=E1101
-                typer.secho(
-                    "❌ Node ID not found for this account.",
-                    fg=typer.colors.RED,
-                    bold=True,
-                    err=True,
-                )
-                raise typer.Exit(code=1) from None
+                raise click.ClickException(
+                    "Node ID not found for this account."
+                ) from None
         if e.code() == grpc.StatusCode.FAILED_PRECONDITION:
             if e.details() == PULL_UNFINISHED_RUN_MESSAGE:  # pylint: disable=E1101
-                typer.secho(
-                    "❌ Run is not finished yet. Artifacts can only be pulled after "
-                    "the run is finished. You can check the run status with `flwr ls`.",
-                    fg=typer.colors.RED,
-                    bold=True,
-                    err=True,
-                )
-                raise typer.Exit(code=1) from None
+                raise click.ClickException(
+                    "Run is not finished yet. Artifacts can only be pulled after "
+                    "the run is finished. You can check the run status with `flwr ls`."
+                ) from None
             if (
                 e.details() == PUBLIC_KEY_ALREADY_IN_USE_MESSAGE
             ):  # pylint: disable=E1101
-                typer.secho(
-                    "❌ The provided public key is already in use by another "
-                    "SuperNode.",
-                    fg=typer.colors.RED,
-                    bold=True,
-                    err=True,
-                )
-                raise typer.Exit(code=1) from None
+                raise click.ClickException(
+                    "The provided public key is already in use by another SuperNode."
+                ) from None
             if e.details() == PUBLIC_KEY_NOT_VALID:  # pylint: disable=E1101
-                typer.secho(
-                    "❌ The provided public key is invalid. Please provide a valid "
-                    "NIST EC public key.",
-                    fg=typer.colors.RED,
-                    bold=True,
-                    err=True,
-                )
-                raise typer.Exit(code=1) from None
+                raise click.ClickException(
+                    "The provided public key is invalid. Please provide a valid "
+                    "NIST EC public key."
+                ) from None
+            if e.details() == FEDERATION_NOT_SPECIFIED_MESSAGE:  # pylint: disable=E1101
+                raise click.ClickException(
+                    "No federation specified. "
+                    "Please use the `--federation` flag or set a default federation "
+                    "in your SuperLink connection configuration."
+                ) from None
+            patten = re.compile(FEDERATION_NOT_FOUND_MESSAGE.replace("%s", "(.+)"))
+            if m := patten.match(e.details()):  # pylint: disable=E1101
+                raise click.ClickException(
+                    f"Federation '{m.group(1)}' does not exist. "
+                    "Please verify the federation name and try again."
+                ) from None
 
             # Log details from grpc error directly
-            typer.secho(
-                f"❌ {e.details()}",
-                fg=typer.colors.RED,
-                bold=True,
-                err=True,
-            )
-            raise typer.Exit(code=1) from None
+            raise click.ClickException(f"{e.details()}") from None
         raise
 
 
@@ -600,23 +491,17 @@ def validate_credentials_content(creds_path: Path) -> str:
     try:
         creds: dict[str, str] = json.loads(creds_path.read_text(encoding="utf-8"))
     except (OSError, json.JSONDecodeError) as err:
-        typer.secho(
-            f"Invalid credentials file at '{creds_path}': {err}",
-            fg=typer.colors.RED,
-            err=True,
-        )
-        raise typer.Exit(code=1) from err
+        raise click.ClickException(
+            f"Invalid credentials file at '{creds_path}': {err}"
+        ) from err
 
     required_keys = [AUTHN_TYPE_JSON_KEY, ACCESS_TOKEN_KEY, REFRESH_TOKEN_KEY]
     missing = [key for key in required_keys if key not in creds]
 
     if missing:
-        typer.secho(
+        raise click.ClickException(
             f"Credentials file '{creds_path}' is missing "
-            f"required key(s): {', '.join(missing)}. Please log in again.",
-            fg=typer.colors.RED,
-            err=True,
+            f"required key(s): {', '.join(missing)}. Please log in again."
         )
-        raise typer.Exit(code=1)
 
     return creds[ACCESS_TOKEN_KEY]

flwr/client/grpc_rere_client/connection.py

@@ -33,7 +33,7 @@ from flwr.common.inflatable_protobuf_utils import (
 )
 from flwr.common.logger import log
 from flwr.common.message import Message, remove_content_from_message
-from flwr.common.retry_invoker import RetryInvoker, _wrap_stub
+from flwr.common.retry_invoker import RetryInvoker, wrap_stub
 from flwr.common.serde import (
     fab_from_proto,
     message_from_proto,
@@ -136,7 +136,7 @@ def grpc_request_response(  # pylint: disable=R0913,R0914,R0915,R0917
     confirm_message_received : Callable[[str], None]
     """
     if isinstance(root_certificates, str):
-        root_certificates = Path(root_certificates).read_bytes()
+        root_certificates = Path(root_certificates).expanduser().read_bytes()
 
     # Automatic node auth: generate keys if user didn't provide any
     self_registered = False
@@ -165,7 +165,7 @@ def grpc_request_response(  # pylint: disable=R0913,R0914,R0915,R0917
     node: Node | None = None
 
     # Wrap stub
-    _wrap_stub(stub, retry_invoker)
+    wrap_stub(stub, retry_invoker)
     ###########################################################################
     # SuperNode functions
     ###########################################################################

flwr/client/grpc_rere_client/grpc_adapter.py

@@ -32,7 +32,6 @@ from flwr.common.constant import (
     GRPC_ADAPTER_METADATA_MESSAGE_QUALNAME_KEY,
     GRPC_ADAPTER_METADATA_SHOULD_EXIT_KEY,
 )
-from flwr.common.version import package_name, package_version
 from flwr.proto.fab_pb2 import GetFabRequest, GetFabResponse  # pylint: disable=E0611
 from flwr.proto.fleet_pb2 import (  # pylint: disable=E0611
     ActivateNodeRequest,
@@ -64,6 +63,7 @@ from flwr.proto.message_pb2 import (  # pylint: disable=E0611
 )
 from flwr.proto.run_pb2 import GetRunRequest, GetRunResponse  # pylint: disable=E0611
 from flwr.supercore.constant import FORCE_EXIT_TIMEOUT_SECONDS
+from flwr.supercore.version import package_name, package_version
 
 T = TypeVar("T", bound=GrpcMessage)
 

flwr/client/message_handler/message_handler.py

@@ -18,6 +18,7 @@
 from logging import WARN
 from typing import cast
 
+from flwr.app.message_type import MessageType
 from flwr.client.client import (
     maybe_call_evaluate,
     maybe_call_fit,
@@ -27,7 +28,7 @@ from flwr.client.client import (
 from flwr.client.numpy_client import NumPyClient
 from flwr.client.typing import ClientFnExt
 from flwr.common import ConfigRecord, Context, Message, Metadata, RecordDict, log
-from flwr.common.constant import MessageType, MessageTypeLegacy
+from flwr.common.constant import MessageTypeLegacy
 from flwr.common.recorddict_compat import (
     evaluateres_to_recorddict,
     fitres_to_recorddict,

flwr/client/mod/centraldp_mods.py

@@ -17,10 +17,10 @@
 
 from logging import INFO
 
+from flwr.app.message_type import MessageType
 from flwr.client.typing import ClientAppCallable
 from flwr.common import ndarrays_to_parameters, parameters_to_ndarrays
 from flwr.common import recorddict_compat as compat
-from flwr.common.constant import MessageType
 from flwr.common.context import Context
 from flwr.common.differential_privacy import (
     compute_adaptive_clip_model_update,

flwr/client/mod/localdp_mod.py

@@ -19,10 +19,10 @@ from logging import INFO
 
 import numpy as np
 
+from flwr.app.message_type import MessageType
 from flwr.client.typing import ClientAppCallable
 from flwr.common import ndarrays_to_parameters, parameters_to_ndarrays
 from flwr.common import recorddict_compat as compat
-from flwr.common.constant import MessageType
 from flwr.common.context import Context
 from flwr.common.differential_privacy import (
     add_localdp_gaussian_noise_to_params,

flwr/client/mod/secure_aggregation/secaggplus_mod.py

@@ -20,6 +20,7 @@ from dataclasses import dataclass, field
 from logging import DEBUG, WARNING
 from typing import Any, cast
 
+from flwr.app.message_type import MessageType
 from flwr.client.typing import ClientAppCallable
 from flwr.common import (
     ConfigRecord,
@@ -31,7 +32,6 @@ from flwr.common import (
     parameters_to_ndarrays,
 )
 from flwr.common import recorddict_compat as compat
-from flwr.common.constant import MessageType
 from flwr.common.logger import log
 from flwr.common.secure_aggregation.crypto.shamir import create_shares
 from flwr.common.secure_aggregation.crypto.symmetric_encryption import (