_initialize.py

from __future__ import annotations

import functools
import os
import sys
import threading
import typing
from contextlib import contextmanager
from contextvars import ContextVar
from dataclasses import dataclass, field, replace
from pathlib import Path
from typing import TYPE_CHECKING, Callable, Generator, List, Literal, Optional, TypeVar

from flyte.errors import InitializationError
from flyte.syncify import syncify

from ._logging import LogFormat, initialize_logger, logger

if TYPE_CHECKING:
    from types import FunctionType

    from flyte._internal.imagebuild import ImageBuildEngine
    from flyte.config import Config
    from flyte.config._config import PlatformConfig
    from flyte.remote._client.auth import AuthType, ClientConfig
    from flyte.remote._client.controlplane import ClientSet
    from flyte.storage import Storage

Mode = Literal["local", "remote"]


@dataclass(init=True, repr=True, eq=True, frozen=True, kw_only=True)
class CommonInit:
    """
    Common initialization configuration for Flyte.
    """

    root_dir: Path
    org: str | None = None
    project: str | None = None
    domain: str | None = None
    batch_size: int = 1000
    source_config_path: Optional[Path] = None  # Only used for documentation
    sync_local_sys_paths: bool = True
    local_persistence: bool = False
    local_tracked: bool = False
    local_tracked_strict: bool = False


@dataclass(init=True, kw_only=True, repr=True, eq=True, frozen=True)
class _InitConfig(CommonInit):
    client: Optional[ClientSet] = None
    storage: Optional[Storage] = None
    image_builder: "ImageBuildEngine.ImageBuilderType" = "local"
    images: typing.Dict[str, str] = field(default_factory=dict)
    image_registry: str | None = None

    def replace(self, **kwargs) -> _InitConfig:
        return replace(self, **kwargs)


# Global singleton to store initialization configuration
_init_config: _InitConfig | None = None
_init_lock = threading.RLock()  # Reentrant lock for thread safety

# Per-context override of the module-global config. A server that has to vary the config per
# inbound request or per concurrent call -- the MCP server scopes each tool call to the
# project/domain it was given -- sets this for the duration of that call instead of mutating the
# process-wide global. Unset by default, so nothing changes for an ordinary process: every
# reader goes through ``_get_init_config()``, which falls back to the global.
_context_init_config: ContextVar[Optional[_InitConfig]] = ContextVar("_context_init_config", default=None)


def _platform_to_client_kwargs(pc: "PlatformConfig") -> dict[str, typing.Any]:
    """Translate PlatformConfig fields into kwargs accepted by both
    _initialize_client (via init) and create_remote_controller. Single
    source of truth — extending the HTTP/auth stack with a new
    PlatformConfig field means editing this one function.

    Callers may need to layer on caller-specific kwargs that are NOT
    derived from PlatformConfig (api_key, headless, in-cluster env-var
    overrides); those stay at the call site.

    ca_cert_file_path takes precedence over insecure_skip_verify when
    both are present. _resolve_tls_ca_cert otherwise prefers the
    bootstrap path, which fetches only what the server's leaf presents
    and trips "UnknownIssuer" on chains nginx serves without
    intermediates.
    """
    kw: dict[str, typing.Any] = {}
    if pc.endpoint:
        kw["endpoint"] = pc.endpoint
    if pc.insecure:
        kw["insecure"] = True
    if pc.ca_cert_file_path:
        kw["ca_cert_file_path"] = pc.ca_cert_file_path
    elif pc.insecure_skip_verify:
        kw["insecure_skip_verify"] = True
    if pc.client_id:
        kw["client_id"] = pc.client_id
    if pc.client_credentials_secret:
        kw["client_credentials_secret"] = pc.client_credentials_secret
    if pc.auth_mode:
        kw["auth_type"] = pc.auth_mode
    if pc.command:
        kw["command"] = pc.command
    if pc.proxy_command:
        kw["proxy_command"] = pc.proxy_command
    if pc.http_proxy_url:
        kw["http_proxy_url"] = pc.http_proxy_url
    if pc.rpc_retries:
        kw["rpc_retries"] = pc.rpc_retries
    # NOTE: disable_keyring is intentionally NOT emitted here. The helper
    # output flows both into _initialize_client (via init) AND into
    # create_remote_controller (via init_in_cluster). The controller's
    # constructor doesn't accept disable_keyring, so emitting it would
    # raise TypeError on the in-cluster path. init_from_config threads
    # disable_keyring through to init separately, alongside this helper.
    return kw


async def _initialize_client(
    api_key: str | None = None,
    auth_type: AuthType = "Pkce",
    endpoint: str | None = None,
    client_config: ClientConfig | None = None,
    headless: bool = False,
    insecure: bool = False,
    insecure_skip_verify: bool = False,
    ca_cert_file_path: str | None = None,
    command: List[str] | None = None,
    proxy_command: List[str] | None = None,
    client_id: str | None = None,
    client_credentials_secret: str | None = None,
    rpc_retries: int = 3,
    http_proxy_url: str | None = None,
    disable_keyring: bool = False,
) -> ClientSet:
    """
    Initialize the client based on the execution mode.
    Returns:
        The initialized client
    """
    from flyte.remote._client.controlplane import ClientSet

    if endpoint and api_key is None:
        return await ClientSet.for_endpoint(
            endpoint,
            insecure=insecure,
            insecure_skip_verify=insecure_skip_verify,
            auth_type=auth_type,
            headless=headless,
            ca_cert_file_path=ca_cert_file_path,
            command=command,
            proxy_command=proxy_command,
            client_id=client_id,
            client_credentials_secret=client_credentials_secret,
            client_config=client_config,
            rpc_retries=rpc_retries,
            http_proxy_url=http_proxy_url,
            disable_keyring=disable_keyring,
        )
    elif api_key:
        return await ClientSet.for_api_key(
            api_key,
            insecure=insecure,
            insecure_skip_verify=insecure_skip_verify,
            auth_type=auth_type,
            headless=headless,
            ca_cert_file_path=ca_cert_file_path,
            command=command,
            proxy_command=proxy_command,
            client_id=client_id,
            client_credentials_secret=client_credentials_secret,
            client_config=client_config,
            rpc_retries=rpc_retries,
            http_proxy_url=http_proxy_url,
            disable_keyring=disable_keyring,
        )

    raise InitializationError(
        "MissingEndpointOrApiKeyError", "user", "Either endpoint or api_key must be provided to initialize the client."
    )


def _initialize_logger(
    log_level: int | None = None,
    log_format: LogFormat | None = None,
    reset_root_logger: bool = False,
    user_log_level: int | None = None,
) -> None:
    # In-cluster runtimes never render Rich output (stdout is captured), so skip the Rich handler
    # — this avoids rich.logging and the transitive ipython_check -> IPython import at startup.
    enable_rich = os.environ.get("FLYTE_INTERNAL_EXECUTION_PROJECT") is None
    initialize_logger(
        log_level=log_level,
        log_format=log_format,
        enable_rich=enable_rich,
        reset_root_logger=reset_root_logger,
        user_log_level=user_log_level,
    )


@syncify
async def init(
    org: str | None = None,
    project: str | None = None,
    domain: str | None = None,
    root_dir: Path | None = None,
    log_level: int | None = None,
    log_format: LogFormat | None = None,
    reset_root_logger: bool = False,
    user_log_level: int | None = None,
    endpoint: str | None = None,
    headless: bool = False,
    insecure: bool = False,
    insecure_skip_verify: bool = False,
    ca_cert_file_path: str | None = None,
    auth_type: AuthType = "Pkce",
    command: List[str] | None = None,
    proxy_command: List[str] | None = None,
    api_key: str | None = None,
    client_id: str | None = None,
    client_credentials_secret: str | None = None,
    auth_client_config: ClientConfig | None = None,
    rpc_retries: int = 3,
    http_proxy_url: str | None = None,
    disable_keyring: bool = False,
    storage: Storage | None = None,
    batch_size: int = 1000,
    image_builder: ImageBuildEngine.ImageBuilderType = "local",
    images: typing.Dict[str, str] | None = None,
    image_registry: str | None = None,
    source_config_path: Optional[Path] = None,
    sync_local_sys_paths: bool = True,
    load_plugin_type_transformers: bool = True,
    local_persistence: bool = False,
    local_tracked: bool = False,
    local_tracked_strict: bool = False,
) -> None:
    """
    Initialize the Flyte system with the given configuration. This method should be called before any other Flyte
    remote API methods are called. Thread-safe implementation.

    Args:
        project: Optional project name (not used in this implementation)
        domain: Optional domain name (not used in this implementation)
        root_dir: Optional root directory from which to determine how to load files, and find paths to files.
            This is useful for determining the root directory for the current project, and for locating
            files like config etc.
            also use to determine all the code that needs to be copied to the remote location.
            defaults to the editable install directory if the cwd is in a Python editable install, else just the cwd.
        log_level: Optional logging level for the logger, default is set using the default initialization policies
        log_format: Optional logging format for the logger, default is "console"
        reset_root_logger: By default, we clear out root logger handlers and set up our own.
        api_key: Optional API key for authentication
        endpoint: Optional API endpoint URL
        headless: Optional Whether to run in headless mode
        insecure_skip_verify: Whether to skip SSL certificate verification
        auth_client_config: Optional client configuration for authentication
        auth_type: The authentication type to use (Pkce, ClientSecret, ExternalCommand, DeviceFlow)
        command: This command is executed to return a token using an external process
        proxy_command: This command is executed to return a token for proxy authorization using an external process
        client_id: This is the public identifier for the app which handles authorization for a Flyte deployment.
            More details here: https://www.oauth.com/oauth2-servers/client-registration/client-id-secret/.
        client_credentials_secret: Used for service auth, which is automatically called during pyflyte. This will
            allow the Flyte engine to read the password directly from the environment variable. Note that this is
            less secure! Please only use this if mounting the secret as a file is impossible
        ca_cert_file_path: [optional] str Root Cert to be loaded and used to verify admin
        http_proxy_url: [optional] HTTP Proxy to be used for OAuth requests
        rpc_retries: [optional] int Number of times to retry the platform calls
        insecure: insecure flag for the client
        storage: Optional blob store (S3, GCS, Azure) configuration if needed to access (i.e. using Minio)
        org: Optional organization override for the client. Should be set by auth instead.
        batch_size: Optional batch size for operations that use listings, defaults to 1000, so limit larger than
            batch_size will be split into multiple requests.
        image_builder: Optional image builder configuration, if not provided, the default image builder will be used.
        images: Optional dict of images that can be used by referencing the image name.
        image_registry: Optional container registry to push built images to, overriding the
            built-in default base registry. Equivalent to the `image.registry` config entry.
        source_config_path: Optional path to the source configuration file (This is only used for documentation)
        sync_local_sys_paths: Whether to include and synchronize local sys.path entries under the root directory
            into the remote container (default: True).
        load_plugin_type_transformers: If enabled (default True), load the type transformer plugins registered under
            the "flyte.plugins.types" entry point group.
        local_persistence: Whether to enable SQLite persistence for local run metadata (default: False).
        local_tracked: Whether to report tracked run state to the Flyte control plane
            (default: False). Requires an initialized client and a configured project/domain.
        local_tracked_strict: Strict tracked-run reporting for debugging (default: False). Any
            reporting failure fails the run loudly instead of being logged and swallowed. Only takes
            effect when reporting is enabled.
        disable_keyring: Disable storage of tokens in local keyring.

    Returns:
        None
    """
    from flyte._utils import org_from_endpoint, sanitize_endpoint
    from flyte.types import _load_custom_type_transformers

    _initialize_logger(
        log_level=log_level,
        log_format=log_format,
        reset_root_logger=reset_root_logger,
        user_log_level=user_log_level,
    )
    if load_plugin_type_transformers:
        _load_custom_type_transformers()

    global _init_config  # noqa: PLW0603

    endpoint = sanitize_endpoint(endpoint)

    with _init_lock:
        client = None
        if endpoint or api_key:
            client = await _initialize_client(
                api_key=api_key,
                auth_type=auth_type,
                endpoint=endpoint,
                headless=headless,
                insecure=insecure,
                insecure_skip_verify=insecure_skip_verify,
                ca_cert_file_path=ca_cert_file_path,
                command=command,
                proxy_command=proxy_command,
                client_id=client_id,
                client_credentials_secret=client_credentials_secret,
                client_config=auth_client_config,
                rpc_retries=rpc_retries,
                http_proxy_url=http_proxy_url,
                disable_keyring=disable_keyring,
            )

        if not root_dir:
            root_dir = Path.cwd()
        # We will inject the root_dir into the sys,path for module resolution
        sys.path.append(str(root_dir))

        _init_config = _InitConfig(
            root_dir=root_dir,
            project=project,
            domain=domain,
            client=client,
            storage=storage,
            org=org or org_from_endpoint(endpoint),
            batch_size=batch_size,
            image_builder=image_builder,
            images=images or {},
            image_registry=image_registry,
            source_config_path=source_config_path,
            sync_local_sys_paths=sync_local_sys_paths,
            local_persistence=local_persistence,
            local_tracked=local_tracked,
            local_tracked_strict=local_tracked_strict,
        )

        logger.info(f"Flyte initialized with config: {_init_config}")


@syncify
async def init_from_config(
    path_or_config: str | Path | Config | None = None,
    root_dir: Path | None = None,
    log_level: int | None = None,
    log_format: LogFormat = "console",
    user_log_level: int | None = None,
    org: str | None = None,
    project: str | None = None,
    domain: str | None = None,
    storage: Storage | None = None,
    batch_size: int = 1000,
    image_builder: ImageBuildEngine.ImageBuilderType | None = None,
    images: tuple[str, ...] | None = None,
    sync_local_sys_paths: bool = True,
) -> None:
    """
    Initialize the Flyte system using a configuration file or Config object. This method should be called before any
    other Flyte remote API methods are called. Thread-safe implementation.

    Args:
        path_or_config: Path to the configuration file or Config object
        org: Org name, this will override the org in the configuration file when non-empty
        project: Project name, this will override any project names in the configuration file
        domain: Domain name, this will override any domain names in the configuration file
        root_dir: Optional root directory from which to determine how to load files, and find paths to
            files like config etc. For example if one uses the copy-style=="all", it is essential to determine the
            root directory for the current project. If not provided, it defaults to the editable install directory or
            if not available, the current working directory.
        log_level: Optional logging level for the framework logger,
            default is set using the default initialization policies
        log_format: Optional logging format for the logger, default is "console"
        storage: Optional blob store (S3, GCS, Azure) configuration if needed to access (i.e. using Minio)
        images: List of image strings in format "imagename=imageuri" or just "imageuri".
        sync_local_sys_paths: Whether to include and synchronize local sys.path entries under the root directory
            into the remote container (default: True).
        batch_size: Optional batch size for operations that use listings, defaults to 1000
        image_builder: Optional image builder configuration, if provided,
            will override any defaults set in the configuration.

    Returns:
        None
    """
    from rich.highlighter import ReprHighlighter

    import flyte.config as config
    from flyte.cli._common import parse_images

    cfg: config.Config
    cfg_path: Optional[Path] = None
    if path_or_config is None:
        # If no path is provided, use the default config file
        cfg = config.auto()
    elif isinstance(path_or_config, (str, Path)):
        if root_dir:
            cfg_path = root_dir.expanduser() / path_or_config
        else:
            cfg_path = Path(path_or_config).expanduser()
        if not Path(cfg_path).exists():
            raise InitializationError(
                "ConfigFileNotFoundError",
                "user",
                f"Configuration file '{cfg_path}' does not exist., current working directory is {Path.cwd()}",
            )
        cfg = config.auto(cfg_path)
    else:
        cfg = path_or_config

    logger.info(f"Flyte config initialized as {cfg}", extra={"highlighter": ReprHighlighter()})

    # parse image, this will overwrite the image_refs set in the config file
    parse_images(cfg, images)

    await init.aio(
        org=org or cfg.task.org,
        project=project or cfg.task.project,
        domain=domain or cfg.task.domain,
        root_dir=root_dir,
        log_level=log_level,
        log_format=log_format,
        user_log_level=user_log_level,
        image_builder=image_builder or cfg.image.builder or "local",
        batch_size=batch_size,
        images=cfg.image.image_refs,
        image_registry=cfg.image.registry,
        storage=storage,
        source_config_path=cfg_path,
        sync_local_sys_paths=sync_local_sys_paths,
        local_persistence=cfg.local.persistence,
        local_tracked=cfg.local.tracked,
        local_tracked_strict=cfg.local.tracked_strict,
        # disable_keyring is threaded outside _platform_to_client_kwargs
        # because the helper output is also spread into
        # create_remote_controller from init_in_cluster, and the
        # controller's constructor doesn't accept this kwarg.
        disable_keyring=cfg.platform.disable_keyring,
        **_platform_to_client_kwargs(cfg.platform),
    )


@syncify
async def init_from_api_key(
    api_key: str | None = None,
    project: str | None = None,
    domain: str | None = None,
    root_dir: Path | None = None,
    log_level: int | None = None,
    log_format: LogFormat | None = None,
    storage: Storage | None = None,
    batch_size: int = 1000,
    image_builder: ImageBuildEngine.ImageBuilderType = "local",
    images: typing.Dict[str, str] | None = None,
    sync_local_sys_paths: bool = True,
    insecure_skip_verify: bool = False,
    ca_cert_file_path: str | None = None,
) -> None:
    """
    Initialize the Flyte system using an API key for authentication. This is a convenience
    method for API key-based authentication. Thread-safe implementation.

    The API key should be an encoded API key that contains the endpoint, client ID, client secret,
    and organization information. You can obtain this encoded API key from your Flyte administrator
    or cloud provider.

    Args:
        api_key: Optional encoded API key for authentication. If None, reads from FLYTE_API_KEY
            environment variable. The API key is a base64-encoded string containing endpoint, client_id,
            client_secret, and org information.
        project: Optional project name
        domain: Optional domain name
        root_dir: Optional root directory from which to determine how to load files, and find paths to files.
            defaults to the editable install directory if the cwd is in a Python editable install, else just the cwd.
        log_level: Optional logging level for the logger
        log_format: Optional logging format for the logger, default is "console"
        storage: Optional blob store (S3, GCS, Azure) configuration
        batch_size: Optional batch size for operations that use listings, defaults to 1000
        image_builder: Optional image builder configuration
        images: Optional dict of images that can be used by referencing the image name
        sync_local_sys_paths: Whether to include and synchronize local sys.path entries under the root directory
            into the remote container (default: True)
        insecure_skip_verify: Whether to skip SSL certificate verification
        ca_cert_file_path: Optional path to a CA certificate bundle used to verify the server certificate.
            Useful behind TLS-intercepting corporate proxies that re-sign traffic with a private CA.
            Note that if insecure_skip_verify is also set, it takes precedence and certificate
            verification against this bundle is skipped.

    Returns:
        None
    """
    from flyte._utils import sanitize_endpoint
    from flyte.remote._client.auth._auth_utils import decode_api_key

    # If api_key is not provided, read from environment variable
    if api_key is None:
        api_key = os.getenv("FLYTE_API_KEY")
        if api_key is None:
            raise InitializationError(
                "MissingApiKeyError",
                "user",
                "API key must be provided either as a parameter or via the FLYTE_API_KEY environment variable.",
            )

    # Decode the API key to extract endpoint, client_id, client_secret, and org
    endpoint, client_id, client_secret, org = decode_api_key(api_key)

    # Sanitize the endpoint
    endpoint = sanitize_endpoint(endpoint)  # type: ignore[assignment]

    await init.aio(
        org=None if org == "None" else org,
        project=project,
        domain=domain,
        endpoint=endpoint,
        api_key=api_key,
        client_id=client_id,
        client_credentials_secret=client_secret,
        auth_type="ClientSecret",  # API keys use client credentials flow
        root_dir=root_dir,
        log_level=log_level,
        log_format=log_format,
        insecure=False,
        insecure_skip_verify=insecure_skip_verify,
        ca_cert_file_path=ca_cert_file_path,
        storage=storage,
        batch_size=batch_size,
        image_builder=image_builder,
        images=images,
        sync_local_sys_paths=sync_local_sys_paths,
    )


@syncify
async def init_in_cluster(
    org: str | None = None,
    project: str | None = None,
    domain: str | None = None,
    api_key: str | None = None,
    endpoint: str | None = None,
    insecure: bool = False,
) -> dict[str, typing.Any]:
    from flyte._utils import str2bool

    PROJECT_NAME = "FLYTE_INTERNAL_EXECUTION_PROJECT"
    DOMAIN_NAME = "FLYTE_INTERNAL_EXECUTION_DOMAIN"
    ORG_NAME = "_U_ORG_NAME"
    ENDPOINT_OVERRIDE = "_U_EP_OVERRIDE"
    INSECURE_SKIP_VERIFY_OVERRIDE = "_U_INSECURE_SKIP_VERIFY"
    INSECURE_OVERRIDE = "_U_INSECURE"
    _UNION_EAGER_API_KEY_ENV_VAR = "_UNION_EAGER_API_KEY"
    EAGER_API_KEY = "EAGER_API_KEY"

    # When the cluster mounts a credentials config file (typical for the
    # file-mounted client-secret deploy path), prefer it over the legacy
    # env-var-injected api key. The chart sets FLYTECTL_CONFIG on the
    # task pod pointing at the mount, so a direct env-var probe is all
    # we need — no point walking resolve_config_path's full precedence
    # chain (which includes a `git rev-parse` subprocess that is wasted
    # work in a task pod and shows up in subprocess-mock tests).
    # If api_key is supplied by the caller explicitly, that wins.
    if api_key is None:
        from flyte.config._reader import FLYTECTL_CONFIG_ENV_VAR, UCTL_CONFIG_ENV_VAR

        cfg_path_str = os.getenv(UCTL_CONFIG_ENV_VAR) or os.getenv(FLYTECTL_CONFIG_ENV_VAR)
        if cfg_path_str:
            # Existence is intentionally NOT pre-checked: a typo in the
            # env var should surface as the FileNotFoundError that
            # init_from_config raises when it tries to open the path,
            # not silently fall back to the legacy api-key branch.
            logger.info(f"init_in_cluster: delegating to init_from_config({cfg_path_str})")
            # init_from_config is @syncify-decorated; call its async form so we don't
            # block the syncify thread.
            await init_from_config.aio(
                path_or_config=cfg_path_str,
                org=org or os.getenv(ORG_NAME),
                project=project or os.getenv(PROJECT_NAME),
                domain=domain or os.getenv(DOMAIN_NAME),
            )
            # runtime._run_action spreads our return as **kwargs into
            # create_remote_controller. Build that dict from the same
            # PlatformConfig mapping init_from_config used above so the
            # two clients can't drift apart on a future field addition.
            from flyte.config import Config

            cfg = Config.auto(cfg_path_str)
            kwargs = _platform_to_client_kwargs(cfg.platform)
            kwargs["headless"] = True  # task pod never has a browser available
            if ep := os.getenv(ENDPOINT_OVERRIDE):
                kwargs["endpoint"] = ep
            return kwargs

    org = org or os.getenv(ORG_NAME)
    project = project or os.getenv(PROJECT_NAME)
    domain = domain or os.getenv(DOMAIN_NAME)
    api_key = api_key or os.getenv(_UNION_EAGER_API_KEY_ENV_VAR) or os.getenv(EAGER_API_KEY)

    remote_kwargs: dict[str, typing.Any] = {"insecure": insecure}
    if api_key:
        logger.info("Using api key from environment")
        remote_kwargs["api_key"] = api_key
    else:
        ep = endpoint or os.environ.get(ENDPOINT_OVERRIDE, "host.docker.internal:8090")
        remote_kwargs["endpoint"] = ep
        if not insecure:
            if "localhost" in ep or "docker" in ep:
                remote_kwargs["insecure"] = True
        if str2bool(os.getenv(INSECURE_OVERRIDE, "")):
            remote_kwargs["insecure"] = True
        logger.debug(f"Using controller endpoint: {ep} with kwargs: {remote_kwargs}")

    # Check for insecure_skip_verify override (e.g. for self-signed certs)
    insecure_skip_verify_str = os.getenv(INSECURE_SKIP_VERIFY_OVERRIDE, "")
    if str2bool(insecure_skip_verify_str):
        remote_kwargs["insecure_skip_verify"] = True
        logger.info("SSL certificate verification disabled (insecure_skip_verify=True)")

    # Cluster runtime never benefits from keyring storage: tokens are short-lived, the pod is
    # ephemeral, and there is no human keychain to read from. Disabling keyring also avoids the
    # ~180ms cold-start hit from `keyring`'s backend enumeration (incl. the `keyring.backends.macOS.api`
    # C-extension probe that runs even on Linux). Passed directly to ``init`` rather than via
    # ``remote_kwargs`` because the returned dict is also used to construct the controller, which
    # does not accept ``disable_keyring``.
    await init.aio(
        org=org,
        project=project,
        domain=domain,
        root_dir=Path.cwd(),
        image_builder="remote",
        disable_keyring=True,
        **remote_kwargs,
    )
    return remote_kwargs


@syncify
async def init_passthrough(
    endpoint: str | None = None,
    org: str | None = None,
    project: str | None = None,
    domain: str | None = None,
    insecure: bool = False,
) -> dict[str, typing.Any]:
    """
    Initialize the Flyte system with passthrough authentication.

    This authentication mode allows you to pass custom authentication metadata
    using the `flyte.remote.auth_metadata()` context manager.

    The endpoint is automatically configured from the environment if in a flyte cluster with endpoint injected.

    Args:
        org: Optional organization name
        project: Optional project name
        domain: Optional domain name
        endpoint: Optional API endpoint URL
        insecure: Whether to use an insecure channel

    Returns:
        Dictionary of remote kwargs used for initialization
    """
    ENDPOINT_OVERRIDE = "_U_EP_OVERRIDE"
    ep = endpoint or os.environ.get(ENDPOINT_OVERRIDE, None)

    await init.aio(
        org=org,
        project=project,
        domain=domain,
        root_dir=Path.cwd(),
        image_builder="remote",
        endpoint=ep,
        insecure=insecure,
        auth_type="Passthrough",
    )
    return {"endpoint": endpoint, "insecure": insecure}


@contextmanager
def init_config_context(cfg: _InitConfig) -> Generator[None, None, None]:
    """
    Override the initialization configuration for the current context (thread / asyncio task).

    Internal API. Intended for servers that vary the config per request or per concurrent call
    (for example the MCP server scoping a tool call to a given project/domain): the override is
    scoped to the `with` block and to the current context, so concurrent callers never see
    each other's config. The module-global config is untouched.

    Args:
        cfg: The configuration to use for the duration of the block
    """
    token = _context_init_config.set(cfg)
    try:
        yield
    finally:
        _context_init_config.reset(token)


def _get_init_config() -> Optional[_InitConfig]:
    """
    Get the current initialization configuration. Thread-safe implementation.

    A context-scoped override installed by `init_config_context` wins over the
    module-global config; otherwise the global is returned.

    Returns:
        The current InitData if initialized, None otherwise
    """
    ctx_cfg = _context_init_config.get()
    if ctx_cfg is not None:
        return ctx_cfg
    with _init_lock:
        return _init_config


def get_init_config() -> _InitConfig:
    """
    Get the current initialization configuration. Thread-safe implementation.

    Returns:
        The current InitData if initialized, None otherwise
    """
    cfg = _get_init_config()
    if cfg is None:
        raise InitializationError(
            "ClientNotInitializedError",
            "user",
            "Configuration has not been initialized. Call flyte.init() with a valid endpoint/api-key before",
            " using this function or Call flyte.init_from_config() with a valid path to the config file",
        )
    return cfg


def get_storage() -> Storage | None:
    """
    Get the current storage configuration. Thread-safe implementation.

    Returns:
        The current storage configuration
    """
    cfg = _get_init_config()
    if cfg is None:
        raise InitializationError(
            "StorageNotInitializedError",
            "user",
            "Configuration has not been initialized. Call flyte.init() with a valid"
            " storage configuration before using this function or Call flyte.init_from_config()"
            " with a valid path to the config file",
        )
    return cfg.storage


def get_client() -> ClientSet:
    """
    Get the current client. Thread-safe implementation.

    Returns:
        The current client
    """
    cfg = _get_init_config()
    if cfg is None or cfg.client is None:
        raise InitializationError(
            "ClientNotInitializedError",
            "user",
            "Client has not been initialized. Call flyte.init() with a valid endpoint/api-key "
            "before using this function or Call flyte.init_from_config() with a valid path to the config file",
        )
    return cfg.client


def is_initialized() -> bool:
    """
    Check if the system has been initialized.

    Returns:
        True if initialized, False otherwise
    """
    return _get_init_config() is not None


def is_persistence_enabled() -> bool:
    """Check if local run persistence is enabled."""
    cfg = _get_init_config()
    if cfg is None:
        return False
    return cfg.local_persistence


def is_local_tracked_enabled() -> bool:
    """Check if reporting tracked run state to the control plane is enabled."""
    cfg = _get_init_config()
    if cfg is None:
        return False
    return cfg.local_tracked


def is_local_tracked_strict() -> bool:
    """Check if strict tracked-run reporting (fail the run on any reporting failure) is enabled."""
    cfg = _get_init_config()
    if cfg is None:
        return False
    return cfg.local_tracked_strict


def initialize_in_cluster() -> None:
    """
    Initialize the system for in-cluster execution. This is a placeholder function and does not perform any actions.

    Returns:
        None
    """
    init()


# Define a generic type variable for the decorated function
T = TypeVar("T", bound=Callable)


def ensure_client():
    """
    Ensure that the client is initialized. If not, raise an InitializationError.
    This function is used to check if the client is initialized before executing any Flyte remote API methods.
    """
    cfg = _get_init_config()
    if cfg is None or cfg.client is None:
        raise InitializationError(
            "ClientNotInitializedError",
            "user",
            "Client has not been initialized. Call flyte.init() with a valid endpoint/api-key before using"
            " this function or Call flyte.init_from_config() with a valid path to the config file",
        )


def requires_storage(func: T) -> T:
    """
    Decorator that checks if the storage has been initialized before executing the function.
    Raises InitializationError if the storage is not initialized.

    Args:
        func: Function to decorate

    Returns:
        Decorated function that checks for initialization
    """

    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        cfg = _get_init_config()
        if cfg is None or cfg.storage is None:
            raise InitializationError(
                "StorageNotInitializedError",
                "user",
                f"Function '{typing.cast('FunctionType', func).__name__}' requires storage to be initialized. "
                "Call flyte.init() with a valid storage configuration before using this function."
                "or Call flyte.init_from_config() with a valid path to the config file",
            )
        return func(*args, **kwargs)

    return typing.cast(T, wrapper)


def requires_upload_location(func: T) -> T:
    """
    Decorator that checks if the storage has been initialized before executing the function.
    Raises InitializationError if the storage is not initialized.

    Args:
        func: Function to decorate

    Returns:
        Decorated function that checks for initialization
    """

    @functools.wraps(func)
    def wrapper(*args, **kwargs) -> T:
        from ._context import internal_ctx

        ctx = internal_ctx()
        if not ctx.raw_data:
            raise InitializationError(
                "No upload path configured",
                "user",
                f"Function '{typing.cast('FunctionType', func).__name__}' requires client to be initialized. "
                "Call flyte.init() with storage configuration before using this function."
                "or Call flyte.init_from_config() with a valid path to the config file.",
            )
        return func(*args, **kwargs)

    return typing.cast(T, wrapper)


def requires_initialization(func: T) -> T:
    """
    Decorator that checks if the system has been initialized before executing the function.
    Raises InitializationError if the system is not initialized.

    Args:
        func: Function to decorate

    Returns:
        Decorated function that checks for initialization
    """

    @functools.wraps(func)
    def wrapper(*args, **kwargs) -> T:
        if not is_initialized():
            raise InitializationError(
                "NotInitConfiguredError",
                "user",
                f"Function '{typing.cast('FunctionType', func).__name__}' requires initialization. "
                "Call flyte.init() before using this function"
                " or Call flyte.init_from_config() with a valid path to the config file.",
            )
        return func(*args, **kwargs)

    return typing.cast(T, wrapper)


def require_project_and_domain(func):
    """
    Decorator that ensures the current Flyte configuration defines
    both 'project' and 'domain'. Raises a clear error if not found.
    """

    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        cfg = get_init_config()
        if cfg.project is None:
            raise InitializationError(
                "ProjectNotConfigured",
                "user",
                "Project must be provided to initialize the client. "
                "Please set 'project' in the 'task' section of your config file, "
                "or pass it directly to flyte.init(project='your-project-name').",
            )

        if cfg.domain is None:
            raise InitializationError(
                "DomainNotConfigured",
                "user",
                "Domain must be provided to initialize the client. "
                "Please set 'domain' in the 'task' section of your config file, "
                "or pass it directly to flyte.init(domain='your-domain-name').",
            )

        return func(*args, **kwargs)

    return wrapper


async def _init_for_testing(
    project: str | None = None,
    domain: str | None = None,
    root_dir: Path | None = None,
    log_level: int | None = None,
    client: ClientSet | None = None,
    org: str | None = None,
):
    global _init_config  # noqa: PLW0603

    if log_level:
        initialize_logger(log_level=log_level)

    with _init_lock:
        root_dir = root_dir or Path.cwd()
        _init_config = _InitConfig(
            root_dir=root_dir,
            project=project,
            domain=domain,
            client=client,
            org=org,
        )


def replace_client(client):
    global _init_config  # noqa: PLW0603

    with _init_lock:
        _init_config = typing.cast(_InitConfig, _init_config).replace(client=client)


def current_domain() -> str:
    """
    Returns the current domain from Runtime environment (on the cluster) or from the initialized configuration.
    This is safe to be used during `deploy`, `run` and within `task` code.

    NOTE: This will not work if you deploy a task to a domain and then run it in another domain.

    Raises InitializationError if the configuration is not initialized or domain is not set.
    Returns:
        The current domain
    """
    from ._context import ctx

    tctx = ctx()
    if tctx:
        domain = tctx.action.domain
        if domain is not None:
            return domain

    cfg = _get_init_config()
    if cfg is None or cfg.domain is None:
        raise InitializationError(
            "DomainNotInitializedError",
            "user",
            "Domain has not been initialized. Call flyte.init() with a valid domain before using this function"
            " or Call flyte.init_from_config() with a valid path to the config file",
        )
    return cfg.domain


def current_project() -> str:
    """
    Returns the current project from the Runtime environment (on the cluster) or from the initialized configuration.
    This is safe to be used during `deploy`, `run` and within `task` code.

    NOTE: This will not work if you deploy a task to a project and then run it in another project.

    Raises InitializationError if the configuration is not initialized or project is not set.
    Returns:
        The current project
    """
    from ._context import ctx

    tctx = ctx()
    if tctx:
        project = tctx.action.project
        if project is not None:
            return project

    cfg = _get_init_config()
    if cfg is None or cfg.project is None:
        raise InitializationError(
            "ProjectNotInitializedError",
            "user",
            "Project has not been initialized. Call flyte.init() with a valid project before using this function"
            " or Call flyte.init_from_config() with a valid path to the config file",
        )
    return cfg.project