mlflow--mlflow/mlflow/server/handlers.py

# Define all the service endpoint handlers here.
import io
import json
import logging
import os
import pathlib
import posixpath
import re
import tempfile
import threading
import time
import unicodedata
import urllib
from functools import partial, wraps
from typing import Any, Callable

import requests
from cachetools import TTLCache
from flask import Request, Response, current_app, g, jsonify, request, send_file
from google.protobuf import descriptor
from google.protobuf.json_format import ParseError
from werkzeug.http import quote_header_value

import mlflow
from mlflow.client import MlflowClient
from mlflow.entities import (
    Assessment,
    DatasetInput,
    Expectation,
    ExperimentTag,
    FallbackConfig,
    FallbackStrategy,
    Feedback,
    FileInfo,
    GatewayEndpointModelConfig,
    GatewayEndpointTag,
    GatewayResourceType,
    InputTag,
    IssueSeverity,
    IssueStatus,
    Metric,
    Param,
    RunStatus,
    RunTag,
    ViewType,
    Workspace,
    WorkspaceDeletionMode,
)
from mlflow.entities import (
    RoutingStrategy as RoutingStrategyEntity,
)
from mlflow.entities.gateway_budget_policy import (
    BudgetAction,
    BudgetDuration,
    BudgetDurationUnit,
    BudgetTargetScope,
    BudgetUnit,
)
from mlflow.entities.logged_model import LoggedModel
from mlflow.entities.logged_model_input import LoggedModelInput
from mlflow.entities.logged_model_output import LoggedModelOutput
from mlflow.entities.logged_model_parameter import LoggedModelParameter
from mlflow.entities.logged_model_status import LoggedModelStatus
from mlflow.entities.logged_model_tag import LoggedModelTag
from mlflow.entities.model_registry import ModelVersionTag, RegisteredModelTag
from mlflow.entities.model_registry.prompt_version import IS_PROMPT_TAG_KEY
from mlflow.entities.multipart_upload import MultipartUploadPart
from mlflow.entities.trace_info import TraceInfo
from mlflow.entities.trace_info_v2 import TraceInfoV2
from mlflow.entities.trace_metrics import MetricAggregation, MetricViewType
from mlflow.entities.trace_status import TraceStatus
from mlflow.entities.webhook import WebhookAction, WebhookEntity, WebhookEvent, WebhookStatus
from mlflow.environment_variables import (
    MLFLOW_CREATE_MODEL_VERSION_SOURCE_VALIDATION_REGEX,
    MLFLOW_DEPLOYMENTS_TARGET,
    MLFLOW_ENABLE_WORKSPACES,
    MLFLOW_PRESIGNED_DOWNLOAD_URL_TTL_SECONDS,
)
from mlflow.exceptions import (
    MlflowException,
    MlflowNotImplementedException,
    MlflowTraceDataException,
    MlflowTracingException,
    _UnsupportedMultipartDownloadException,
    _UnsupportedMultipartUploadException,
    _UnsupportedPresignedUploadException,
)
from mlflow.gateway.budget import maybe_refresh_budget_policies
from mlflow.gateway.budget_tracker import _policy_applies, get_budget_tracker
from mlflow.gateway.utils import is_valid_endpoint_name
from mlflow.genai.label_schemas.label_schemas import LabelSchemaType, _input_from_proto
from mlflow.genai.review_queues import ReviewItemType, ReviewQueueType, ReviewStatus
from mlflow.genai.review_queues.validation import validate_item_ids_for_attach
from mlflow.genai.scorers.scorer_utils import DECORATOR_SCORER_REGISTRATION_NOT_SUPPORTED_ERROR
from mlflow.models import Model
from mlflow.prompt.constants import PROMPT_TEXT_TAG_KEY, PROMPT_TYPE_TAG_KEY
from mlflow.protos import databricks_pb2
from mlflow.protos.databricks_pb2 import (
    BAD_REQUEST,
    FEATURE_DISABLED,
    INTERNAL_ERROR,
    INVALID_PARAMETER_VALUE,
    INVALID_STATE,
    RESOURCE_DOES_NOT_EXIST,
)
from mlflow.protos.issues_pb2 import (
    CreateIssue,
    GetIssue,
    SearchIssues,
    UpdateIssue,
)
from mlflow.protos.jobs_pb2 import JobStatus
from mlflow.protos.label_schemas_pb2 import (
    CreateLabelSchema,
    DeleteLabelSchema,
    GetLabelSchema,
    GetLabelSchemaByName,
    ListLabelSchemas,
    UpdateLabelSchema,
)
from mlflow.protos.mlflow_artifacts_pb2 import (
    AbortMultipartUpload,
    CompleteMultipartUpload,
    CreateMultipartUpload,
    DeleteArtifact,
    DownloadArtifact,
    GetPresignedDownloadUrl,
    MlflowArtifactsService,
    UploadArtifact,
)
from mlflow.protos.mlflow_artifacts_pb2 import (
    ListArtifacts as ListArtifactsMlflowArtifacts,
)
from mlflow.protos.model_registry_pb2 import (
    CreateModelVersion,
    CreateRegisteredModel,
    DeleteModelVersion,
    DeleteModelVersionTag,
    DeleteRegisteredModel,
    DeleteRegisteredModelAlias,
    DeleteRegisteredModelTag,
    GetLatestVersions,
    GetModelVersion,
    GetModelVersionByAlias,
    GetModelVersionDownloadUri,
    GetRegisteredModel,
    ModelRegistryService,
    RenameRegisteredModel,
    SearchModelVersions,
    SearchRegisteredModels,
    SetModelVersionTag,
    SetRegisteredModelAlias,
    SetRegisteredModelTag,
    TransitionModelVersionStage,
    UpdateModelVersion,
    UpdateRegisteredModel,
)
from mlflow.protos.prompt_optimization_pb2 import (
    PromptOptimizationJob as PromptOptimizationJobProto,
)
from mlflow.protos.review_queues_pb2 import (
    REVIEW_ITEM_TYPE_UNSPECIFIED,
    REVIEW_STATUS_UNSPECIFIED,
    AddItemsToReviewQueue,
    CreateReviewQueue,
    DeleteReviewQueue,
    GetOrCreateUserQueue,
    GetReviewQueue,
    GetReviewQueueByName,
    ListReviewQueueItems,
    ListReviewQueues,
    RemoveItemsFromReviewQueue,
    SetReviewQueueItemStatus,
    UpdateReviewQueue,
)
from mlflow.protos.service_pb2 import (
    AddDatasetToExperiments,
    AddGuardrailToEndpoint,
    AttachModelToGatewayEndpoint,
    BatchGetTraceInfos,
    BatchGetTraces,
    CalculateTraceFilterCorrelation,
    CancelPromptOptimizationJob,
    CreateAssessment,
    CreateDataset,
    CreateExperiment,
    CreateGatewayBudgetPolicy,
    CreateGatewayEndpoint,
    CreateGatewayEndpointBinding,
    CreateGatewayGuardrail,
    CreateGatewayModelDefinition,
    CreateGatewaySecret,
    CreateLoggedModel,
    CreatePresignedUploadUrl,
    CreatePromptOptimizationJob,
    CreateRun,
    CreateWorkspace,
    DeleteAssessment,
    DeleteDataset,
    DeleteDatasetRecords,
    DeleteDatasetTag,
    DeleteExperiment,
    DeleteExperimentTag,
    DeleteGatewayBudgetPolicy,
    DeleteGatewayEndpoint,
    DeleteGatewayEndpointBinding,
    DeleteGatewayEndpointTag,
    DeleteGatewayGuardrail,
    DeleteGatewayModelDefinition,
    DeleteGatewaySecret,
    DeleteLoggedModel,
    DeleteLoggedModelTag,
    DeletePromptOptimizationJob,
    DeleteRun,
    DeleteScorer,
    DeleteTag,
    DeleteTraces,
    DeleteTracesV3,
    DeleteTraceTag,
    DeleteTraceTagV3,
    DeleteWorkspace,
    DetachModelFromGatewayEndpoint,
    EndTrace,
    FinalizeLoggedModel,
    GetAssessmentRequest,
    GetDataset,
    GetDatasetExperimentIds,
    GetDatasetRecords,
    GetExperiment,
    GetExperimentByName,
    GetGatewayBudgetPolicy,
    GetGatewayEndpoint,
    GetGatewayGuardrail,
    GetGatewayModelDefinition,
    GetGatewaySecretInfo,
    GetLoggedModel,
    GetMetricHistory,
    GetMetricHistoryBulkInterval,
    GetPromptOptimizationJob,
    GetRun,
    GetScorer,
    GetTrace,
    GetTraceInfo,
    GetTraceInfoV3,
    GetWorkspace,
    LinkPromptsToTrace,
    LinkTracesToRun,
    ListArtifacts,
    ListEndpointGuardrailConfigs,
    ListGatewayBudgetPolicies,
    ListGatewayBudgetWindows,
    ListGatewayEndpointBindings,
    ListGatewayEndpoints,
    ListGatewayGuardrails,
    ListGatewayModelDefinitions,
    ListGatewaySecretInfos,
    ListLoggedModelArtifacts,
    ListScorers,
    ListScorerVersions,
    ListWorkspaces,
    LogBatch,
    LogInputs,
    LogLoggedModelParamsRequest,
    LogMetric,
    LogModel,
    LogOutputs,
    LogParam,
    MlflowService,
    QueryTraceMetrics,
    RegisterScorer,
    RemoveDatasetFromExperiments,
    RemoveGuardrailFromEndpoint,
    RestoreExperiment,
    RestoreRun,
    SearchDatasets,
    SearchEvaluationDatasets,
    SearchExperiments,
    SearchLoggedModels,
    SearchPromptOptimizationJobs,
    SearchRuns,
    SearchTraces,
    SearchTracesV3,
    SetDatasetTags,
    SetExperimentTag,
    SetGatewayEndpointTag,
    SetLoggedModelTags,
    SetTag,
    SetTraceTag,
    SetTraceTagV3,
    StartTrace,
    StartTraceV3,
    UpdateAssessment,
    UpdateEndpointGuardrailConfig,
    UpdateExperiment,
    UpdateGatewayBudgetPolicy,
    UpdateGatewayEndpoint,
    UpdateGatewayModelDefinition,
    UpdateGatewaySecret,
    UpdateRun,
    UpdateWorkspace,
    UpsertDatasetRecords,
)
from mlflow.protos.service_pb2 import Trace as ProtoTrace
from mlflow.protos.webhooks_pb2 import (
    CreateWebhook,
    DeleteWebhook,
    GetWebhook,
    ListWebhooks,
    TestWebhook,
    UpdateWebhook,
    WebhookService,
)
from mlflow.server.validation import _validate_content_type
from mlflow.server.workspace_helpers import (
    _get_workspace_store,
)
from mlflow.store.artifact.artifact_repo import (
    ARTIFACT_STREAM_CHUNK_SIZE,
    MultipartDownloadMixin,
    MultipartUploadMixin,
    PresignedUploadMixin,
    StreamUploadMixin,
)
from mlflow.store.artifact.artifact_repository_registry import get_artifact_repository
from mlflow.store.db.db_types import DATABASE_ENGINES
from mlflow.store.jobs.abstract_store import AbstractJobStore
from mlflow.store.model_registry.abstract_store import AbstractStore as AbstractModelRegistryStore
from mlflow.store.model_registry.rest_store import RestStore as ModelRegistryRestStore
from mlflow.store.tracking import (
    MAX_RESULTS_QUERY_TRACE_METRICS,
    SEARCH_MAX_RESULTS_DEFAULT,
    SEARCH_MAX_RESULTS_THRESHOLD,
)
from mlflow.store.tracking.abstract_store import AbstractStore as AbstractTrackingStore
from mlflow.store.tracking.databricks_rest_store import DatabricksTracingRestStore
from mlflow.store.workspace.abstract_store import WorkspaceNameValidator
from mlflow.telemetry import get_telemetry_client
from mlflow.telemetry.installation_id import get_or_create_installation_id
from mlflow.telemetry.schemas import Record, Status
from mlflow.telemetry.utils import (
    FALLBACK_UI_CONFIG,
    fetch_ui_telemetry_config,
    is_telemetry_disabled,
)
from mlflow.tracing.constant import SpansLocation, TraceTagKey
from mlflow.tracing.trace_archival_config import get_trace_archival_server_config
from mlflow.tracing.utils.artifact_utils import (
    TRACE_DATA_FILE_NAME,
    get_archive_uri_for_trace,
    get_artifact_uri_for_trace,
)
from mlflow.tracking._model_registry import utils as registry_utils
from mlflow.tracking._model_registry.registry import ModelRegistryStoreRegistry
from mlflow.tracking._tracking_service import utils
from mlflow.tracking._tracking_service.registry import TrackingStoreRegistry
from mlflow.tracking.context.default_context import _get_user
from mlflow.tracking.registry import UnsupportedModelRegistryStoreURIException
from mlflow.utils import workspace_context
from mlflow.utils.crypto import CRYPTO_KEK_PASSPHRASE_ENV_VAR
from mlflow.utils.databricks_utils import get_databricks_host_creds
from mlflow.utils.file_utils import local_file_uri_to_path
from mlflow.utils.mime_type_utils import _guess_mime_type
from mlflow.utils.mlflow_tags import (
    MLFLOW_GENAI_EVALUATE_JOB_ID,
    MLFLOW_ISSUE_DETECTION_JOB_ID,
    MLFLOW_RUN_TYPE,
    MLFLOW_RUN_TYPE_GENAI_EVALUATE,
    MLFLOW_RUN_TYPE_ISSUE_DETECTION,
    MLFLOW_TRACE_ARCHIVAL_FAILURE,
    MLFLOW_TRACE_ARCHIVE_LOCATION,
    MLFLOW_TRACE_SPANS_LOCATION,
)
from mlflow.utils.promptlab_utils import _create_promptlab_run_impl
from mlflow.utils.proto_json_utils import message_to_json, parse_dict
from mlflow.utils.providers import (
    get_all_providers,
    get_models,
    get_provider_config_response,
)
from mlflow.utils.string_utils import is_string_type
from mlflow.utils.time import get_current_time_millis
from mlflow.utils.uri import is_local_uri, validate_path_is_safe, validate_query_string
from mlflow.utils.validation import (
    _validate_batch_log_api_req,
    _validate_experiment_artifact_location,
    _validate_experiment_artifact_location_length,
    _validate_trace_archival_location,
    _validate_trace_archival_retention_string,
    invalid_value,
    missing_value,
)
from mlflow.utils.workspace_utils import DEFAULT_WORKSPACE_NAME
from mlflow.webhooks.delivery import deliver_webhook, test_webhook
from mlflow.webhooks.types import (
    ModelVersionAliasCreatedPayload,
    ModelVersionAliasDeletedPayload,
    ModelVersionCreatedPayload,
    ModelVersionTagDeletedPayload,
    ModelVersionTagSetPayload,
    PromptAliasCreatedPayload,
    PromptAliasDeletedPayload,
    PromptCreatedPayload,
    PromptTagDeletedPayload,
    PromptTagSetPayload,
    PromptVersionCreatedPayload,
    PromptVersionTagDeletedPayload,
    PromptVersionTagSetPayload,
    RegisteredModelCreatedPayload,
)

_logger = logging.getLogger(__name__)
_tracking_store = None
_model_registry_store = None
_job_store = None
_artifact_repo = None
STATIC_PREFIX_ENV_VAR = "_MLFLOW_STATIC_PREFIX"
MAX_RUNS_GET_METRIC_HISTORY_BULK = 100
MAX_RESULTS_PER_RUN = 2500


class TrackingStoreRegistryWrapper(TrackingStoreRegistry):
    def __init__(self):
        super().__init__()
        self.register("", self._get_file_store)
        self.register("file", self._get_file_store)
        for scheme in DATABASE_ENGINES:
            self.register(scheme, self._get_sqlalchemy_store)
        # Add support for Databricks tracking store
        self.register("databricks", self._get_databricks_rest_store)
        self.register_entrypoints()

    @classmethod
    def _get_file_store(cls, store_uri, artifact_uri):
        from mlflow.store.tracking.file_store import FileStore

        return FileStore(store_uri, artifact_uri)

    @classmethod
    def _get_sqlalchemy_store(cls, store_uri, artifact_uri):
        from mlflow.server.constants import READ_REPLICA_BACKEND_STORE_URI_ENV_VAR
        from mlflow.store.tracking.sqlalchemy_store import SqlAlchemyStore
        from mlflow.store.tracking.sqlalchemy_workspace_store import (
            WorkspaceAwareSqlAlchemyStore,
        )

        read_db_uri = os.environ.get(READ_REPLICA_BACKEND_STORE_URI_ENV_VAR, None)
        store_cls = (
            WorkspaceAwareSqlAlchemyStore if MLFLOW_ENABLE_WORKSPACES.get() else SqlAlchemyStore
        )
        return store_cls(store_uri, artifact_uri, read_db_uri=read_db_uri)

    @classmethod
    def _get_databricks_rest_store(cls, store_uri, artifact_uri):
        return DatabricksTracingRestStore(partial(get_databricks_host_creds, store_uri))


class ModelRegistryStoreRegistryWrapper(ModelRegistryStoreRegistry):
    def __init__(self):
        super().__init__()
        self.register("", self._get_file_store)
        self.register("file", self._get_file_store)
        for scheme in DATABASE_ENGINES:
            self.register(scheme, self._get_sqlalchemy_store)
        # Add support for Databricks registries
        self.register("databricks", self._get_databricks_rest_store)
        self.register("databricks-uc", self._get_databricks_uc_rest_store)
        self.register_entrypoints()

    @classmethod
    def _get_file_store(cls, store_uri):
        from mlflow.store.model_registry.file_store import FileStore

        return FileStore(store_uri)

    @classmethod
    def _get_sqlalchemy_store(cls, store_uri):
        from mlflow.server.constants import READ_REPLICA_BACKEND_STORE_URI_ENV_VAR
        from mlflow.store.model_registry.sqlalchemy_store import SqlAlchemyStore
        from mlflow.store.model_registry.sqlalchemy_workspace_store import (
            WorkspaceAwareSqlAlchemyStore,
        )

        read_db_uri = os.environ.get(READ_REPLICA_BACKEND_STORE_URI_ENV_VAR, None)
        store_cls = (
            WorkspaceAwareSqlAlchemyStore if MLFLOW_ENABLE_WORKSPACES.get() else SqlAlchemyStore
        )
        return store_cls(store_uri, read_db_uri=read_db_uri)

    @classmethod
    def _get_databricks_rest_store(cls, store_uri):
        return ModelRegistryRestStore(partial(get_databricks_host_creds, store_uri))

    @classmethod
    def _get_databricks_uc_rest_store(cls, store_uri):
        from mlflow.environment_variables import MLFLOW_TRACKING_URI
        from mlflow.store._unity_catalog.registry.rest_store import UcModelRegistryStore

        # Get tracking URI from environment or use "databricks-uc" as default
        tracking_uri = MLFLOW_TRACKING_URI.get() or "databricks-uc"
        return UcModelRegistryStore(store_uri, tracking_uri)


_tracking_store_registry = TrackingStoreRegistryWrapper()
_model_registry_store_registry = ModelRegistryStoreRegistryWrapper()


def _get_artifact_repo_mlflow_artifacts():
    """
    Get an artifact repository specified by ``--artifacts-destination`` option for ``mlflow server``
    command.
    """
    from mlflow.server import ARTIFACTS_DESTINATION_ENV_VAR

    global _artifact_repo
    if _artifact_repo is None:
        _artifact_repo = get_artifact_repository(os.environ[ARTIFACTS_DESTINATION_ENV_VAR])
    return _artifact_repo


def _get_trace_artifact_repo(trace_info: TraceInfo):
    """
    Resolve the artifact repository for fetching data for the given trace.

    Args:
        trace_info: The trace info object containing metadata about the trace.
    """
    return _get_trace_repo_from_uri(get_artifact_uri_for_trace(trace_info))


def _get_trace_archive_repo(trace_info: TraceInfo):
    """
    Resolve the artifact repository that stores archived trace payloads.
    """
    return _get_trace_repo_from_uri(get_archive_uri_for_trace(trace_info))


def _get_trace_repo_from_uri(artifact_uri: str):
    if _is_servable_proxied_run_artifact_root(artifact_uri):
        # If the artifact location is a proxied run artifact root (e.g. mlflow-artifacts://...),
        # we need to resolve it to the actual artifact location.
        from mlflow.server import ARTIFACTS_DESTINATION_ENV_VAR

        path = _get_proxied_run_artifact_destination_path(artifact_uri)
        if not path:
            raise MlflowException(
                f"Failed to resolve the proxied run artifact URI: {artifact_uri}. ",
                "Trace artifact URI must contain subpath to the trace data directory.",
                error_code=BAD_REQUEST,
            )
        root = os.environ[ARTIFACTS_DESTINATION_ENV_VAR]
        artifact_uri = posixpath.join(root, path)

        # We don't set it to global var unlike run artifact, because the artifact repo has
        # to be created with full trace artifact URI including request_id.
        # e.g. s3://<experiment_id>/traces/<request_id>
        artifact_repo = get_artifact_repository(artifact_uri)
    else:
        artifact_repo = get_artifact_repository(artifact_uri)
    return artifact_repo


def _is_serving_proxied_artifacts():
    """
    Returns:
        True if the MLflow server is serving proxied artifacts (i.e. acting as a proxy for
        artifact upload / download / list operations), as would be enabled by specifying the
        --serve-artifacts configuration option. False otherwise.
    """
    from mlflow.server import SERVE_ARTIFACTS_ENV_VAR

    return os.environ.get(SERVE_ARTIFACTS_ENV_VAR, "false") == "true"


def _is_servable_proxied_run_artifact_root(run_artifact_root):
    """
    Determines whether or not the following are true:

    - The specified Run artifact root is a proxied artifact root (i.e. an artifact root with scheme
      ``http``, ``https``, or ``mlflow-artifacts``).

    - The MLflow server is capable of resolving and accessing the underlying storage location
      corresponding to the proxied artifact root, allowing it to fulfill artifact list and
      download requests by using this storage location directly.

    Args:
        run_artifact_root: The Run artifact root location (URI).

    Returns:
        True if the specified Run artifact root refers to proxied artifacts that can be
        served by this MLflow server (i.e. the server has access to the destination and
        can respond to list and download requests for the artifact). False otherwise.
    """
    parsed_run_artifact_root = urllib.parse.urlparse(run_artifact_root)
    # NB: If the run artifact root is a proxied artifact root (has scheme `http`, `https`, or
    # `mlflow-artifacts`) *and* the MLflow server is configured to serve artifacts, the MLflow
    # server always assumes that it has access to the underlying storage location for the proxied
    # artifacts. This may not always be accurate. For example:
    #
    # An organization may initially use the MLflow server to serve Tracking API requests and proxy
    # access to artifacts stored in Location A (via `mlflow server --serve-artifacts`). Then, for
    # scalability and / or security purposes, the organization may decide to store artifacts in a
    # new location B and set up a separate server (e.g. `mlflow server --artifacts-only`) to proxy
    # access to artifacts stored in Location B.
    #
    # In this scenario, requests for artifacts stored in Location B that are sent to the original
    # MLflow server will fail if the original MLflow server does not have access to Location B
    # because it will assume that it can serve all proxied artifacts regardless of the underlying
    # location. Such failures can be remediated by granting the original MLflow server access to
    # Location B.
    return (
        parsed_run_artifact_root.scheme in ["http", "https", "mlflow-artifacts"]
        and _is_serving_proxied_artifacts()
    )


def _get_proxied_run_artifact_destination_path(proxied_artifact_root, relative_path=None):
    """
    Resolves the specified proxied artifact location within a Run to a concrete storage location.

    Args:
        proxied_artifact_root: The Run artifact root location (URI) with scheme ``http``,
            ``https``, or `mlflow-artifacts` that can be resolved by the MLflow server to a
            concrete storage location.
        relative_path: The relative path of the destination within the specified
            ``proxied_artifact_root``. If ``None``, the destination is assumed to be
            the resolved ``proxied_artifact_root``.

    Returns:
        The storage location of the specified artifact.
    """
    parsed_proxied_artifact_root = urllib.parse.urlparse(proxied_artifact_root)
    assert parsed_proxied_artifact_root.scheme in ["http", "https", "mlflow-artifacts"]

    if parsed_proxied_artifact_root.scheme == "mlflow-artifacts":
        # If the proxied artifact root is an `mlflow-artifacts` URI, the run artifact root path is
        # simply the path component of the URI, since the fully-qualified format of an
        # `mlflow-artifacts` URI is `mlflow-artifacts://<netloc>/path/to/artifact`
        proxied_run_artifact_root_path = parsed_proxied_artifact_root.path.lstrip("/")
    else:
        # In this case, the proxied artifact root is an HTTP(S) URL referring to an mlflow-artifacts
        # API route that can be used to download the artifact. These routes are always anchored at
        # `/api/2.0/mlflow-artifacts/artifacts`. Accordingly, we split the path on this route anchor
        # and interpret the rest of the path (everything after the route anchor) as the run artifact
        # root path
        mlflow_artifacts_http_route_anchor = "/api/2.0/mlflow-artifacts/artifacts/"
        assert mlflow_artifacts_http_route_anchor in parsed_proxied_artifact_root.path

        proxied_run_artifact_root_path = parsed_proxied_artifact_root.path.split(
            mlflow_artifacts_http_route_anchor
        )[1].lstrip("/")

    return (
        posixpath.join(proxied_run_artifact_root_path, relative_path)
        if relative_path is not None
        else proxied_run_artifact_root_path
    )


def _get_tracking_store(
    backend_store_uri: str | None = None,
    default_artifact_root: str | None = None,
) -> AbstractTrackingStore:
    from mlflow.server import ARTIFACT_ROOT_ENV_VAR, BACKEND_STORE_URI_ENV_VAR

    global _tracking_store
    if _tracking_store is None:
        store_uri = backend_store_uri or os.environ.get(BACKEND_STORE_URI_ENV_VAR, None)
        artifact_root = default_artifact_root or os.environ.get(ARTIFACT_ROOT_ENV_VAR, None)
        _tracking_store = _tracking_store_registry.get_store(store_uri, artifact_root)
        utils.set_tracking_uri(store_uri)
    return _tracking_store


def _get_model_registry_store(registry_store_uri: str | None = None) -> AbstractModelRegistryStore:
    from mlflow.server import BACKEND_STORE_URI_ENV_VAR, REGISTRY_STORE_URI_ENV_VAR

    global _model_registry_store
    if _model_registry_store is None:
        store_uri = (
            registry_store_uri
            or os.environ.get(REGISTRY_STORE_URI_ENV_VAR, None)
            or os.environ.get(BACKEND_STORE_URI_ENV_VAR, None)
        )
        _model_registry_store = _model_registry_store_registry.get_store(store_uri)
        registry_utils.set_registry_uri(store_uri)
    return _model_registry_store


def _get_job_store(backend_store_uri: str | None = None) -> AbstractJobStore:
    """
    Get a job store instance based on the backend store URI.

    Args:
        backend_store_uri: Optional backend store URI. If not provided,
                          uses environment variable.

    Returns:
        An instance of AbstractJobStore
    """
    from mlflow.server import BACKEND_STORE_URI_ENV_VAR
    from mlflow.store.jobs.sqlalchemy_store import SqlAlchemyJobStore
    from mlflow.store.jobs.sqlalchemy_workspace_store import WorkspaceAwareSqlAlchemyJobStore
    from mlflow.utils.uri import extract_db_type_from_uri

    global _job_store
    if _job_store is None:
        store_uri = backend_store_uri or os.environ.get(BACKEND_STORE_URI_ENV_VAR, None)
        if not store_uri:
            raise MlflowException.invalid_parameter_value("Job store requires a backend store URI")
        try:
            extract_db_type_from_uri(store_uri)
        except (MlflowException, ValueError):
            # Require a database backend URI for the job store
            # Raise MlflowException so the CLI/REST layer returns a structured 400
            # instead of surfacing a generic 500 from ValueError
            raise MlflowException.invalid_parameter_value("Job store requires a backend store URI")

        store_cls = (
            WorkspaceAwareSqlAlchemyJobStore
            if MLFLOW_ENABLE_WORKSPACES.get()
            else SqlAlchemyJobStore
        )
        _job_store = store_cls(store_uri)

        if MLFLOW_ENABLE_WORKSPACES.get():
            _verify_job_store_workspace_support(_job_store)

    return _job_store


def initialize_backend_stores(
    backend_store_uri: str | None = None,
    registry_store_uri: str | None = None,
    default_artifact_root: str | None = None,
    workspace_store_uri: str | None = None,
    read_replica_backend_store_uri: str | None = None,
) -> None:
    from mlflow.server.constants import READ_REPLICA_BACKEND_STORE_URI_ENV_VAR

    # Set the read backend store URI env var so _get_sqlalchemy_store can pick it up
    if read_replica_backend_store_uri:
        os.environ[READ_REPLICA_BACKEND_STORE_URI_ENV_VAR] = read_replica_backend_store_uri

    tracking_store = _get_tracking_store(backend_store_uri, default_artifact_root)
    registry_store = None
    try:
        registry_store = _get_model_registry_store(registry_store_uri)
    except UnsupportedModelRegistryStoreURIException:
        pass

    if MLFLOW_ENABLE_WORKSPACES.get():
        # Initialize the workspace store to verify it's correctly configured
        _get_workspace_store(
            workspace_uri=workspace_store_uri,
            tracking_uri=backend_store_uri,
        )
        _verify_tracking_store_workspace_support(tracking_store)
        _verify_model_registry_store_workspace_support(registry_store)

    _verify_tracking_store_trace_archival_support(tracking_store)


def _store_supports_workspaces(
    store: AbstractTrackingStore | AbstractModelRegistryStore | AbstractJobStore,
) -> bool:
    """Return whether the provided store reports workspace support."""
    return bool(getattr(store, "supports_workspaces", False))


def _store_supports_trace_archival(store: AbstractTrackingStore) -> bool:
    """Return whether the provided tracking store reports trace archival support."""
    return bool(getattr(store, "supports_trace_archival", False))


def _verify_tracking_store_workspace_support(tracking_store: AbstractTrackingStore) -> None:
    if not _store_supports_workspaces(tracking_store):
        raise MlflowException(
            "The configured tracking store does not support workspace-aware operations. "
            "Remove the --enable-workspaces flag or configure a workspace-capable backend store.",
            error_code=INVALID_STATE,
        )


def _verify_tracking_store_trace_archival_support(tracking_store: AbstractTrackingStore) -> None:
    trace_archival_config = get_trace_archival_server_config()
    if trace_archival_config is None or not trace_archival_config.enabled:
        return

    if not _store_supports_trace_archival(tracking_store):
        raise MlflowException(
            "The configured tracking store does not support server-owned trace archival. "
            "Remove the trace archival config or configure a trace-archival-capable backend store.",
            error_code=INVALID_STATE,
        )


def _verify_model_registry_store_workspace_support(
    registry_store: AbstractModelRegistryStore,
) -> None:
    if registry_store is None:
        return

    if not _store_supports_workspaces(registry_store):
        raise MlflowException(
            "The configured model registry store does not support workspace-aware operations. "
            "Remove the --enable-workspaces flag or configure a workspace-capable backend store.",
            error_code=INVALID_STATE,
        )


def _verify_job_store_workspace_support(job_store: AbstractJobStore) -> None:
    if not _store_supports_workspaces(job_store):
        raise MlflowException(
            "The configured job store does not support workspace-aware operations. "
            "Remove the --enable-workspaces flag or configure a workspace-capable backend store.",
            error_code=INVALID_STATE,
        )


def _assert_string(x):
    assert isinstance(x, str)


def _assert_intlike(x):
    try:
        x = int(x)
    except ValueError:
        pass

    assert isinstance(x, int)


def _assert_bool(x):
    assert isinstance(x, bool)


def _assert_floatlike(x):
    try:
        x = float(x)
    except ValueError:
        pass

    assert isinstance(x, float)


def _assert_array(x):
    assert isinstance(x, list)


def _assert_dict(x):
    assert isinstance(x, dict)


def _assert_map_key_present(x):
    _assert_array(x)
    for entry in x:
        _assert_required(entry.get("key"))


def _assert_required(x, path=None):
    if path is None:
        assert x is not None
        # When parsing JSON payloads via proto, absent string fields
        # are expressed as empty strings
        assert x != ""
    else:
        assert x is not None, missing_value(path)
        assert x != "", missing_value(path)


def _assert_less_than_or_equal(x, max_value, message=None):
    if x > max_value:
        raise AssertionError(message) if message else AssertionError()


def _assert_intlike_within_range(x, min_value, max_value, message=None):
    if not min_value <= x <= max_value:
        raise AssertionError(message) if message else AssertionError()


def _assert_item_type_string(x):
    assert all(isinstance(item, str) for item in x)


def _assert_secret_value(x):
    """Validate secret_value is present. Does not print values in errors."""
    if x is None:
        raise MlflowException(
            message="Missing value for required parameter 'secret_value'.",
            error_code=INVALID_PARAMETER_VALUE,
        )


_TYPE_VALIDATORS = {
    _assert_intlike,
    _assert_string,
    _assert_bool,
    _assert_floatlike,
    _assert_array,
    _assert_item_type_string,
}


def _validate_param_against_schema(schema, param, value, proto_parsing_succeeded=False):
    """
    Attempts to validate a single parameter against a specified schema. Examples of the elements of
    the schema are type assertions and checks for required parameters. Returns None on validation
    success.  Otherwise, raises an MLFlowException if an assertion fails. This method is intended
    to be called for side effects.

    Args:
        schema: A list of functions to validate the parameter against.
        param: The string name of the parameter being validated.
        value: The corresponding value of the `param` being validated.
        proto_parsing_succeeded: A boolean value indicating whether proto parsing succeeded.
            If the proto was successfully parsed, we assume all of the types of the parameters in
            the request body were correctly specified, and thus we skip validating types. If proto
            parsing failed, then we validate types in addition to the rest of the schema. For
            details, see https://github.com/mlflow/mlflow/pull/5458#issuecomment-1080880870.
    """

    for f in schema:
        if f in _TYPE_VALIDATORS and proto_parsing_succeeded:
            continue

        try:
            f(value)
        except AssertionError as e:
            if e.args:
                message = e.args[0]
            elif f == _assert_required:
                message = f"Missing value for required parameter '{param}'."
            else:
                message = invalid_value(
                    param, value, f" Hint: Value was of type '{type(value).__name__}'."
                )
            raise MlflowException(
                message=(
                    message + " See the API docs for more information about request parameters."
                ),
                error_code=INVALID_PARAMETER_VALUE,
            )

    return None


def _get_request_json(flask_request=request):
    _validate_content_type(flask_request, ["application/json"])
    return flask_request.get_json(force=True, silent=True)


def _get_normalized_request_json(flask_request: Request = request) -> dict[str, Any]:
    """
    Get request JSON with normalization for legacy clients.

    Handles double-encoded JSON strings from older clients and empty request bodies.

    Args:
        flask_request: The Flask request object.

    Returns:
        The request data as a dictionary (empty dict if no body).
    """
    request_json = _get_request_json(flask_request)

    # Older clients may post their JSON double-encoded as strings, so the get_json
    # above actually converts it to a string. Therefore, we check this condition
    # (which we can tell for sure because any proper request should be a dictionary),
    # and decode it a second time.
    if is_string_type(request_json):
        request_json = json.loads(request_json)

    # If request doesn't have json body then assume it's empty.
    if request_json is None:
        request_json = {}

    return request_json


def _validate_request_json_with_schema(
    request_json: dict[str, Any],
    schema: dict[str, list[Callable[..., Any]]] | None,
    proto_parsing_succeeded: bool | None,
) -> None:
    """
    Validate request JSON against a schema without requiring protobuf messages.

    Args:
        request_json: The request data as a dictionary.
        schema: Dictionary mapping parameter names to lists of validation functions.
        proto_parsing_succeeded: Whether protobuf parsing succeeded. None indicates the
            request was not parsed from protobuf.
    """
    schema = schema or {}
    for schema_key, schema_validation_fns in schema.items():
        if schema_key in request_json or _assert_required in schema_validation_fns:
            value = request_json.get(schema_key)
            if schema_key == "run_id" and value is None and "run_uuid" in request_json:
                value = request_json.get("run_uuid")
            _validate_param_against_schema(
                schema=schema_validation_fns,
                param=schema_key,
                value=value,
                proto_parsing_succeeded=proto_parsing_succeeded,
            )


def _get_request_message(request_message, flask_request=request, schema=None):
    if flask_request.method == "GET" and flask_request.args:
        # Convert atomic values of repeated fields to lists before calling protobuf deserialization.
        # Context: We parse the parameter string into a dictionary outside of protobuf since
        # protobuf does not know how to read the query parameters directly. The query parser above
        # has no type information and hence any parameter that occurs exactly once is parsed as an
        # atomic value. Since protobuf requires that the values of repeated fields are lists,
        # deserialization will fail unless we do the fix below.
        request_json = {}
        for field in request_message.DESCRIPTOR.fields:
            if field.name not in flask_request.args:
                continue

            # Use is_repeated property (preferred) with fallback to deprecated label
            try:
                is_repeated = field.is_repeated
            except AttributeError:
                is_repeated = field.label == descriptor.FieldDescriptor.LABEL_REPEATED

            if is_repeated:
                request_json[field.name] = flask_request.args.getlist(field.name)
            else:
                value = flask_request.args.get(field.name)
                if field.type == descriptor.FieldDescriptor.TYPE_BOOL and isinstance(value, str):
                    if value.lower() not in ["true", "false"]:
                        raise MlflowException.invalid_parameter_value(
                            f"Invalid boolean value: {value}, must be 'true' or 'false'.",
                        )
                    value = value.lower() == "true"
                request_json[field.name] = value
    else:
        request_json = _get_normalized_request_json(flask_request)

    proto_parsing_succeeded = True
    try:
        parse_dict(request_json, request_message)
    except ParseError:
        proto_parsing_succeeded = False

    _validate_request_json_with_schema(request_json, schema, proto_parsing_succeeded)

    return request_message


def _get_validated_flask_request_json(
    flask_request: Request = request,
    schema: dict[str, list[Callable[..., Any]]] | None = None,
) -> dict[str, Any]:
    """
    Get and validate request data without protobuf parsing.

    This is an alternative to _get_request_message for endpoints that don't
    use protobuf message definitions. Supports both GET and POST/PUT requests.

    Args:
        flask_request: The Flask request object.
        schema: Dictionary mapping parameter names to lists of validation functions.

    Returns:
        The validated request data as a dictionary.

    Raises:
        MlflowException: If validation fails.
    """
    if flask_request.method == "GET" and flask_request.args:
        # Extract query parameters for GET requests
        request_json = {}
        schema = schema or {}
        for key in flask_request.args:
            # Get all values for this key (supports repeated parameters)
            values = flask_request.args.getlist(key)
            # Check if this field is a list type by looking for _assert_array validator
            is_list_type = _assert_array in schema.get(key, [])
            # If list type, always keep as list; otherwise use scalar if only one value
            request_json[key] = (
                values if is_list_type else (values[0] if len(values) == 1 else values)
            )
    else:
        request_json = _get_normalized_request_json(flask_request)

    _validate_request_json_with_schema(request_json, schema, proto_parsing_succeeded=None)

    return request_json


def _content_disposition_attachment(filename: str) -> str:
    """
    Build an RFC 6266 / RFC 5987 ``Content-Disposition`` value for an attachment.

    HTTP headers must be ASCII-encodable; ASGI adapters such as starlette's
    ``WSGIMiddleware`` raise ``UnicodeEncodeError`` on raw non-ASCII bytes. For
    filenames containing non-ASCII characters, emit an ASCII ``filename=``
    fallback alongside a percent-encoded ``filename*=UTF-8''…`` parameter.
    """
    try:
        filename.encode("ascii")
    except UnicodeEncodeError:
        # ``or "download"`` ensures a well-formed ``filename=<value>`` parameter
        # even when normalization strips every character (e.g. ``日本語`` with no
        # extension). Clients that ignore ``filename*`` still get a usable name.
        ascii_fallback = (
            unicodedata.normalize("NFKD", filename).encode("ascii", "ignore").decode("ascii")
            or "download"
        )
        # safe = RFC 5987 attr-char
        quoted = urllib.parse.quote(filename, safe="!#$&+-.^_`|~")
        quoted_ascii_fallback = quote_header_value(ascii_fallback, allow_token=True)
        return f"attachment; filename={quoted_ascii_fallback}; filename*=UTF-8''{quoted}"
    quoted_filename = quote_header_value(filename, allow_token=True)
    return f"attachment; filename={quoted_filename}"


def _response_with_file_attachment_headers(file_path, response):
    mime_type = _guess_mime_type(file_path)
    filename = pathlib.Path(file_path).name
    response.mimetype = mime_type
    content_disposition_header_name = "Content-Disposition"
    if content_disposition_header_name not in response.headers:
        response.headers[content_disposition_header_name] = _content_disposition_attachment(
            filename
        )
    response.headers["X-Content-Type-Options"] = "nosniff"
    response.headers["Content-Type"] = mime_type
    return response


def _create_artifact_file_response(file_path: str, artifact_name: str) -> Response:
    """Serve a local file while preserving the logical artifact name for downloads."""
    if os.path.isdir(file_path):
        raise MlflowException.invalid_parameter_value(
            f"Artifact path refers to a directory, not a file: '{artifact_name}'"
        )
    file_sender_response = send_file(file_path, mimetype=_guess_mime_type(file_path))
    file_sender_response.headers["Content-Disposition"] = _content_disposition_attachment(
        pathlib.Path(artifact_name).name
    )
    return _response_with_file_attachment_headers(file_path, file_sender_response)


def _send_artifact(artifact_repository, path):
    # Always send artifacts as attachments to prevent the browser from displaying them on our web
    # server's domain, which might enable XSS.
    if (local_path := artifact_repository.get_local_path(path)) is not None:
        return _create_artifact_file_response(os.path.abspath(local_path), path)

    tmp_dir = tempfile.TemporaryDirectory()
    try:
        file_path = os.path.abspath(
            artifact_repository.download_artifacts(path, dst_path=tmp_dir.name)
        )
        response = _create_artifact_file_response(file_path, path)
        response.call_on_close(tmp_dir.cleanup)
        return response
    except Exception:
        tmp_dir.cleanup()
        raise


def catch_mlflow_exception(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        try:
            return func(*args, **kwargs)
        except MlflowException as e:
            response = Response(mimetype="application/json")
            response.set_data(e.serialize_as_json())
            response.status_code = e.get_http_status_code()
            if response.status_code >= 500:
                is_debug = _logger.isEnabledFor(logging.DEBUG)
                msg = f"Error in {func.__name__}: {e}"
                if not is_debug:
                    msg += ". Set MLFLOW_LOGGING_LEVEL=DEBUG for traceback."
                _logger.error(msg, exc_info=is_debug)
            return response

    return wrapper


def _disable_unless_serve_artifacts(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        if not _is_serving_proxied_artifacts():
            return Response(
                (
                    f"Endpoint: {request.url_rule} disabled due to the mlflow server running "
                    "with `--no-serve-artifacts`. To enable artifacts server functionality, "
                    "run `mlflow server` with `--serve-artifacts`"
                ),
                503,
            )
        return func(*args, **kwargs)

    return wrapper


def _disable_if_artifacts_only(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        from mlflow.server import ARTIFACTS_ONLY_ENV_VAR

        if os.environ.get(ARTIFACTS_ONLY_ENV_VAR):
            return Response(
                (
                    f"Endpoint: {request.url_rule} disabled due to the mlflow server running "
                    "in `--artifacts-only` mode. To enable tracking server functionality, run "
                    "`mlflow server` without `--artifacts-only`"
                ),
                503,
            )
        return func(*args, **kwargs)

    return wrapper


def _disable_if_workspaces_disabled(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        if not MLFLOW_ENABLE_WORKSPACES.get():
            return Response(
                (
                    f"Endpoint: {request.url_rule} disabled because the server is running "
                    "without workspaces support. To enable workspace, run "
                    "`mlflow server` with `--enable-workspaces`"
                ),
                503,
            )
        return func(*args, **kwargs)

    return wrapper


def _workspace_not_supported(message: str) -> MlflowException:
    return MlflowException(message, FEATURE_DISABLED)


def _validate_storage_location_uri(value: str, field_name: str) -> str:
    """Validate a storage URI shared by experiment and workspace settings."""
    parsed = urllib.parse.urlparse(value)
    if parsed.fragment or parsed.params:
        raise MlflowException.invalid_parameter_value(
            f"'{field_name}' URL can't include fragments or params."
        )

    validate_query_string(parsed.query)
    _validate_experiment_artifact_location(value)
    _validate_experiment_artifact_location_length(value)
    return value


def _validate_optional_workspace_storage_location(value: str | None, field_name: str) -> str | None:
    if value is None:
        return None

    trimmed = value.strip()
    if not trimmed:
        return ""

    return _validate_storage_location_uri(trimmed, field_name)


def _validate_workspace_default_artifact_root(value: str | None) -> str | None:
    return _validate_optional_workspace_storage_location(value, "default_artifact_root")


def _validate_workspace_trace_archival_location(value: str | None) -> str | None:
    validated = _validate_optional_workspace_storage_location(
        value, "trace_archival_config.location"
    )
    if validated in (None, ""):
        return validated
    return _validate_trace_archival_location(
        validated, parameter_name="trace_archival_config.location"
    )


def _validate_workspace_trace_archival_retention(value: str | None) -> str | None:
    if value is None:
        return None

    trimmed = value.strip()
    if not trimmed:
        return ""

    return _validate_trace_archival_retention_string(
        trimmed, parameter_name="trace_archival_config.retention"
    )


def _get_workspace_request_message(
    request_message, schema: dict[str, list[Callable[..., Any]]] | None = None
) -> tuple[Any, dict[str, Any]]:
    request_json = _get_normalized_request_json()
    _validate_request_json_with_schema(request_json, schema, proto_parsing_succeeded=None)

    trace_archival_config_json = request_json.get("trace_archival_config")
    if trace_archival_config_json is not None:
        _validate_param_against_schema(
            schema=[_assert_dict],
            param="trace_archival_config",
            value=trace_archival_config_json,
            proto_parsing_succeeded=False,
        )
        _validate_request_json_with_schema(
            trace_archival_config_json,
            {
                "location": [_assert_string],
                "retention": [_assert_string],
            },
            proto_parsing_succeeded=None,
        )

    parse_dict(request_json, request_message)
    return request_message, request_json


def _ensure_artifact_root_available(workspace_artifact_root: str | None) -> None:
    """Ensure an artifact root is available either at workspace or server level.

    Args:
        workspace_artifact_root: The workspace's default_artifact_root value.
            - None means "not specified" (fallback to server default)
            - "" means "clear/unset" (fallback to server default)
            - non-empty string means "use this workspace-specific root"

    Raises:
        MlflowException: If neither workspace nor server has an artifact root configured.
    """
    # If workspace has a non-empty artifact root, it's valid
    if workspace_artifact_root:
        return

    # Otherwise, check if server has a default artifact root
    server_artifact_root = _get_tracking_store().artifact_root_uri
    if not server_artifact_root:
        raise MlflowException.invalid_parameter_value(
            "Cannot create or update workspace without an artifact root. Either specify "
            "'default_artifact_root' for this workspace or start the server with "
            "'--default-artifact-root'."
        )


@catch_mlflow_exception
@_disable_if_workspaces_disabled
def _list_workspaces_handler():
    _get_request_message(ListWorkspaces())
    workspaces = _get_workspace_store().list_workspaces()
    response_message = ListWorkspaces.Response()
    response_message.workspaces.extend([ws.to_proto() for ws in workspaces])
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_workspaces_disabled
def _create_workspace_handler():
    request_message, request_json = _get_workspace_request_message(
        CreateWorkspace(),
        schema={
            "name": [_assert_required, _assert_string],
            "description": [_assert_string],
            "default_artifact_root": [_assert_string],
        },
    )

    if request_message.name == DEFAULT_WORKSPACE_NAME:
        raise MlflowException.invalid_parameter_value(
            f"The '{DEFAULT_WORKSPACE_NAME}' workspace is reserved and cannot be created"
        )
    WorkspaceNameValidator.validate(request_message.name)
    description = request_message.description if request_message.HasField("description") else None
    default_artifact_root = (
        request_message.default_artifact_root
        if request_message.HasField("default_artifact_root")
        else None
    )
    trace_archival_config_json = request_json.get("trace_archival_config") or {}
    has_trace_archival_location = "location" in trace_archival_config_json
    has_trace_archival_retention = "retention" in trace_archival_config_json
    trace_archival_location = (
        request_message.trace_archival_config.location if has_trace_archival_location else None
    )
    trace_archival_retention = (
        request_message.trace_archival_config.retention if has_trace_archival_retention else None
    )
    default_artifact_root = _validate_workspace_default_artifact_root(default_artifact_root)
    trace_archival_location = _validate_workspace_trace_archival_location(trace_archival_location)
    trace_archival_retention = _validate_workspace_trace_archival_retention(
        trace_archival_retention
    )
    _ensure_artifact_root_available(default_artifact_root)
    store = _get_workspace_store()
    try:
        workspace = store.create_workspace(
            Workspace(
                name=request_message.name,
                description=description,
                default_artifact_root=default_artifact_root,
                trace_archival_location=trace_archival_location,
                trace_archival_retention=trace_archival_retention,
            )
        )
    except NotImplementedError:
        raise _workspace_not_supported("Workspace creation is not supported by this provider")

    response_message = CreateWorkspace.Response()
    response_message.workspace.MergeFrom(workspace.to_proto())
    response = _wrap_response(response_message)
    response.status_code = 201
    return response


@catch_mlflow_exception
@_disable_if_workspaces_disabled
def _get_workspace_handler(workspace_name: str):
    if workspace_name != DEFAULT_WORKSPACE_NAME:
        WorkspaceNameValidator.validate(workspace_name)
    workspace = _get_workspace_store().get_workspace(workspace_name)
    response_message = GetWorkspace.Response()
    response_message.workspace.MergeFrom(workspace.to_proto())
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_workspaces_disabled
def _update_workspace_handler(workspace_name: str):
    if workspace_name != DEFAULT_WORKSPACE_NAME:
        WorkspaceNameValidator.validate(workspace_name)
    request_message, request_json = _get_workspace_request_message(
        UpdateWorkspace(),
        schema={
            "description": [_assert_string],
            "default_artifact_root": [_assert_string],
        },
    )

    has_description = "description" in request_json
    has_artifact_root = "default_artifact_root" in request_json
    trace_archival_config_json = request_json.get("trace_archival_config") or {}
    has_trace_archival_location = "location" in trace_archival_config_json
    has_trace_archival_retention = "retention" in trace_archival_config_json

    if (
        not has_description
        and not has_artifact_root
        and not has_trace_archival_location
        and not has_trace_archival_retention
    ):
        raise MlflowException.invalid_parameter_value("Workspace update must have at least one key")

    description = request_message.description if has_description else None
    default_artifact_root = request_message.default_artifact_root if has_artifact_root else None
    trace_archival_location = (
        request_message.trace_archival_config.location if has_trace_archival_location else None
    )
    trace_archival_retention = (
        request_message.trace_archival_config.retention if has_trace_archival_retention else None
    )
    default_artifact_root = _validate_workspace_default_artifact_root(default_artifact_root)
    trace_archival_location = _validate_workspace_trace_archival_location(trace_archival_location)
    trace_archival_retention = _validate_workspace_trace_archival_retention(
        trace_archival_retention
    )

    # If the user is clearing the workspace artifact root (empty string), ensure the server
    # has a default artifact root configured
    if default_artifact_root == "":
        _ensure_artifact_root_available(default_artifact_root)

    store = _get_workspace_store()
    try:
        workspace = store.update_workspace(
            Workspace(
                name=workspace_name,
                description=description,
                default_artifact_root=default_artifact_root,
                trace_archival_location=trace_archival_location,
                trace_archival_retention=trace_archival_retention,
            )
        )
    except NotImplementedError:
        raise _workspace_not_supported("Workspace updates are not supported by this provider")

    response_message = UpdateWorkspace.Response()
    response_message.workspace.MergeFrom(workspace.to_proto())
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_workspaces_disabled
def _delete_workspace_handler(workspace_name: str):
    if workspace_name == DEFAULT_WORKSPACE_NAME:
        raise MlflowException.invalid_parameter_value(
            f"The '{DEFAULT_WORKSPACE_NAME}' workspace is reserved and cannot be deleted"
        )
    WorkspaceNameValidator.validate(workspace_name)
    mode_str = request.args.get("mode", WorkspaceDeletionMode.RESTRICT.value)
    try:
        mode = WorkspaceDeletionMode(mode_str)
    except ValueError:
        raise MlflowException.invalid_parameter_value(
            f"Invalid deletion mode '{mode_str}'. "
            f"Must be one of: {', '.join(m.value for m in WorkspaceDeletionMode)}"
        )
    store = _get_workspace_store()
    try:
        store.delete_workspace(workspace_name, mode=mode)
    except NotImplementedError:
        raise _workspace_not_supported("Workspace deletion is not supported by this provider")
    return Response(status=204)


@catch_mlflow_exception
def get_artifact_handler():
    run_id = request.args.get("run_id") or request.args.get("run_uuid")
    path = request.args["path"]
    path = validate_path_is_safe(path)
    run = _get_tracking_store().get_run(run_id)

    if _is_servable_proxied_run_artifact_root(run.info.artifact_uri):
        artifact_repo = _get_artifact_repo_mlflow_artifacts()
        artifact_path = _get_proxied_run_artifact_destination_path(
            proxied_artifact_root=run.info.artifact_uri,
            relative_path=path,
        )
        artifact_path = _get_workspace_scoped_repo_path_if_enabled(artifact_path)
    else:
        artifact_repo = _get_artifact_repo(run)
        artifact_path = path

    return _send_artifact(artifact_repo, artifact_path)


def _not_implemented():
    response = Response()
    response.status_code = 404
    return response


# Tracking Server APIs


@catch_mlflow_exception
@_disable_if_artifacts_only
def _create_experiment():
    request_message = _get_request_message(
        CreateExperiment(),
        schema={
            "name": [_assert_required, _assert_string],
            "artifact_location": [_assert_string],
            "tags": [_assert_array],
        },
    )

    tags = [ExperimentTag(tag.key, tag.value) for tag in request_message.tags]

    if request_message.artifact_location:
        _validate_storage_location_uri(request_message.artifact_location, "artifact_location")
    experiment_id = _get_tracking_store().create_experiment(
        request_message.name, request_message.artifact_location, tags
    )
    response_message = CreateExperiment.Response()
    response_message.experiment_id = experiment_id
    response = Response(mimetype="application/json")
    response.set_data(message_to_json(response_message))
    return response


@catch_mlflow_exception
@_disable_if_artifacts_only
def _get_experiment():
    request_message = _get_request_message(
        GetExperiment(), schema={"experiment_id": [_assert_required, _assert_string]}
    )
    response_message = get_experiment_impl(request_message)
    response = Response(mimetype="application/json")
    response.set_data(message_to_json(response_message))
    return response


def get_experiment_impl(request_message):
    response_message = GetExperiment.Response()
    experiment = _get_tracking_store().get_experiment(request_message.experiment_id).to_proto()
    response_message.experiment.MergeFrom(experiment)
    return response_message


@catch_mlflow_exception
@_disable_if_artifacts_only
def _get_experiment_by_name():
    request_message = _get_request_message(
        GetExperimentByName(),
        schema={"experiment_name": [_assert_required, _assert_string]},
    )
    response_message = GetExperimentByName.Response()
    store_exp = _get_tracking_store().get_experiment_by_name(request_message.experiment_name)
    if store_exp is None:
        raise MlflowException(
            f"Could not find experiment with name '{request_message.experiment_name}'",
            error_code=RESOURCE_DOES_NOT_EXIST,
        )
    experiment = store_exp.to_proto()
    response_message.experiment.MergeFrom(experiment)
    response = Response(mimetype="application/json")
    response.set_data(message_to_json(response_message))
    return response


@catch_mlflow_exception
@_disable_if_artifacts_only
def _delete_experiment():
    request_message = _get_request_message(
        DeleteExperiment(), schema={"experiment_id": [_assert_required, _assert_string]}
    )
    _get_tracking_store().delete_experiment(request_message.experiment_id)
    response_message = DeleteExperiment.Response()
    response = Response(mimetype="application/json")
    response.set_data(message_to_json(response_message))
    return response


@catch_mlflow_exception
@_disable_if_artifacts_only
def _restore_experiment():
    request_message = _get_request_message(
        RestoreExperiment(),
        schema={"experiment_id": [_assert_required, _assert_string]},
    )
    _get_tracking_store().restore_experiment(request_message.experiment_id)
    response_message = RestoreExperiment.Response()
    response = Response(mimetype="application/json")
    response.set_data(message_to_json(response_message))
    return response


@catch_mlflow_exception
@_disable_if_artifacts_only
def _update_experiment():
    request_message = _get_request_message(
        UpdateExperiment(),
        schema={
            "experiment_id": [_assert_required, _assert_string],
            "new_name": [_assert_string, _assert_required],
        },
    )
    if request_message.new_name:
        _get_tracking_store().rename_experiment(
            request_message.experiment_id, request_message.new_name
        )
    response_message = UpdateExperiment.Response()
    response = Response(mimetype="application/json")
    response.set_data(message_to_json(response_message))
    return response


@catch_mlflow_exception
@_disable_if_artifacts_only
def _create_run():
    request_message = _get_request_message(
        CreateRun(),
        schema={
            "experiment_id": [_assert_string],
            "start_time": [_assert_intlike],
            "run_name": [_assert_string],
        },
    )

    tags = [RunTag(tag.key, tag.value) for tag in request_message.tags]
    run = _get_tracking_store().create_run(
        experiment_id=request_message.experiment_id,
        user_id=request_message.user_id,
        start_time=request_message.start_time,
        tags=tags,
        run_name=request_message.run_name,
    )

    response_message = CreateRun.Response()
    response_message.run.MergeFrom(run.to_proto())
    response = Response(mimetype="application/json")
    response.set_data(message_to_json(response_message))
    return response


@catch_mlflow_exception
@_disable_if_artifacts_only
def _update_run():
    request_message = _get_request_message(
        UpdateRun(),
        schema={
            "run_id": [_assert_required, _assert_string],
            "end_time": [_assert_intlike],
            "status": [_assert_string],
            "run_name": [_assert_string],
        },
    )
    run_id = request_message.run_id or request_message.run_uuid
    run_name = request_message.run_name if request_message.HasField("run_name") else None
    end_time = request_message.end_time if request_message.HasField("end_time") else None
    status = request_message.status if request_message.HasField("status") else None
    updated_info = _get_tracking_store().update_run_info(run_id, status, end_time, run_name)
    response_message = UpdateRun.Response(run_info=updated_info.to_proto())
    response = Response(mimetype="application/json")
    response.set_data(message_to_json(response_message))
    return response


@catch_mlflow_exception
@_disable_if_artifacts_only
def _delete_run():
    request_message = _get_request_message(
        DeleteRun(), schema={"run_id": [_assert_required, _assert_string]}
    )
    _get_tracking_store().delete_run(request_message.run_id)
    response_message = DeleteRun.Response()
    response = Response(mimetype="application/json")
    response.set_data(message_to_json(response_message))
    return response


@catch_mlflow_exception
@_disable_if_artifacts_only
def _restore_run():
    request_message = _get_request_message(
        RestoreRun(), schema={"run_id": [_assert_required, _assert_string]}
    )
    _get_tracking_store().restore_run(request_message.run_id)
    response_message = RestoreRun.Response()
    response = Response(mimetype="application/json")
    response.set_data(message_to_json(response_message))
    return response


@catch_mlflow_exception
@_disable_if_artifacts_only
def _log_metric():
    request_message = _get_request_message(
        LogMetric(),
        schema={
            "run_id": [_assert_required, _assert_string],
            "key": [_assert_required, _assert_string],
            "value": [_assert_required, _assert_floatlike],
            "timestamp": [_assert_intlike, _assert_required],
            "step": [_assert_intlike],
            "model_id": [_assert_string],
            "dataset_name": [_assert_string],
            "dataset_digest": [_assert_string],
        },
    )
    metric = Metric(
        request_message.key,
        request_message.value,
        request_message.timestamp,
        request_message.step,
        request_message.model_id or None,
        request_message.dataset_name or None,
        request_message.dataset_digest or None,
        request_message.run_id or None,
    )
    run_id = request_message.run_id or request_message.run_uuid
    _get_tracking_store().log_metric(run_id, metric)
    response_message = LogMetric.Response()
    response = Response(mimetype="application/json")
    response.set_data(message_to_json(response_message))
    return response


@catch_mlflow_exception
@_disable_if_artifacts_only
def _log_param():
    request_message = _get_request_message(
        LogParam(),
        schema={
            "run_id": [_assert_required, _assert_string],
            "key": [_assert_required, _assert_string],
            "value": [_assert_string],
        },
    )
    param = Param(request_message.key, request_message.value)
    run_id = request_message.run_id or request_message.run_uuid
    _get_tracking_store().log_param(run_id, param)
    response_message = LogParam.Response()
    response = Response(mimetype="application/json")
    response.set_data(message_to_json(response_message))
    return response


@catch_mlflow_exception
@_disable_if_artifacts_only
def _log_inputs():
    request_message = _get_request_message(
        LogInputs(),
        schema={
            "run_id": [_assert_required, _assert_string],
            "datasets": [_assert_array],
            "models": [_assert_array],
        },
    )
    run_id = request_message.run_id
    datasets = [
        DatasetInput.from_proto(proto_dataset_input)
        for proto_dataset_input in request_message.datasets
    ]
    models = (
        [
            LoggedModelInput.from_proto(proto_logged_model_input)
            for proto_logged_model_input in request_message.models
        ]
        if request_message.models
        else None
    )

    _get_tracking_store().log_inputs(run_id, datasets=datasets, models=models)
    response_message = LogInputs.Response()
    response = Response(mimetype="application/json")
    response.set_data(message_to_json(response_message))
    return response


@catch_mlflow_exception
@_disable_if_artifacts_only
def _log_outputs():
    request_message = _get_request_message(
        LogOutputs(),
        schema={
            "run_id": [_assert_required, _assert_string],
            "models": [_assert_required, _assert_array],
        },
    )
    models = [LoggedModelOutput.from_proto(p) for p in request_message.models]
    _get_tracking_store().log_outputs(run_id=request_message.run_id, models=models)
    response_message = LogOutputs.Response()
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _set_experiment_tag():
    request_message = _get_request_message(
        SetExperimentTag(),
        schema={
            "experiment_id": [_assert_required, _assert_string],
            "key": [_assert_required, _assert_string],
            "value": [_assert_string],
        },
    )
    tag = ExperimentTag(request_message.key, request_message.value)
    _get_tracking_store().set_experiment_tag(request_message.experiment_id, tag)
    response_message = SetExperimentTag.Response()
    response = Response(mimetype="application/json")
    response.set_data(message_to_json(response_message))
    return response


@catch_mlflow_exception
@_disable_if_artifacts_only
def _delete_experiment_tag():
    request_message = _get_request_message(
        DeleteExperimentTag(),
        schema={
            "experiment_id": [_assert_required, _assert_string],
            "key": [_assert_required, _assert_string],
        },
    )
    _get_tracking_store().delete_experiment_tag(request_message.experiment_id, request_message.key)
    response_message = DeleteExperimentTag.Response()
    response = Response(mimetype="application/json")
    response.set_data(message_to_json(response_message))
    return response


@catch_mlflow_exception
@_disable_if_artifacts_only
def _set_tag():
    request_message = _get_request_message(
        SetTag(),
        schema={
            "run_id": [_assert_required, _assert_string],
            "key": [_assert_required, _assert_string],
            "value": [_assert_string],
        },
    )
    tag = RunTag(request_message.key, request_message.value)
    run_id = request_message.run_id or request_message.run_uuid
    _get_tracking_store().set_tag(run_id, tag)
    response_message = SetTag.Response()
    response = Response(mimetype="application/json")
    response.set_data(message_to_json(response_message))
    return response


@catch_mlflow_exception
@_disable_if_artifacts_only
def _delete_tag():
    request_message = _get_request_message(
        DeleteTag(),
        schema={
            "run_id": [_assert_required, _assert_string],
            "key": [_assert_required, _assert_string],
        },
    )
    _get_tracking_store().delete_tag(request_message.run_id, request_message.key)
    response_message = DeleteTag.Response()
    response = Response(mimetype="application/json")
    response.set_data(message_to_json(response_message))
    return response


@catch_mlflow_exception
@_disable_if_artifacts_only
def _get_run():
    request_message = _get_request_message(
        GetRun(), schema={"run_id": [_assert_required, _assert_string]}
    )
    response_message = get_run_impl(request_message)
    response = Response(mimetype="application/json")
    response.set_data(message_to_json(response_message))
    return response


def get_run_impl(request_message):
    response_message = GetRun.Response()
    run_id = request_message.run_id or request_message.run_uuid
    response_message.run.MergeFrom(_get_tracking_store().get_run(run_id).to_proto())
    return response_message


@catch_mlflow_exception
@_disable_if_artifacts_only
def _search_runs():
    request_message = _get_request_message(
        SearchRuns(),
        schema={
            "experiment_ids": [_assert_array],
            "filter": [_assert_string],
            "max_results": [
                _assert_intlike,
                lambda x: _assert_less_than_or_equal(int(x), 50000),
            ],
            "order_by": [_assert_array, _assert_item_type_string],
        },
    )
    response_message = search_runs_impl(request_message)
    response = Response(mimetype="application/json")
    response.set_data(message_to_json(response_message))
    return response


def search_runs_impl(request_message):
    response_message = SearchRuns.Response()
    run_view_type = ViewType.ACTIVE_ONLY
    if request_message.HasField("run_view_type"):
        run_view_type = ViewType.from_proto(request_message.run_view_type)
    filter_string = request_message.filter
    max_results = request_message.max_results
    experiment_ids = list(request_message.experiment_ids)

    # NB: Local import to avoid circular dependency (auth imports from handlers)
    try:
        from mlflow.server import auth

        if auth.auth_config:
            experiment_ids = auth.filter_experiment_ids(experiment_ids)
    except ImportError:
        # Auth module not available (Flask-WTF not installed), skip filtering
        pass

    order_by = request_message.order_by
    run_entities = _get_tracking_store().search_runs(
        experiment_ids=experiment_ids,
        filter_string=filter_string,
        run_view_type=run_view_type,
        max_results=max_results,
        order_by=order_by,
        page_token=request_message.page_token or None,
    )
    response_message.runs.extend([r.to_proto() for r in run_entities])
    if run_entities.token:
        response_message.next_page_token = run_entities.token
    return response_message


@catch_mlflow_exception
@_disable_if_artifacts_only
def _list_artifacts():
    request_message = _get_request_message(
        ListArtifacts(),
        schema={
            "run_id": [_assert_string, _assert_required],
            "path": [_assert_string],
            "page_token": [_assert_string],
        },
    )
    response_message = list_artifacts_impl(request_message)
    response = Response(mimetype="application/json")
    response.set_data(message_to_json(response_message))
    return response


def list_artifacts_impl(request_message):
    response_message = ListArtifacts.Response()
    if request_message.HasField("path"):
        path = request_message.path
        path = validate_path_is_safe(path)
    else:
        path = None
    run_id = request_message.run_id or request_message.run_uuid
    run = _get_tracking_store().get_run(run_id)

    if _is_servable_proxied_run_artifact_root(run.info.artifact_uri):
        artifact_entities = _list_artifacts_for_proxied_run_artifact_root(
            proxied_artifact_root=run.info.artifact_uri,
            relative_path=path,
        )
    else:
        artifact_entities = _get_artifact_repo(run).list_artifacts(path)

    response_message.files.extend([a.to_proto() for a in artifact_entities])
    response_message.root_uri = run.info.artifact_uri
    return response_message


def _list_artifacts_for_proxied_run_artifact_root(proxied_artifact_root, relative_path=None):
    """
    Lists artifacts from the specified ``relative_path`` within the specified proxied Run artifact
    root (i.e. a Run artifact root with scheme ``http``, ``https``, or ``mlflow-artifacts``).

    Args:
        proxied_artifact_root: The Run artifact root location (URI) with scheme ``http``,
                               ``https``, or ``mlflow-artifacts`` that can be resolved by the
                               MLflow server to a concrete storage location.
        relative_path: The relative path within the specified ``proxied_artifact_root`` under
                       which to list artifact contents. If ``None``, artifacts are listed from
                       the ``proxied_artifact_root`` directory.
    """
    parsed_proxied_artifact_root = urllib.parse.urlparse(proxied_artifact_root)
    assert parsed_proxied_artifact_root.scheme in ["http", "https", "mlflow-artifacts"]

    artifact_destination_repo = _get_artifact_repo_mlflow_artifacts()
    artifact_destination_path = _get_proxied_run_artifact_destination_path(
        proxied_artifact_root=proxied_artifact_root,
        relative_path=relative_path,
    )
    artifact_destination_path = _get_workspace_scoped_repo_path_if_enabled(
        artifact_destination_path
    )

    artifact_entities = []
    for file_info in artifact_destination_repo.list_artifacts(artifact_destination_path):
        basename = posixpath.basename(file_info.path)
        run_relative_artifact_path = (
            posixpath.join(relative_path, basename) if relative_path else basename
        )
        artifact_entities.append(
            FileInfo(run_relative_artifact_path, file_info.is_dir, file_info.file_size)
        )

    return artifact_entities


@catch_mlflow_exception
@_disable_if_artifacts_only
def _get_metric_history():
    request_message = _get_request_message(
        GetMetricHistory(),
        schema={
            "run_id": [_assert_string, _assert_required],
            "metric_key": [_assert_string, _assert_required],
            "page_token": [_assert_string],
        },
    )
    response_message = GetMetricHistory.Response()
    run_id = request_message.run_id or request_message.run_uuid

    # NB: An unset proto2 int field reads as 0 (never None), so a `max_results is not None`
    # check would treat requests without `max_results` as `max_results=0`: the store queries
    # one row beyond the requested page size (LIMIT 1), concludes more results exist,
    # truncates the page to zero metrics, and emits a token for `offset + 0` that points back
    # at the same position forever. Use HasField to keep requests without `max_results` on the
    # documented non-paginated path, and reject explicit non-positive page sizes.
    max_results = request_message.max_results if request_message.HasField("max_results") else None
    if max_results is not None and max_results <= 0:
        raise MlflowException(
            f"Invalid value {max_results} for parameter 'max_results' supplied. "
            "It must be a positive integer.",
            error_code=INVALID_PARAMETER_VALUE,
        )

    metric_entities = _get_tracking_store().get_metric_history(
        run_id,
        request_message.metric_key,
        max_results=max_results,
        page_token=request_message.page_token or None,
    )
    response_message.metrics.extend([m.to_proto() for m in metric_entities])

    # Set next_page_token if available
    if next_page_token := metric_entities.token:
        response_message.next_page_token = next_page_token

    response = Response(mimetype="application/json")
    response.set_data(message_to_json(response_message))
    return response


@catch_mlflow_exception
@_disable_if_artifacts_only
def get_metric_history_bulk_handler():
    MAX_HISTORY_RESULTS = 25000
    MAX_RUN_IDS_PER_REQUEST = 100
    run_ids = request.args.to_dict(flat=False).get("run_id", [])
    if not run_ids:
        raise MlflowException(
            message="GetMetricHistoryBulk request must specify at least one run_id.",
            error_code=INVALID_PARAMETER_VALUE,
        )
    if len(run_ids) > MAX_RUN_IDS_PER_REQUEST:
        raise MlflowException(
            message=(
                f"GetMetricHistoryBulk request cannot specify more than {MAX_RUN_IDS_PER_REQUEST}"
                f" run_ids. Received {len(run_ids)} run_ids."
            ),
            error_code=INVALID_PARAMETER_VALUE,
        )

    metric_key = request.args.get("metric_key")
    if metric_key is None:
        raise MlflowException(
            message="GetMetricHistoryBulk request must specify a metric_key.",
            error_code=INVALID_PARAMETER_VALUE,
        )

    max_results = int(request.args.get("max_results", MAX_HISTORY_RESULTS))
    max_results = min(max_results, MAX_HISTORY_RESULTS)

    store = _get_tracking_store()

    def _default_history_bulk_impl():
        metrics_with_run_ids = []
        for run_id in sorted(run_ids):
            metrics_for_run = sorted(
                store.get_metric_history(
                    run_id=run_id,
                    metric_key=metric_key,
                    max_results=max_results,
                ),
                key=lambda metric: (metric.timestamp, metric.step, metric.value),
            )
            metrics_with_run_ids.extend([
                {
                    "key": metric.key,
                    "value": metric.value,
                    "timestamp": metric.timestamp,
                    "step": metric.step,
                    "run_id": run_id,
                }
                for metric in metrics_for_run
            ])
        return metrics_with_run_ids

    if hasattr(store, "get_metric_history_bulk"):
        metrics_with_run_ids = [
            metric.to_dict()
            for metric in store.get_metric_history_bulk(
                run_ids=run_ids,
                metric_key=metric_key,
                max_results=max_results,
            )
        ]
    else:
        metrics_with_run_ids = _default_history_bulk_impl()

    return {
        "metrics": metrics_with_run_ids[:max_results],
    }


@catch_mlflow_exception
@_disable_if_artifacts_only
def get_metric_history_bulk_interval_handler():
    request_message = _get_request_message(
        GetMetricHistoryBulkInterval(),
        schema={
            "run_ids": [
                _assert_required,
                _assert_array,
                _assert_item_type_string,
                lambda x: _assert_less_than_or_equal(
                    len(x),
                    MAX_RUNS_GET_METRIC_HISTORY_BULK,
                    message=f"GetMetricHistoryBulkInterval request must specify at most "
                    f"{MAX_RUNS_GET_METRIC_HISTORY_BULK} run_ids. Received {len(x)} run_ids.",
                ),
            ],
            "metric_key": [_assert_required, _assert_string],
            "start_step": [_assert_intlike],
            "end_step": [_assert_intlike],
            "max_results": [
                _assert_intlike,
                lambda x: _assert_intlike_within_range(
                    int(x),
                    1,
                    MAX_RESULTS_PER_RUN,
                    message=f"max_results must be between 1 and {MAX_RESULTS_PER_RUN}.",
                ),
            ],
        },
    )
    response_message = get_metric_history_bulk_interval_impl(request_message)
    response = Response(mimetype="application/json")
    response.set_data(message_to_json(response_message))
    return response


def get_metric_history_bulk_interval_impl(request_message):
    args = request.args
    run_ids = request_message.run_ids
    metric_key = request_message.metric_key
    max_results = int(args.get("max_results", MAX_RESULTS_PER_RUN))
    start_step = args.get("start_step")
    end_step = args.get("end_step")
    if start_step is not None and end_step is not None:
        start_step = int(start_step)
        end_step = int(end_step)
        if start_step > end_step:
            raise MlflowException.invalid_parameter_value(
                "end_step must be greater than start_step. "
                f"Found start_step={start_step} and end_step={end_step}."
            )
    elif start_step is not None or end_step is not None:
        raise MlflowException.invalid_parameter_value(
            "If either start step or end step are specified, both must be specified."
        )

    store = _get_tracking_store()
    metrics_with_run_ids = store.get_metric_history_bulk_interval(
        run_ids=run_ids,
        metric_key=metric_key,
        max_results=max_results,
        start_step=start_step,
        end_step=end_step,
    )

    response_message = GetMetricHistoryBulkInterval.Response()
    response_message.metrics.extend([m.to_proto() for m in metrics_with_run_ids])
    return response_message


@catch_mlflow_exception
@_disable_if_artifacts_only
def _search_datasets_handler():
    request_message = _get_request_message(
        SearchDatasets(),
    )
    response_message = search_datasets_impl(request_message)
    response = Response(mimetype="application/json")
    response.set_data(message_to_json(response_message))
    return response


def search_datasets_impl(request_message):
    MAX_EXPERIMENT_IDS_PER_REQUEST = 20
    _validate_content_type(request, ["application/json"])
    experiment_ids = request_message.experiment_ids or []
    if not experiment_ids:
        raise MlflowException(
            message="SearchDatasets request must specify at least one experiment_id.",
            error_code=INVALID_PARAMETER_VALUE,
        )
    if len(experiment_ids) > MAX_EXPERIMENT_IDS_PER_REQUEST:
        raise MlflowException(
            message=(
                f"SearchDatasets request cannot specify more than {MAX_EXPERIMENT_IDS_PER_REQUEST}"
                f" experiment_ids. Received {len(experiment_ids)} experiment_ids."
            ),
            error_code=INVALID_PARAMETER_VALUE,
        )

    store = _get_tracking_store()

    if hasattr(store, "_search_datasets"):
        response_message = SearchDatasets.Response()
        response_message.dataset_summaries.extend([
            summary.to_proto() for summary in store._search_datasets(experiment_ids)
        ])
        return response_message
    else:
        return _not_implemented()


def _validate_gateway_path(method: str, gateway_path: str) -> None:
    if not gateway_path:
        raise MlflowException(
            message="Deployments proxy request must specify a gateway_path.",
            error_code=INVALID_PARAMETER_VALUE,
        )
    elif method == "GET":
        if gateway_path.strip("/") != "api/2.0/endpoints":
            raise MlflowException(
                message=f"Invalid gateway_path: {gateway_path} for method: {method}",
                error_code=INVALID_PARAMETER_VALUE,
            )
    elif method == "POST":
        # For POST, gateway_path must be in the form of "gateway/{name}/invocations"
        if not re.fullmatch(r"gateway/[^/]+/invocations", gateway_path.strip("/")):
            raise MlflowException(
                message=f"Invalid gateway_path: {gateway_path} for method: {method}",
                error_code=INVALID_PARAMETER_VALUE,
            )


@catch_mlflow_exception
def gateway_proxy_handler():
    target_uri = MLFLOW_DEPLOYMENTS_TARGET.get()
    if not target_uri:
        # Pretend an empty gateway service is running
        return {"endpoints": []}

    args = request.args if request.method == "GET" else request.json
    gateway_path = args.get("gateway_path")
    _validate_gateway_path(request.method, gateway_path)
    json_data = args.get("json_data", None)
    response = requests.request(request.method, f"{target_uri}/{gateway_path}", json=json_data)
    if response.status_code == 200:
        return response.json()
    else:
        raise MlflowException(
            message=f"Deployments proxy request failed with error code {response.status_code}. "
            f"Error message: {response.text}",
            error_code=response.status_code,
        )


@catch_mlflow_exception
@_disable_if_artifacts_only
def create_promptlab_run_handler():
    def assert_arg_exists(arg_name, arg):
        if not arg:
            raise MlflowException(
                message=f"CreatePromptlabRun request must specify {arg_name}.",
                error_code=INVALID_PARAMETER_VALUE,
            )

    _validate_content_type(request, ["application/json"])

    args = request.json
    experiment_id = args.get("experiment_id")
    assert_arg_exists("experiment_id", experiment_id)
    run_name = args.get("run_name", None)
    tags = args.get("tags", [])
    prompt_template = args.get("prompt_template")
    assert_arg_exists("prompt_template", prompt_template)
    raw_prompt_parameters = args.get("prompt_parameters")
    assert_arg_exists("prompt_parameters", raw_prompt_parameters)
    prompt_parameters = [
        Param(param.get("key"), param.get("value")) for param in args.get("prompt_parameters")
    ]
    model_route = args.get("model_route")
    assert_arg_exists("model_route", model_route)
    raw_model_parameters = args.get("model_parameters", [])
    model_parameters = [
        Param(param.get("key"), param.get("value")) for param in raw_model_parameters
    ]
    model_input = args.get("model_input")
    assert_arg_exists("model_input", model_input)
    model_output = args.get("model_output", None)
    raw_model_output_parameters = args.get("model_output_parameters", [])
    model_output_parameters = [
        Param(param.get("key"), param.get("value")) for param in raw_model_output_parameters
    ]
    mlflow_version = args.get("mlflow_version")
    assert_arg_exists("mlflow_version", mlflow_version)
    user_id = args.get("user_id", "unknown")

    # use current time if not provided
    start_time = args.get("start_time", int(time.time() * 1000))

    store = _get_tracking_store()

    run = _create_promptlab_run_impl(
        store,
        experiment_id=experiment_id,
        run_name=run_name,
        tags=tags,
        prompt_template=prompt_template,
        prompt_parameters=prompt_parameters,
        model_route=model_route,
        model_parameters=model_parameters,
        model_input=model_input,
        model_output=model_output,
        model_output_parameters=model_output_parameters,
        mlflow_version=mlflow_version,
        user_id=user_id,
        start_time=start_time,
    )
    response_message = CreateRun.Response()
    response_message.run.MergeFrom(run.to_proto())
    response = Response(mimetype="application/json")
    response.set_data(message_to_json(response_message))
    return response


@catch_mlflow_exception
def upload_artifact_handler():
    args = request.args
    run_uuid = args.get("run_uuid")
    if not run_uuid:
        raise MlflowException(
            message="Request must specify run_uuid.",
            error_code=INVALID_PARAMETER_VALUE,
        )
    path = args.get("path")
    if not path:
        raise MlflowException(
            message="Request must specify path.",
            error_code=INVALID_PARAMETER_VALUE,
        )
    path = validate_path_is_safe(path)

    if request.content_length and request.content_length > 10 * 1024 * 1024:
        raise MlflowException(
            message="Artifact size is too large. Max size is 10MB.",
            error_code=INVALID_PARAMETER_VALUE,
        )

    data = request.data
    if not data:
        raise MlflowException(
            message="Request must specify data.",
            error_code=INVALID_PARAMETER_VALUE,
        )

    if len(data) > 10 * 1024 * 1024:
        raise MlflowException(
            message="Artifact size is too large. Max size is 10MB.",
            error_code=INVALID_PARAMETER_VALUE,
        )

    run = _get_tracking_store().get_run(run_uuid)
    artifact_dir = run.info.artifact_uri

    basename = posixpath.basename(path)
    dirname = posixpath.dirname(path)

    def _log_artifact_to_repo(file, run, dirname, artifact_dir):
        if _is_servable_proxied_run_artifact_root(run.info.artifact_uri):
            artifact_repo = _get_artifact_repo_mlflow_artifacts()
            # Use posixpath.join since these are logical artifact paths (not local filesystem paths)
            # that should always use forward slashes regardless of the platform.
            path_to_log = (
                posixpath.join(run.info.experiment_id, run.info.run_id, "artifacts", dirname)
                if dirname
                else posixpath.join(run.info.experiment_id, run.info.run_id, "artifacts")
            )
            path_to_log = _get_workspace_scoped_repo_path_if_enabled(path_to_log)
        else:
            artifact_repo = get_artifact_repository(artifact_dir)
            path_to_log = dirname

        artifact_repo.log_artifact(file, path_to_log)

    with tempfile.TemporaryDirectory() as tmpdir:
        dir_path = os.path.join(tmpdir, dirname) if dirname else tmpdir
        file_path = os.path.join(dir_path, basename)

        os.makedirs(dir_path, exist_ok=True)

        with open(file_path, "wb") as f:
            f.write(data)

        _log_artifact_to_repo(file_path, run, dirname, artifact_dir)

    return Response(mimetype="application/json")


@catch_mlflow_exception
@_disable_if_artifacts_only
def _search_experiments():
    request_message = _get_request_message(
        SearchExperiments(),
        schema={
            "view_type": [_assert_intlike],
            "max_results": [_assert_intlike],
            "order_by": [_assert_array],
            "filter": [_assert_string],
            "page_token": [_assert_string],
        },
    )

    experiment_entities = _get_tracking_store().search_experiments(
        view_type=request_message.view_type,
        max_results=request_message.max_results,
        order_by=request_message.order_by,
        filter_string=request_message.filter,
        page_token=request_message.page_token or None,
    )
    response_message = SearchExperiments.Response()
    response_message.experiments.extend([e.to_proto() for e in experiment_entities])
    if experiment_entities.token:
        response_message.next_page_token = experiment_entities.token
    response = Response(mimetype="application/json")
    response.set_data(message_to_json(response_message))
    return response


@catch_mlflow_exception
def _get_artifact_repo(run):
    return get_artifact_repository(run.info.artifact_uri)


_HANDLER_BLOCKED_TRACE_TAGS = frozenset({
    MLFLOW_TRACE_SPANS_LOCATION,
    MLFLOW_TRACE_ARCHIVE_LOCATION,
    MLFLOW_TRACE_ARCHIVAL_FAILURE,
})
_HANDLER_TRACE_TAGS_MUTABLE_ON_DELETE = frozenset({MLFLOW_TRACE_ARCHIVAL_FAILURE})


def _validate_trace_tag_handler_mutation(key: str, operation: str) -> None:
    # `mlflow.trace.archivalFailure` is system-managed for writes, but deleting it is the
    # supported way to clear a terminal archival failure and allow a later retry.
    if key in _HANDLER_BLOCKED_TRACE_TAGS and not (
        operation == "deleted" and key in _HANDLER_TRACE_TAGS_MUTABLE_ON_DELETE
    ):
        raise MlflowException.invalid_parameter_value(
            f"Tag '{key}' is immutable and cannot be {operation} on a trace."
        )


@catch_mlflow_exception
@_disable_if_artifacts_only
def _log_batch():
    def _assert_metrics_fields_present(metrics):
        for idx, m in enumerate(metrics):
            _assert_required(m.get("key"), path=f"metrics[{idx}].key")
            _assert_required(m.get("value"), path=f"metrics[{idx}].value")
            _assert_required(m.get("timestamp"), path=f"metrics[{idx}].timestamp")

    def _assert_params_fields_present(params):
        for idx, param in enumerate(params):
            _assert_required(param.get("key"), path=f"params[{idx}].key")

    def _assert_tags_fields_present(tags):
        for idx, tag in enumerate(tags):
            _assert_required(tag.get("key"), path=f"tags[{idx}].key")

    _validate_batch_log_api_req(_get_request_json())
    request_message = _get_request_message(
        LogBatch(),
        schema={
            "run_id": [_assert_string, _assert_required],
            "metrics": [_assert_array, _assert_metrics_fields_present],
            "params": [_assert_array, _assert_params_fields_present],
            "tags": [_assert_array, _assert_tags_fields_present],
        },
    )
    metrics = [Metric.from_proto(proto_metric) for proto_metric in request_message.metrics]
    params = [Param.from_proto(proto_param) for proto_param in request_message.params]
    tags = [RunTag.from_proto(proto_tag) for proto_tag in request_message.tags]
    _get_tracking_store().log_batch(
        run_id=request_message.run_id, metrics=metrics, params=params, tags=tags
    )
    response_message = LogBatch.Response()
    response = Response(mimetype="application/json")
    response.set_data(message_to_json(response_message))
    return response


@catch_mlflow_exception
@_disable_if_artifacts_only
def _log_model():
    request_message = _get_request_message(
        LogModel(),
        schema={
            "run_id": [_assert_string, _assert_required],
            "model_json": [_assert_string, _assert_required],
        },
    )
    try:
        model = json.loads(request_message.model_json)
    except Exception:
        raise MlflowException(
            f"Malformed model info. \n {request_message.model_json} \n is not a valid JSON.",
            error_code=INVALID_PARAMETER_VALUE,
        )

    missing_fields = {"artifact_path", "flavors", "utc_time_created", "run_id"} - set(model.keys())

    if missing_fields:
        raise MlflowException(
            f"Model json is missing mandatory fields: {missing_fields}",
            error_code=INVALID_PARAMETER_VALUE,
        )
    _get_tracking_store().record_logged_model(
        run_id=request_message.run_id, mlflow_model=Model.from_dict(model)
    )
    response_message = LogModel.Response()
    response = Response(mimetype="application/json")
    response.set_data(message_to_json(response_message))
    return response


def _wrap_response(response_message):
    response = Response(mimetype="application/json")
    response.set_data(message_to_json(response_message))
    return response


# Model Registry APIs


@catch_mlflow_exception
@_disable_if_artifacts_only
def _create_registered_model():
    request_message = _get_request_message(
        CreateRegisteredModel(),
        schema={
            "name": [_assert_string, _assert_required],
            "tags": [_assert_array],
            "description": [_assert_string],
        },
    )
    store = _get_model_registry_store()
    registered_model = store.create_registered_model(
        name=request_message.name,
        tags=request_message.tags,
        description=request_message.description,
    )
    response_message = CreateRegisteredModel.Response(registered_model=registered_model.to_proto())

    # Determine if this is a prompt based on the tags
    if _is_prompt_request(request_message):
        # Send prompt creation webhook
        deliver_webhook(
            event=WebhookEvent(WebhookEntity.PROMPT, WebhookAction.CREATED),
            payload=PromptCreatedPayload(
                name=request_message.name,
                tags={
                    t.key: t.value
                    for t in request_message.tags
                    if t.key not in {IS_PROMPT_TAG_KEY, PROMPT_TYPE_TAG_KEY}
                },
                description=request_message.description,
            ),
            store=store,
        )
    else:
        # Send regular model creation webhook
        deliver_webhook(
            event=WebhookEvent(WebhookEntity.REGISTERED_MODEL, WebhookAction.CREATED),
            payload=RegisteredModelCreatedPayload(
                name=request_message.name,
                tags={t.key: t.value for t in request_message.tags},
                description=request_message.description,
            ),
            store=store,
        )

    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _get_registered_model():
    request_message = _get_request_message(
        GetRegisteredModel(), schema={"name": [_assert_string, _assert_required]}
    )
    registered_model = _get_model_registry_store().get_registered_model(name=request_message.name)
    response_message = GetRegisteredModel.Response(registered_model=registered_model.to_proto())
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _update_registered_model():
    request_message = _get_request_message(
        UpdateRegisteredModel(),
        schema={
            "name": [_assert_string, _assert_required],
            "description": [_assert_string],
        },
    )
    name = request_message.name
    new_description = request_message.description
    registered_model = _get_model_registry_store().update_registered_model(
        name=name, description=new_description
    )
    response_message = UpdateRegisteredModel.Response(registered_model=registered_model.to_proto())
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _rename_registered_model():
    request_message = _get_request_message(
        RenameRegisteredModel(),
        schema={
            "name": [_assert_string, _assert_required],
            "new_name": [_assert_string, _assert_required],
        },
    )
    name = request_message.name
    new_name = request_message.new_name
    registered_model = _get_model_registry_store().rename_registered_model(
        name=name, new_name=new_name
    )
    response_message = RenameRegisteredModel.Response(registered_model=registered_model.to_proto())
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _delete_registered_model():
    request_message = _get_request_message(
        DeleteRegisteredModel(), schema={"name": [_assert_string, _assert_required]}
    )
    _get_model_registry_store().delete_registered_model(name=request_message.name)
    return _wrap_response(DeleteRegisteredModel.Response())


@catch_mlflow_exception
@_disable_if_artifacts_only
def _search_registered_models():
    request_message = _get_request_message(
        SearchRegisteredModels(),
        schema={
            "filter": [_assert_string],
            "max_results": [
                _assert_intlike,
                lambda x: _assert_less_than_or_equal(int(x), 1000),
            ],
            "order_by": [_assert_array, _assert_item_type_string],
            "page_token": [_assert_string],
        },
    )
    store = _get_model_registry_store()
    registered_models = store.search_registered_models(
        filter_string=request_message.filter,
        max_results=request_message.max_results,
        order_by=request_message.order_by,
        page_token=request_message.page_token or None,
    )
    response_message = SearchRegisteredModels.Response()
    response_message.registered_models.extend([e.to_proto() for e in registered_models])
    if registered_models.token:
        response_message.next_page_token = registered_models.token
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _get_latest_versions():
    request_message = _get_request_message(
        GetLatestVersions(),
        schema={
            "name": [_assert_string, _assert_required],
            "stages": [_assert_array, _assert_item_type_string],
        },
    )
    latest_versions = _get_model_registry_store().get_latest_versions(
        name=request_message.name, stages=request_message.stages
    )
    response_message = GetLatestVersions.Response()
    response_message.model_versions.extend([e.to_proto() for e in latest_versions])
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _set_registered_model_tag():
    request_message = _get_request_message(
        SetRegisteredModelTag(),
        schema={
            "name": [_assert_string, _assert_required],
            "key": [_assert_string, _assert_required],
            "value": [_assert_string],
        },
    )
    tag = RegisteredModelTag(key=request_message.key, value=request_message.value)
    store = _get_model_registry_store()
    store.set_registered_model_tag(name=request_message.name, tag=tag)

    if _is_prompt(request_message.name):
        # Send prompt tag set webhook
        deliver_webhook(
            event=WebhookEvent(WebhookEntity.PROMPT_TAG, WebhookAction.SET),
            payload=PromptTagSetPayload(
                name=request_message.name,
                key=request_message.key,
                value=request_message.value,
            ),
            store=store,
        )

    return _wrap_response(SetRegisteredModelTag.Response())


@catch_mlflow_exception
@_disable_if_artifacts_only
def _delete_registered_model_tag():
    request_message = _get_request_message(
        DeleteRegisteredModelTag(),
        schema={
            "name": [_assert_string, _assert_required],
            "key": [_assert_string, _assert_required],
        },
    )
    store = _get_model_registry_store()
    store.delete_registered_model_tag(name=request_message.name, key=request_message.key)

    if _is_prompt(request_message.name):
        # Send prompt tag deleted webhook
        deliver_webhook(
            event=WebhookEvent(WebhookEntity.PROMPT_TAG, WebhookAction.DELETED),
            payload=PromptTagDeletedPayload(
                name=request_message.name,
                key=request_message.key,
            ),
            store=store,
        )

    return _wrap_response(DeleteRegisteredModelTag.Response())


def _validate_non_local_source_contains_relative_paths(source: str):
    """
    Validation check to ensure that sources that are provided that conform to the schemes:
    http, https, or mlflow-artifacts do not contain relative path designations that are intended
    to access local file system paths on the tracking server.

    Example paths that this validation function is intended to find and raise an Exception if
    passed:
    "mlflow-artifacts://host:port/../../../../"
    "http://host:port/api/2.0/mlflow-artifacts/artifacts/../../../../"
    "https://host:port/api/2.0/mlflow-artifacts/artifacts/../../../../"
    "/models/artifacts/../../../"
    "s3:/my_bucket/models/path/../../other/path"
    "file://path/to/../../../../some/where/you/should/not/be"
    "mlflow-artifacts://host:port/..%2f..%2f..%2f..%2f"
    "http://host:port/api/2.0/mlflow-artifacts/artifacts%00"
    """
    invalid_source_error_message = (
        f"Invalid model version source: '{source}'. If supplying a source as an http, https, "
        "local file path, ftp, objectstore, or mlflow-artifacts uri, an absolute path must be "
        "provided without relative path references present. "
        "Please provide an absolute path."
    )

    while (unquoted := urllib.parse.unquote_plus(source)) != source:
        source = unquoted
    source_path = re.sub(r"/+", "/", urllib.parse.urlparse(source).path.rstrip("/"))
    if "\x00" in source_path or any(p == ".." for p in source.split("/")):
        raise MlflowException(invalid_source_error_message, INVALID_PARAMETER_VALUE)
    resolved_source = pathlib.Path(source_path).resolve().as_posix()
    # NB: drive split is specifically for Windows since WindowsPath.resolve() will append the
    # drive path of the pwd to a given path. We don't care about the drive here, though.
    _, resolved_path = os.path.splitdrive(resolved_source)

    if resolved_path != source_path:
        raise MlflowException(invalid_source_error_message, INVALID_PARAMETER_VALUE)


def _validate_source_run(source: str, run_id: str) -> None:
    if is_local_uri(source):
        if run_id:
            store = _get_tracking_store()
            run = store.get_run(run_id)
            source = pathlib.Path(local_file_uri_to_path(source)).resolve()
            if is_local_uri(run.info.artifact_uri):
                run_artifact_dir = pathlib.Path(
                    local_file_uri_to_path(run.info.artifact_uri)
                ).resolve()
                if run_artifact_dir in [source, *source.parents]:
                    return

        raise MlflowException(
            f"Invalid model version source: '{source}'. To use a local path as a model version "
            "source, the run_id request parameter has to be specified and the local path has to be "
            "contained within the artifact directory of the run specified by the run_id.",
            INVALID_PARAMETER_VALUE,
        )

    # Checks if relative paths are present in the source (a security threat). If any are present,
    # raises an Exception.
    _validate_non_local_source_contains_relative_paths(source)


def _validate_source_model(source: str, model_id: str) -> None:
    if is_local_uri(source):
        if model_id:
            store = _get_tracking_store()
            model = store.get_logged_model(model_id)
            source = pathlib.Path(local_file_uri_to_path(source)).resolve()
            if is_local_uri(model.artifact_location):
                run_artifact_dir = pathlib.Path(
                    local_file_uri_to_path(model.artifact_location)
                ).resolve()
                if run_artifact_dir in [source, *source.parents]:
                    return

        raise MlflowException(
            f"Invalid model version source: '{source}'. To use a local path as a model version "
            "source, the model_id request parameter has to be specified and the local path has to "
            "be contained within the artifact directory of the run specified by the model_id.",
            INVALID_PARAMETER_VALUE,
        )

    # Checks if relative paths are present in the source (a security threat). If any are present,
    # raises an Exception.
    _validate_non_local_source_contains_relative_paths(source)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _create_model_version():
    request_message = _get_request_message(
        CreateModelVersion(),
        schema={
            "name": [_assert_string, _assert_required],
            "source": [_assert_string, _assert_required],
            "run_id": [_assert_string],
            "tags": [_assert_array],
            "run_link": [_assert_string],
            "description": [_assert_string],
            "model_id": [_assert_string],
        },
    )

    if request_message.source and (
        regex := MLFLOW_CREATE_MODEL_VERSION_SOURCE_VALIDATION_REGEX.get()
    ):
        if not re.search(regex, request_message.source):
            raise MlflowException(
                f"Invalid model version source: '{request_message.source}'.",
                error_code=INVALID_PARAMETER_VALUE,
            )

    is_prompt = _is_prompt_request(request_message)
    if is_prompt:
        # Prompt sources must not point to local filesystem paths.
        # Block file:// URIs and absolute paths (e.g. /etc/passwd) but allow
        # the legitimate schemeless placeholder sources used internally
        # (e.g. "prompt-template", "dummy-source").
        source = request_message.source
        parsed = urllib.parse.urlparse(source)
        if parsed.scheme == "file" or (parsed.scheme == "" and source.startswith("/")):
            raise MlflowException(
                f"Invalid prompt source: '{source}'. "
                "Local source paths are not allowed for prompts.",
                INVALID_PARAMETER_VALUE,
            )
        # Only validate traversal for sources with a URL scheme (http, https, etc.)
        if parsed.scheme:
            _validate_non_local_source_contains_relative_paths(source)
    else:
        if request_message.model_id:
            _validate_source_model(request_message.source, request_message.model_id)
        else:
            _validate_source_run(request_message.source, request_message.run_id)

    store = _get_model_registry_store()
    model_version = store.create_model_version(
        name=request_message.name,
        source=request_message.source,
        run_id=request_message.run_id,
        run_link=request_message.run_link,
        tags=request_message.tags,
        description=request_message.description,
        model_id=request_message.model_id,
    )
    if not is_prompt and request_message.model_id:
        tracking_store = _get_tracking_store()
        tracking_store.set_model_versions_tags(
            name=request_message.name,
            version=model_version.version,
            model_id=request_message.model_id,
        )
    response_message = CreateModelVersion.Response(model_version=model_version.to_proto())

    if is_prompt:
        # Convert tags to dict and extract template text efficiently
        tags_dict = {t.key: t.value for t in request_message.tags}
        template_text = tags_dict.pop(PROMPT_TEXT_TAG_KEY, None)
        # Remove internal prompt identification and type tags
        tags_dict.pop(IS_PROMPT_TAG_KEY, None)
        tags_dict.pop(PROMPT_TYPE_TAG_KEY, None)

        # Send prompt version creation webhook
        deliver_webhook(
            event=WebhookEvent(WebhookEntity.PROMPT_VERSION, WebhookAction.CREATED),
            payload=PromptVersionCreatedPayload(
                name=request_message.name,
                version=str(model_version.version),
                template=template_text,
                tags=tags_dict,
                description=request_message.description or None,
            ),
            store=store,
        )
    else:
        # Send regular model version creation webhook
        deliver_webhook(
            event=WebhookEvent(WebhookEntity.MODEL_VERSION, WebhookAction.CREATED),
            payload=ModelVersionCreatedPayload(
                name=request_message.name,
                version=str(model_version.version),
                source=request_message.source,
                run_id=request_message.run_id or None,
                tags={t.key: t.value for t in request_message.tags},
                description=request_message.description or None,
            ),
            store=store,
        )

    return _wrap_response(response_message)


def _is_prompt_request(request_message):
    return any(tag.key == IS_PROMPT_TAG_KEY for tag in request_message.tags)


def _is_prompt(name: str) -> bool:
    rm = _get_model_registry_store().get_registered_model(name=name)
    return rm._is_prompt()


@catch_mlflow_exception
@_disable_if_artifacts_only
def get_model_version_artifact_handler():
    name = request.args.get("name")
    version = request.args.get("version")
    path = request.args["path"]
    path = validate_path_is_safe(path)
    artifact_uri = _get_model_registry_store().get_model_version_download_uri(name, version)
    if _is_servable_proxied_run_artifact_root(artifact_uri):
        artifact_repo = _get_artifact_repo_mlflow_artifacts()
        artifact_path = _get_proxied_run_artifact_destination_path(
            proxied_artifact_root=artifact_uri,
            relative_path=path,
        )
        artifact_path = _get_workspace_scoped_repo_path_if_enabled(artifact_path)
    else:
        artifact_repo = get_artifact_repository(artifact_uri)
        artifact_path = path

    return _send_artifact(artifact_repo, artifact_path)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _get_model_version():
    request_message = _get_request_message(
        GetModelVersion(),
        schema={
            "name": [_assert_string, _assert_required],
            "version": [_assert_string, _assert_required],
        },
    )
    model_version = _get_model_registry_store().get_model_version(
        name=request_message.name, version=request_message.version
    )
    response_proto = model_version.to_proto()
    response_message = GetModelVersion.Response(model_version=response_proto)
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _update_model_version():
    request_message = _get_request_message(
        UpdateModelVersion(),
        schema={
            "name": [_assert_string, _assert_required],
            "version": [_assert_string, _assert_required],
            "description": [_assert_string],
        },
    )
    new_description = None
    if request_message.HasField("description"):
        new_description = request_message.description
    model_version = _get_model_registry_store().update_model_version(
        name=request_message.name,
        version=request_message.version,
        description=new_description,
    )
    return _wrap_response(UpdateModelVersion.Response(model_version=model_version.to_proto()))


@catch_mlflow_exception
@_disable_if_artifacts_only
def _transition_stage():
    request_message = _get_request_message(
        TransitionModelVersionStage(),
        schema={
            "name": [_assert_string, _assert_required],
            "version": [_assert_string, _assert_required],
            "stage": [_assert_string, _assert_required],
            "archive_existing_versions": [_assert_bool],
        },
    )
    model_version = _get_model_registry_store().transition_model_version_stage(
        name=request_message.name,
        version=request_message.version,
        stage=request_message.stage,
        archive_existing_versions=request_message.archive_existing_versions,
    )
    return _wrap_response(
        TransitionModelVersionStage.Response(model_version=model_version.to_proto())
    )


@catch_mlflow_exception
@_disable_if_artifacts_only
def _delete_model_version():
    request_message = _get_request_message(
        DeleteModelVersion(),
        schema={
            "name": [_assert_string, _assert_required],
            "version": [_assert_string, _assert_required],
        },
    )
    _get_model_registry_store().delete_model_version(
        name=request_message.name, version=request_message.version
    )
    return _wrap_response(DeleteModelVersion.Response())


@catch_mlflow_exception
@_disable_if_artifacts_only
def _get_model_version_download_uri():
    request_message = _get_request_message(GetModelVersionDownloadUri())
    download_uri = _get_model_registry_store().get_model_version_download_uri(
        name=request_message.name, version=request_message.version
    )
    response_message = GetModelVersionDownloadUri.Response(artifact_uri=download_uri)
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _search_model_versions():
    request_message = _get_request_message(
        SearchModelVersions(),
        schema={
            "filter": [_assert_string],
            "max_results": [
                _assert_intlike,
                lambda x: _assert_less_than_or_equal(int(x), 200_000),
            ],
            "order_by": [_assert_array, _assert_item_type_string],
            "page_token": [_assert_string],
        },
    )
    response_message = search_model_versions_impl(request_message)
    return _wrap_response(response_message)


def search_model_versions_impl(request_message):
    store = _get_model_registry_store()
    model_versions = store.search_model_versions(
        filter_string=request_message.filter,
        max_results=request_message.max_results,
        order_by=request_message.order_by,
        page_token=request_message.page_token or None,
    )
    response_message = SearchModelVersions.Response()
    response_message.model_versions.extend([e.to_proto() for e in model_versions])
    if model_versions.token:
        response_message.next_page_token = model_versions.token
    return response_message


@catch_mlflow_exception
@_disable_if_artifacts_only
def _set_model_version_tag():
    request_message = _get_request_message(
        SetModelVersionTag(),
        schema={
            "name": [_assert_string, _assert_required],
            "version": [_assert_string, _assert_required],
            "key": [_assert_string, _assert_required],
            "value": [_assert_string],
        },
    )
    tag = ModelVersionTag(key=request_message.key, value=request_message.value)
    store = _get_model_registry_store()
    store.set_model_version_tag(name=request_message.name, version=request_message.version, tag=tag)

    if _is_prompt(request_message.name):
        # Send prompt version tag set webhook
        deliver_webhook(
            event=WebhookEvent(WebhookEntity.PROMPT_VERSION_TAG, WebhookAction.SET),
            payload=PromptVersionTagSetPayload(
                name=request_message.name,
                version=request_message.version,
                key=request_message.key,
                value=request_message.value,
            ),
            store=store,
        )
    else:
        # Send regular model version tag set webhook
        deliver_webhook(
            event=WebhookEvent(WebhookEntity.MODEL_VERSION_TAG, WebhookAction.SET),
            payload=ModelVersionTagSetPayload(
                name=request_message.name,
                version=request_message.version,
                key=request_message.key,
                value=request_message.value,
            ),
            store=store,
        )

    return _wrap_response(SetModelVersionTag.Response())


@catch_mlflow_exception
@_disable_if_artifacts_only
def _delete_model_version_tag():
    request_message = _get_request_message(
        DeleteModelVersionTag(),
        schema={
            "name": [_assert_string, _assert_required],
            "version": [_assert_string, _assert_required],
            "key": [_assert_string, _assert_required],
        },
    )
    store = _get_model_registry_store()
    store.delete_model_version_tag(
        name=request_message.name,
        version=request_message.version,
        key=request_message.key,
    )

    if _is_prompt(request_message.name):
        # Send prompt version tag deleted webhook
        deliver_webhook(
            event=WebhookEvent(WebhookEntity.PROMPT_VERSION_TAG, WebhookAction.DELETED),
            payload=PromptVersionTagDeletedPayload(
                name=request_message.name,
                version=request_message.version,
                key=request_message.key,
            ),
            store=store,
        )
    else:
        # Send regular model version tag deleted webhook
        deliver_webhook(
            event=WebhookEvent(WebhookEntity.MODEL_VERSION_TAG, WebhookAction.DELETED),
            payload=ModelVersionTagDeletedPayload(
                name=request_message.name,
                version=request_message.version,
                key=request_message.key,
            ),
            store=store,
        )

    return _wrap_response(DeleteModelVersionTag.Response())


@catch_mlflow_exception
@_disable_if_artifacts_only
def _set_registered_model_alias():
    request_message = _get_request_message(
        SetRegisteredModelAlias(),
        schema={
            "name": [_assert_string, _assert_required],
            "alias": [_assert_string, _assert_required],
            "version": [_assert_string, _assert_required],
        },
    )
    store = _get_model_registry_store()
    store.set_registered_model_alias(
        name=request_message.name,
        alias=request_message.alias,
        version=request_message.version,
    )

    if _is_prompt(request_message.name):
        # Send prompt alias created webhook
        deliver_webhook(
            event=WebhookEvent(WebhookEntity.PROMPT_ALIAS, WebhookAction.CREATED),
            payload=PromptAliasCreatedPayload(
                name=request_message.name,
                alias=request_message.alias,
                version=request_message.version,
            ),
            store=store,
        )
    else:
        # Send regular model version alias created webhook
        deliver_webhook(
            event=WebhookEvent(WebhookEntity.MODEL_VERSION_ALIAS, WebhookAction.CREATED),
            payload=ModelVersionAliasCreatedPayload(
                name=request_message.name,
                alias=request_message.alias,
                version=request_message.version,
            ),
            store=store,
        )

    return _wrap_response(SetRegisteredModelAlias.Response())


@catch_mlflow_exception
@_disable_if_artifacts_only
def _delete_registered_model_alias():
    request_message = _get_request_message(
        DeleteRegisteredModelAlias(),
        schema={
            "name": [_assert_string, _assert_required],
            "alias": [_assert_string, _assert_required],
        },
    )
    store = _get_model_registry_store()
    store.delete_registered_model_alias(name=request_message.name, alias=request_message.alias)

    if _is_prompt(request_message.name):
        # Send prompt alias deleted webhook
        deliver_webhook(
            event=WebhookEvent(WebhookEntity.PROMPT_ALIAS, WebhookAction.DELETED),
            payload=PromptAliasDeletedPayload(
                name=request_message.name,
                alias=request_message.alias,
            ),
            store=store,
        )
    else:
        # Send regular model version alias deleted webhook
        deliver_webhook(
            event=WebhookEvent(WebhookEntity.MODEL_VERSION_ALIAS, WebhookAction.DELETED),
            payload=ModelVersionAliasDeletedPayload(
                name=request_message.name,
                alias=request_message.alias,
            ),
            store=store,
        )

    return _wrap_response(DeleteRegisteredModelAlias.Response())


@catch_mlflow_exception
@_disable_if_artifacts_only
def _get_model_version_by_alias():
    request_message = _get_request_message(
        GetModelVersionByAlias(),
        schema={
            "name": [_assert_string, _assert_required],
            "alias": [_assert_string, _assert_required],
        },
    )
    model_version = _get_model_registry_store().get_model_version_by_alias(
        name=request_message.name, alias=request_message.alias
    )
    response_proto = model_version.to_proto()
    response_message = GetModelVersionByAlias.Response(model_version=response_proto)
    return _wrap_response(response_message)


# Webhook APIs
@catch_mlflow_exception
@_disable_if_artifacts_only
def _create_webhook():
    request_message = _get_request_message(
        CreateWebhook(),
        schema={
            "name": [_assert_string, _assert_required],
            "url": [_assert_string, _assert_required],
            "events": [_assert_array, _assert_required],
            "description": [_assert_string],
            "secret": [_assert_string],
            "status": [_assert_string],
        },
    )

    webhook = _get_model_registry_store().create_webhook(
        name=request_message.name,
        url=request_message.url,
        events=[WebhookEvent.from_proto(e) for e in request_message.events],
        description=request_message.description or None,
        secret=request_message.secret or None,
        status=WebhookStatus.from_proto(request_message.status) if request_message.status else None,
    )
    response_message = CreateWebhook.Response(webhook=webhook.to_proto())
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _list_webhooks():
    request_message = _get_request_message(
        ListWebhooks(),
        schema={
            "max_results": [_assert_intlike],
            "page_token": [_assert_string],
        },
    )
    webhooks_page = _get_model_registry_store().list_webhooks(
        max_results=request_message.max_results,
        page_token=request_message.page_token or None,
    )
    response_message = ListWebhooks.Response(
        webhooks=[w.to_proto() for w in webhooks_page],
        next_page_token=webhooks_page.token,
    )
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _get_webhook(webhook_id: str):
    webhook = _get_model_registry_store().get_webhook(webhook_id=webhook_id)
    response_message = GetWebhook.Response(webhook=webhook.to_proto())
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _update_webhook(webhook_id: str):
    request_message = _get_request_message(
        UpdateWebhook(),
        schema={
            "name": [_assert_string],
            "description": [_assert_string],
            "url": [_assert_string],
            "events": [_assert_array],
            "secret": [_assert_string],
            "status": [_assert_string],
        },
    )
    webhook = _get_model_registry_store().update_webhook(
        webhook_id=webhook_id,
        name=request_message.name or None,
        description=request_message.description or None,
        url=request_message.url or None,
        events=(
            [WebhookEvent.from_proto(e) for e in request_message.events]
            if request_message.events
            else None
        ),
        secret=request_message.secret or None,
        status=WebhookStatus.from_proto(request_message.status) if request_message.status else None,
    )
    response_message = UpdateWebhook.Response(webhook=webhook.to_proto())
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _delete_webhook(webhook_id: str):
    _get_model_registry_store().delete_webhook(webhook_id=webhook_id)
    response_message = DeleteWebhook.Response()
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _test_webhook(webhook_id: str):
    request_message = _get_request_message(TestWebhook())
    event = (
        WebhookEvent.from_proto(request_message.event)
        if request_message.HasField("event")
        else None
    )
    store = _get_model_registry_store()
    webhook = store.get_webhook(webhook_id=webhook_id)
    test_result = test_webhook(webhook=webhook, event=event)
    response_message = TestWebhook.Response(result=test_result.to_proto())
    return _wrap_response(response_message)


# MLflow Artifacts APIs


def _get_workspace_scoped_repo_path_if_enabled(artifact_path: str | None) -> str | None:
    """
    Normalize artifact paths for proxied (served) artifacts so they remain workspace-isolated.

    When ``mlflow-artifacts`` proxying is enabled and workspaces are on, every path under the HTTP
    artifact endpoint must be rooted at ``workspaces/<workspace>/...``. Direct artifact repositories
    (e.g., S3, GCS, local URIs) already encode their own isolation, so they bypass this logic by
    calling the underlying store directly. Only the proxied repos need to be rewritten/validated
    here.

    Returns:
        The workspace-scoped path. May return the original path in the following cases:
        - Workspaces are disabled (returns ``artifact_path`` unchanged).
        - Default workspace with no path (returns ``artifact_path`` unchanged to preserve legacy
          root behavior, where artifacts live at the root rather than under ``workspaces/default``).
        For non-default workspaces, always returns a string (``workspaces/<workspace>/...``).
    """
    if not MLFLOW_ENABLE_WORKSPACES.get():
        return artifact_path

    workspace = workspace_context.get_request_workspace()
    if not workspace:
        raise MlflowException.invalid_parameter_value(
            "Active workspace is required for artifact operations. "
            "Ensure X-MLFLOW-WORKSPACE is set or call mlflow.set_workspace()."
        )

    normalized = artifact_path.lstrip("/") if artifact_path else ""
    base = posixpath.join("workspaces", workspace)

    if not normalized:
        # For the default workspace, preserve the legacy root behavior (no prefix),
        # so root operations continue to see the existing layout.
        return base if workspace != DEFAULT_WORKSPACE_NAME else artifact_path

    if workspace == DEFAULT_WORKSPACE_NAME and not normalized.startswith("workspaces/"):
        # Legacy default-workspace artifacts never had the workspace prefix; allow them to be served
        # without rewriting as long as the path isn't trying to opt into the reserved namespace.
        return artifact_path

    leading_segments = normalized.split("/", 2)
    if leading_segments and leading_segments[0] == "workspaces":
        if len(leading_segments) == 1 or not leading_segments[1]:
            raise MlflowException.invalid_parameter_value(
                "Artifact paths prefixed with 'workspaces/' must include a workspace name."
            )
        requested_workspace = leading_segments[1]
        if requested_workspace != workspace:
            raise MlflowException.invalid_parameter_value(
                f"Artifact path targets workspace '{requested_workspace}' "
                f"but the workspace specified in the request is '{workspace}'."
            )
        return normalized

    return posixpath.join(base, normalized)


@catch_mlflow_exception
@_disable_unless_serve_artifacts
def _download_artifact(artifact_path):
    """
    A request handler for `GET /mlflow-artifacts/artifacts/<artifact_path>` to download an artifact
    from `artifact_path` (a relative path from the root artifact directory).
    """
    artifact_path = validate_path_is_safe(artifact_path)
    artifact_path = _get_workspace_scoped_repo_path_if_enabled(artifact_path)
    artifact_repo = _get_artifact_repo_mlflow_artifacts()

    if (local_path := artifact_repo.get_local_path(artifact_path)) is not None:
        return _create_artifact_file_response(os.path.abspath(local_path), artifact_path)

    tmp_dir = tempfile.TemporaryDirectory()
    try:
        dst = os.path.abspath(artifact_repo.download_artifacts(artifact_path, tmp_dir.name))

        # Ref: https://stackoverflow.com/a/24613980/6943581
        file_handle = open(dst, "rb")  # noqa: SIM115
    except Exception:
        tmp_dir.cleanup()
        raise

    def stream_and_remove_file():
        while chunk := file_handle.read(ARTIFACT_STREAM_CHUNK_SIZE):
            yield chunk
        file_handle.close()
        tmp_dir.cleanup()

    file_sender_response = current_app.response_class(stream_and_remove_file())

    return _response_with_file_attachment_headers(artifact_path, file_sender_response)


@catch_mlflow_exception
@_disable_unless_serve_artifacts
def _upload_artifact(artifact_path):
    """
    A request handler for `PUT /mlflow-artifacts/artifacts/<artifact_path>` to upload an artifact
    to `artifact_path` (a relative path from the root artifact directory).
    """
    artifact_path = validate_path_is_safe(artifact_path)
    artifact_path = _get_workspace_scoped_repo_path_if_enabled(artifact_path)
    head, tail = posixpath.split(artifact_path)
    artifact_repo = _get_artifact_repo_mlflow_artifacts()

    if isinstance(artifact_repo, StreamUploadMixin):
        artifact_repo.log_artifact_from_stream(request.stream, tail, artifact_path=head or None)
    else:
        with tempfile.TemporaryDirectory() as tmp_dir:
            tmp_path = os.path.join(tmp_dir, tail)
            with open(tmp_path, "wb") as f:
                while chunk := request.stream.read(ARTIFACT_STREAM_CHUNK_SIZE):
                    f.write(chunk)
            artifact_repo.log_artifact(tmp_path, artifact_path=head or None)

    return _wrap_response(UploadArtifact.Response())


@catch_mlflow_exception
@_disable_unless_serve_artifacts
def _list_artifacts_mlflow_artifacts():
    """
    A request handler for `GET /mlflow-artifacts/artifacts?path=<value>` to list artifacts in `path`
    (a relative path from the root artifact directory).
    """
    request_message = _get_request_message(ListArtifactsMlflowArtifacts())
    path = validate_path_is_safe(request_message.path) if request_message.HasField("path") else None
    path = _get_workspace_scoped_repo_path_if_enabled(path)
    artifact_repo = _get_artifact_repo_mlflow_artifacts()
    files = []
    for file_info in artifact_repo.list_artifacts(path):
        basename = posixpath.basename(file_info.path)
        new_file_info = FileInfo(basename, file_info.is_dir, file_info.file_size)
        files.append(new_file_info.to_proto())
    response_message = ListArtifacts.Response()
    response_message.files.extend(files)
    response = Response(mimetype="application/json")
    response.set_data(message_to_json(response_message))
    return response


@catch_mlflow_exception
@_disable_unless_serve_artifacts
def _delete_artifact_mlflow_artifacts(artifact_path):
    """
    A request handler for `DELETE /mlflow-artifacts/artifacts?path=<value>` to delete artifacts in
    `path` (a relative path from the root artifact directory).
    """
    artifact_path = validate_path_is_safe(artifact_path)
    artifact_path = _get_workspace_scoped_repo_path_if_enabled(artifact_path)
    _get_request_message(DeleteArtifact())
    artifact_repo = _get_artifact_repo_mlflow_artifacts()
    artifact_repo.delete_artifacts(artifact_path)
    response_message = DeleteArtifact.Response()
    response = Response(mimetype="application/json")
    response.set_data(message_to_json(response_message))
    return response


def _get_graphql_auth_middleware():
    """
    Get GraphQL authorization middleware if basic-auth is enabled.

    Returns:
        A list of middleware instances if auth is enabled, empty list otherwise.
    """
    try:
        from mlflow.server.auth import get_graphql_authorization_middleware

        return get_graphql_authorization_middleware()
    except Exception:
        # Auth not configured or other error
        return []


@catch_mlflow_exception
def _graphql():
    from graphql import parse

    from mlflow.server.graphql.graphql_no_batching import check_query_safety
    from mlflow.server.graphql.graphql_schema_extensions import schema

    # Extracting the query, variables, and operationName from the request
    request_json = _get_request_json()
    query = request_json.get("query")
    variables = request_json.get("variables")
    operation_name = request_json.get("operationName")

    node = parse(query)
    if check_result := check_query_safety(node):
        result = check_result
    else:
        # Get auth middleware if basic-auth is enabled
        middleware = _get_graphql_auth_middleware()

        # Executing the GraphQL query using the Graphene schema
        result = schema.execute(
            query,
            variables=variables,
            operation_name=operation_name,
            middleware=middleware,
        )

    # Convert execution result into json.
    result_data = {
        "data": result.data,
        "errors": [error.message for error in result.errors] if result.errors else None,
    }

    # Return the response
    return jsonify(result_data)


def _validate_support_multipart_upload(artifact_repo):
    if not isinstance(artifact_repo, MultipartUploadMixin):
        raise _UnsupportedMultipartUploadException()


def _validate_support_multipart_download(artifact_repo):
    if not isinstance(artifact_repo, MultipartDownloadMixin):
        raise _UnsupportedMultipartDownloadException()


def _validate_support_presigned_upload(artifact_repo):
    if not isinstance(artifact_repo, PresignedUploadMixin):
        raise _UnsupportedPresignedUploadException()


@catch_mlflow_exception
@_disable_if_artifacts_only
def _create_presigned_upload_url():
    """
    Handler for POST /api/2.0/mlflow/artifacts/presigned-upload-url.
    Generates a presigned URL for uploading an artifact directly to cloud storage.

    Client reference: https://github.com/aws/sagemaker-mlflow
    """
    request_message = _get_request_message(
        CreatePresignedUploadUrl(),
        schema={
            "run_id": [_assert_required, _assert_string],
            "path": [_assert_required, _assert_string],
            "expiration": [_assert_intlike],
        },
    )
    run_id = request_message.run_id
    path = validate_path_is_safe(request_message.path)
    expiration = request_message.expiration if request_message.HasField("expiration") else 900

    run = _get_tracking_store().get_run(run_id)
    artifact_uri = run.info.artifact_uri
    artifact_uri_scheme = urllib.parse.urlparse(artifact_uri).scheme
    if artifact_uri_scheme in ("http", "https", "mlflow-artifacts"):
        raise MlflowException(
            "Presigned upload is not supported for runs with proxied artifact storage "
            f"(artifact URI scheme: {artifact_uri_scheme}). "
            "This endpoint requires a run with a direct cloud storage artifact URI.",
            error_code=INVALID_PARAMETER_VALUE,
        )
    artifact_repo = _get_artifact_repo(run)
    _validate_support_presigned_upload(artifact_repo)

    response = artifact_repo.create_presigned_upload_url(path, expiration=expiration)
    response_message = response.to_proto()
    resp = Response(mimetype="application/json")
    resp.set_data(message_to_json(response_message))
    return resp


@catch_mlflow_exception
@_disable_unless_serve_artifacts
def _create_multipart_upload_artifact(artifact_path):
    """
    A request handler for `POST /mlflow-artifacts/mpu/create` to create a multipart upload
    to `artifact_path` (a relative path from the root artifact directory).
    """
    artifact_path = validate_path_is_safe(artifact_path)
    artifact_path = _get_workspace_scoped_repo_path_if_enabled(artifact_path)

    request_message = _get_request_message(
        CreateMultipartUpload(),
        schema={
            "path": [_assert_required, _assert_string],
            "num_parts": [_assert_intlike],
        },
    )
    path = request_message.path
    num_parts = request_message.num_parts

    artifact_repo = _get_artifact_repo_mlflow_artifacts()
    _validate_support_multipart_upload(artifact_repo)

    create_response = artifact_repo.create_multipart_upload(
        path,
        num_parts,
        artifact_path,
    )
    response_message = create_response.to_proto()
    response = Response(mimetype="application/json")
    response.set_data(message_to_json(response_message))
    return response


@catch_mlflow_exception
@_disable_unless_serve_artifacts
def _complete_multipart_upload_artifact(artifact_path):
    """
    A request handler for `POST /mlflow-artifacts/mpu/complete` to complete a multipart upload
    to `artifact_path` (a relative path from the root artifact directory).
    """
    artifact_path = validate_path_is_safe(artifact_path)
    artifact_path = _get_workspace_scoped_repo_path_if_enabled(artifact_path)

    request_message = _get_request_message(
        CompleteMultipartUpload(),
        schema={
            "path": [_assert_required, _assert_string],
            "upload_id": [_assert_string],
            "parts": [_assert_required],
        },
    )
    path = request_message.path
    upload_id = request_message.upload_id
    parts = [MultipartUploadPart.from_proto(part) for part in request_message.parts]

    artifact_repo = _get_artifact_repo_mlflow_artifacts()
    _validate_support_multipart_upload(artifact_repo)

    artifact_repo.complete_multipart_upload(
        path,
        upload_id,
        parts,
        artifact_path,
    )
    return _wrap_response(CompleteMultipartUpload.Response())


@catch_mlflow_exception
@_disable_unless_serve_artifacts
def _abort_multipart_upload_artifact(artifact_path):
    """
    A request handler for `POST /mlflow-artifacts/mpu/abort` to abort a multipart upload
    to `artifact_path` (a relative path from the root artifact directory).
    """
    artifact_path = validate_path_is_safe(artifact_path)
    artifact_path = _get_workspace_scoped_repo_path_if_enabled(artifact_path)

    request_message = _get_request_message(
        AbortMultipartUpload(),
        schema={
            "path": [_assert_required, _assert_string],
            "upload_id": [_assert_string],
        },
    )
    path = request_message.path
    upload_id = request_message.upload_id

    artifact_repo = _get_artifact_repo_mlflow_artifacts()
    _validate_support_multipart_upload(artifact_repo)

    artifact_repo.abort_multipart_upload(
        path,
        upload_id,
        artifact_path,
    )
    return _wrap_response(AbortMultipartUpload.Response())


@catch_mlflow_exception
@_disable_unless_serve_artifacts
def _get_presigned_download_url(artifact_path):
    """
    A request handler for `GET /mlflow-artifacts/presigned/<artifact_path>` to get
    a presigned URL for downloading an artifact directly from cloud storage.
    """
    artifact_path = validate_path_is_safe(artifact_path)

    artifact_repo = _get_artifact_repo_mlflow_artifacts()
    _validate_support_multipart_download(artifact_repo)

    expiration = MLFLOW_PRESIGNED_DOWNLOAD_URL_TTL_SECONDS.get()
    presigned_response = artifact_repo.get_download_presigned_url(
        artifact_path, expiration=expiration
    )
    response = Response(mimetype="application/json")
    response.set_data(json.dumps(presigned_response.to_dict()))
    return response


# MLflow Tracing APIs


@catch_mlflow_exception
@_disable_if_artifacts_only
def _start_trace_v3():
    """
    A request handler for `POST /mlflow/traces` to create a new TraceInfo record in tracking store.
    """
    request_message = _get_request_message(
        StartTraceV3(),
        schema={"trace": [_assert_required]},
    )
    trace_info = TraceInfo.from_proto(request_message.trace.trace_info)
    trace_info = _get_tracking_store().start_trace(trace_info)
    response_message = StartTraceV3.Response(trace=ProtoTrace(trace_info=trace_info.to_proto()))
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _get_trace_info_v3(trace_id):
    """
    A request handler for `GET /mlflow/traces/{trace_id}` to retrieve
    an existing TraceInfo record from tracking store.
    """
    trace_info = _get_tracking_store().get_trace_info(trace_id)
    response_message = GetTraceInfoV3.Response(trace=ProtoTrace(trace_info=trace_info.to_proto()))
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _batch_get_traces() -> Response:
    """
    A request handler for `GET /mlflow/traces/batchGet` to retrieve
    a batch of complete traces with spans for given trace ids.
    """
    request_message = _get_request_message(
        BatchGetTraces(),
        schema={"trace_ids": [_assert_array, _assert_required, _assert_item_type_string]},
    )
    traces = _get_tracking_store().batch_get_traces(request_message.trace_ids, None)
    response_message = BatchGetTraces.Response()
    response_message.traces.extend([t.to_proto() for t in traces])
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _batch_get_trace_infos() -> Response:
    request_message = _get_request_message(
        BatchGetTraceInfos(),
        schema={"trace_ids": [_assert_array, _assert_required, _assert_item_type_string]},
    )
    trace_infos = _get_tracking_store().batch_get_trace_infos(request_message.trace_ids)
    response_message = BatchGetTraceInfos.Response()
    response_message.trace_infos.extend([ti.to_proto() for ti in trace_infos])
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _get_trace() -> Response:
    """
    A request handler for `GET /mlflow/traces/get` to get a trace with spans for given trace id.
    """
    request_message = _get_request_message(
        GetTrace(),
        schema={
            "trace_id": [_assert_string, _assert_required],
            "allow_partial": [_assert_bool],
        },
    )
    trace_id = request_message.trace_id
    allow_partial = request_message.allow_partial
    trace = _get_tracking_store().get_trace(trace_id, allow_partial=allow_partial)
    response_message = GetTrace.Response(trace=trace.to_proto())
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _search_traces_v3():
    """
    A request handler for `GET /mlflow/traces` to search for TraceInfo records in tracking store.
    """
    request_message = _get_request_message(
        SearchTracesV3(),
        schema={
            "locations": [_assert_array, _assert_required],
            "filter": [_assert_string],
            "max_results": [
                _assert_intlike,
                lambda x: _assert_less_than_or_equal(int(x), 500),
            ],
            "order_by": [_assert_array, _assert_item_type_string],
            "page_token": [_assert_string],
        },
    )
    experiment_ids = [
        location.mlflow_experiment.experiment_id
        for location in request_message.locations
        if location.HasField("mlflow_experiment")
    ]

    traces, token = _get_tracking_store().search_traces(
        locations=experiment_ids,
        filter_string=request_message.filter,
        max_results=request_message.max_results,
        order_by=request_message.order_by,
        page_token=request_message.page_token or None,
    )
    response_message = SearchTracesV3.Response()
    response_message.traces.extend([e.to_proto() for e in traces])
    if token:
        response_message.next_page_token = token
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _delete_traces():
    """
    A request handler for `POST /mlflow/traces/delete-traces` to delete TraceInfo records
    from tracking store.
    """
    request_message = _get_request_message(
        DeleteTraces(),
        schema={
            "experiment_id": [_assert_string, _assert_required],
            "max_timestamp_millis": [_assert_intlike],
            "max_traces": [_assert_intlike],
            "request_ids": [_assert_array, _assert_item_type_string],
        },
    )

    # NB: Interestingly, the field accessor for the message object returns the default
    #   value for optional field if it's not set. For example, `request_message.max_traces`
    #   returns 0 if max_traces is not specified in the request. This is not desirable,
    #   because null and 0 means completely opposite i.e. the former is 'delete nothing'
    #   while the latter is 'delete all'. To handle this, we need to explicitly check
    #   if the field is set or not using `HasField` method and return None if not.
    def _get_nullable_field(field):
        if request_message.HasField(field):
            return getattr(request_message, field)
        return None

    traces_deleted = _get_tracking_store().delete_traces(
        experiment_id=request_message.experiment_id,
        max_timestamp_millis=_get_nullable_field("max_timestamp_millis"),
        max_traces=_get_nullable_field("max_traces"),
        trace_ids=request_message.request_ids,
    )
    return _wrap_response(DeleteTraces.Response(traces_deleted=traces_deleted))


@catch_mlflow_exception
@_disable_if_artifacts_only
def _calculate_trace_filter_correlation():
    """
    A request handler for `POST /mlflow/traces/calculate-filter-correlation` to calculate
    NPMI correlation between two trace filter conditions.
    """
    request_message = _get_request_message(
        CalculateTraceFilterCorrelation(),
        schema={
            "experiment_ids": [_assert_array, _assert_required, _assert_item_type_string],
            "filter_string1": [_assert_string, _assert_required],
            "filter_string2": [_assert_string, _assert_required],
            "base_filter": [_assert_string],
        },
    )

    result = _get_tracking_store().calculate_trace_filter_correlation(
        experiment_ids=request_message.experiment_ids,
        filter_string1=request_message.filter_string1,
        filter_string2=request_message.filter_string2,
        base_filter=request_message.base_filter
        if request_message.HasField("base_filter")
        else None,
    )

    return _wrap_response(result.to_proto())


@catch_mlflow_exception
@_disable_if_artifacts_only
def _set_trace_tag(request_id):
    """
    A request handler for `PATCH /mlflow/traces/{request_id}/tags` to set tags on a TraceInfo record
    """
    request_message = _get_request_message(
        SetTraceTag(),
        schema={
            "key": [_assert_string, _assert_required],
            "value": [_assert_string],
        },
    )
    _validate_trace_tag_handler_mutation(request_message.key, "set")
    _get_tracking_store().set_trace_tag(request_id, request_message.key, request_message.value)
    return _wrap_response(SetTraceTag.Response())


@catch_mlflow_exception
@_disable_if_artifacts_only
def _set_trace_tag_v3(trace_id):
    """
    A request handler for `PATCH /mlflow/traces/{trace_id}/tags` to set tags on a TraceInfo record.
    Identical to `_set_trace_tag`, but with request_id renamed to with trace_id.
    """
    request_message = _get_request_message(
        SetTraceTagV3(),
        schema={
            "key": [_assert_string, _assert_required],
            "value": [_assert_string],
        },
    )
    _validate_trace_tag_handler_mutation(request_message.key, "set")
    _get_tracking_store().set_trace_tag(trace_id, request_message.key, request_message.value)
    return _wrap_response(SetTraceTagV3.Response())


@catch_mlflow_exception
@_disable_if_artifacts_only
def _delete_trace_tag(request_id):
    """
    A request handler for `DELETE /mlflow/traces/{request_id}/tags` to delete tags from a TraceInfo
    record.
    """
    request_message = _get_request_message(
        DeleteTraceTag(),
        schema={
            "key": [_assert_string, _assert_required],
        },
    )
    _validate_trace_tag_handler_mutation(request_message.key, "deleted")
    _get_tracking_store().delete_trace_tag(request_id, request_message.key)
    return _wrap_response(DeleteTraceTag.Response())


@catch_mlflow_exception
@_disable_if_artifacts_only
def _delete_trace_tag_v3(trace_id):
    """
    A request handler for `DELETE /mlflow/traces/{trace_id}/tags` to delete tags
    from a TraceInfo record.
    Identical to `_delete_trace_tag`, but with request_id renamed to with trace_id.
    """
    request_message = _get_request_message(
        DeleteTraceTagV3(),
        schema={
            "key": [_assert_string, _assert_required],
        },
    )
    _validate_trace_tag_handler_mutation(request_message.key, "deleted")
    _get_tracking_store().delete_trace_tag(trace_id, request_message.key)
    return _wrap_response(DeleteTraceTagV3.Response())


@catch_mlflow_exception
@_disable_if_artifacts_only
def _link_traces_to_run():
    """
    A request handler for `POST /mlflow/traces/link-to-run` to link traces to a run.
    """
    request_message = _get_request_message(
        LinkTracesToRun(),
        schema={
            "trace_ids": [_assert_array, _assert_required, _assert_item_type_string],
            "run_id": [_assert_string, _assert_required],
        },
    )
    _get_tracking_store().link_traces_to_run(
        trace_ids=request_message.trace_ids,
        run_id=request_message.run_id,
    )
    return _wrap_response(LinkTracesToRun.Response())


@catch_mlflow_exception
@_disable_if_artifacts_only
def _link_prompts_to_trace():
    """
    A request handler for `POST /mlflow/traces/link-prompts` to link prompt versions to a trace.
    """
    from mlflow.entities.model_registry import PromptVersion

    request_message = _get_request_message(
        LinkPromptsToTrace(),
        schema={
            "trace_id": [_assert_string, _assert_required],
            "prompt_versions": [_assert_array, _assert_required],
        },
    )

    # Convert PromptVersionRef proto messages to PromptVersion objects
    # It doesn't load prompt versions since name and version are sufficient for linking
    prompt_versions = [
        PromptVersion(name=pv.name, version=int(pv.version), template="")
        for pv in request_message.prompt_versions
    ]

    _get_tracking_store().link_prompts_to_trace(
        trace_id=request_message.trace_id,
        prompt_versions=prompt_versions,
    )
    return _wrap_response(LinkPromptsToTrace.Response())


def _fetch_trace_data_from_store(
    store: AbstractTrackingStore, request_id: str
) -> dict[str, Any] | None:
    try:
        # allow partial so the frontend can render in-progress traces
        trace = store.get_trace(request_id, allow_partial=True)
        return trace.data.to_dict()
    except MlflowTraceDataException:
        raise
    except MlflowTracingException:
        return None
    except MlflowNotImplementedException:
        # fallback to batch_get_traces if get_trace is not implemented
        pass

    try:
        traces = store.batch_get_traces([request_id], None)
        match traces:
            case [trace]:
                return trace.data.to_dict()
            case _:
                raise MlflowException(
                    f"Trace with id={request_id} not found.",
                    error_code=RESOURCE_DOES_NOT_EXIST,
                )
    except MlflowTraceDataException:
        raise
    # For stores that don't support batch get traces, or if trace data is not in the store,
    # return None to signal fallback to artifact repository
    except (MlflowTracingException, MlflowNotImplementedException):
        return None


@catch_mlflow_exception
@_disable_if_artifacts_only
def get_trace_artifact_handler() -> Response:
    request_id = request.args.get("request_id")
    path = request.args.get("path")

    if not request_id:
        raise MlflowException(
            'Request must include the "request_id" query parameter.',
            error_code=BAD_REQUEST,
        )

    store = _get_tracking_store()

    if path:
        path = validate_path_is_safe(path)
        trace_info = store.get_trace_info(request_id)
        if trace_info is None:
            raise MlflowException(
                f"Trace with ID '{request_id}' not found.",
                error_code=RESOURCE_DOES_NOT_EXIST,
            )
        repo = _get_trace_artifact_repo(trace_info)
        try:
            content_bytes = repo.download_trace_attachment(path)
        except MlflowException:
            raise
        except Exception:
            _logger.warning(
                "Failed to download attachment '%s' for trace '%s'",
                path,
                request_id,
                exc_info=True,
            )
            raise MlflowException(
                f"Failed to download attachment '{path}' for trace '{request_id}'.",
                error_code=INTERNAL_ERROR,
            )
        buf = io.BytesIO(content_bytes)
        file_sender_response = send_file(
            buf,
            mimetype="application/octet-stream",
            as_attachment=True,
            download_name=path,
        )
        return _response_with_file_attachment_headers(path, file_sender_response)

    trace_data = _fetch_trace_data_from_store(store, request_id)
    if trace_data is None:
        trace_info = store.get_trace_info(request_id)
        if trace_info.tags.get(TraceTagKey.SPANS_LOCATION) == SpansLocation.ARCHIVE_REPO.value:
            trace_data = (
                _get_trace_archive_repo(trace_info).download_archived_trace_data().to_dict()
            )
        else:
            trace_data = _get_trace_artifact_repo(trace_info).download_trace_data()

    # Write data to a BytesIO buffer instead of needing to save a temp file
    buf = io.BytesIO()
    buf.write(json.dumps(trace_data).encode())
    buf.seek(0)

    file_sender_response = send_file(
        buf,
        mimetype="application/octet-stream",
        as_attachment=True,
        download_name=TRACE_DATA_FILE_NAME,
    )
    return _response_with_file_attachment_headers(TRACE_DATA_FILE_NAME, file_sender_response)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _query_trace_metrics() -> Response:
    request_message = _get_request_message(
        QueryTraceMetrics(),
        schema={
            "experiment_ids": [_assert_array, _assert_required, _assert_item_type_string],
            "view_type": [_assert_required],
            "metric_name": [_assert_string, _assert_required],
            "aggregations": [_assert_array, _assert_required],
            "dimensions": [_assert_array, _assert_item_type_string],
            "filters": [_assert_array, _assert_item_type_string],
            "time_interval_seconds": [_assert_intlike],
            "start_time_ms": [_assert_intlike],
            "end_time_ms": [_assert_intlike],
            "max_results": [_assert_intlike],
            "page_token": [_assert_string],
        },
    )
    max_results = (
        request_message.max_results
        if request_message.HasField("max_results")
        else MAX_RESULTS_QUERY_TRACE_METRICS
    )
    time_interval_seconds = (
        request_message.time_interval_seconds
        if request_message.HasField("time_interval_seconds")
        else None
    )
    start_time_ms = (
        request_message.start_time_ms if request_message.HasField("start_time_ms") else None
    )
    end_time_ms = request_message.end_time_ms if request_message.HasField("end_time_ms") else None

    result = _get_tracking_store().query_trace_metrics(
        experiment_ids=request_message.experiment_ids,
        view_type=MetricViewType.from_proto(request_message.view_type),
        metric_name=request_message.metric_name,
        aggregations=[MetricAggregation.from_proto(agg) for agg in request_message.aggregations],
        dimensions=request_message.dimensions or None,
        filters=request_message.filters or None,
        time_interval_seconds=time_interval_seconds,
        start_time_ms=start_time_ms,
        end_time_ms=end_time_ms,
        max_results=max_results,
        page_token=request_message.page_token or None,
    )
    response_message = QueryTraceMetrics.Response()
    response_message.data_points.extend([dp.to_proto() for dp in result])
    if result.token:
        response_message.next_page_token = result.token
    return _wrap_response(response_message)


# Assessments API handlers
@catch_mlflow_exception
@_disable_if_artifacts_only
def _create_assessment(trace_id):
    """
    A request handler for `POST /mlflow/traces/{assessment.trace_id}/assessments`
    to create a new assessment.
    """
    request_message = _get_request_message(
        CreateAssessment(),
        schema={
            "assessment": [_assert_required],
        },
    )

    assessment = Assessment.from_proto(request_message.assessment)
    assessment.trace_id = trace_id
    created_assessment = _get_tracking_store().create_assessment(assessment)

    response_message = CreateAssessment.Response(assessment=created_assessment.to_proto())
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _get_assessment(trace_id, assessment_id):
    """
    A request handler for `GET /mlflow/traces/{trace_id}/assessments/{assessment_id}`
    to get an assessment.
    """
    assessment = _get_tracking_store().get_assessment(trace_id, assessment_id)

    response_message = GetAssessmentRequest.Response(assessment=assessment.to_proto())
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _update_assessment(trace_id, assessment_id):
    """
    A request handler for `PATCH /mlflow/traces/{trace_id}/assessments/{assessment_id}`
    to update an assessment.
    """
    request_message = _get_request_message(
        UpdateAssessment(),
        schema={
            "assessment": [_assert_required],
            "update_mask": [_assert_required],
        },
    )

    assessment_proto = request_message.assessment
    update_mask = request_message.update_mask

    kwargs = {}

    for path in update_mask.paths:
        if path == "assessment_name":
            kwargs["name"] = assessment_proto.assessment_name
        elif path == "expectation":
            kwargs["expectation"] = Expectation.from_proto(assessment_proto)
        elif path == "feedback":
            kwargs["feedback"] = Feedback.from_proto(assessment_proto)
        elif path == "rationale":
            kwargs["rationale"] = assessment_proto.rationale
        elif path == "metadata":
            kwargs["metadata"] = dict(assessment_proto.metadata)
        elif path == "valid":
            kwargs["valid"] = assessment_proto.valid

    updated_assessment = _get_tracking_store().update_assessment(
        trace_id=trace_id, assessment_id=assessment_id, **kwargs
    )

    response_message = UpdateAssessment.Response(assessment=updated_assessment.to_proto())
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _delete_assessment(trace_id, assessment_id):
    """
    A request handler for `DELETE /mlflow/traces/{trace_id}/assessments/{assessment_id}`
    to delete an assessment.
    """
    _get_tracking_store().delete_assessment(trace_id, assessment_id)

    response_message = DeleteAssessment.Response()
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _create_issue():
    """
    A request handler for `POST /mlflow/issues` to create a new issue.
    """
    request_message = _get_request_message(
        CreateIssue(),
        schema={
            "name": [_assert_required, _assert_string],
            "description": [_assert_required, _assert_string],
            "experiment_id": [_assert_required, _assert_string],
        },
    )

    # Build kwargs for create_issue
    create_kwargs = {
        "experiment_id": request_message.experiment_id,
        "name": request_message.name,
        "description": request_message.description,
        "source_run_id": request_message.source_run_id or None,
        "root_causes": list(request_message.root_causes) or None,
        "categories": list(request_message.categories) or None,
        "created_by": request_message.created_by or None,
    }

    if request_message.HasField("status"):
        create_kwargs["status"] = IssueStatus(request_message.status)
    if request_message.HasField("severity"):
        create_kwargs["severity"] = IssueSeverity(request_message.severity)

    created_issue = _get_tracking_store().create_issue(**create_kwargs)

    response_message = CreateIssue.Response(issue=created_issue.to_proto())
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _update_issue(issue_id):
    """
    A request handler for `PATCH /mlflow/issues/{issue_id}` to update an issue.
    """
    request_message = _get_request_message(
        UpdateIssue(),
        schema={
            "issue_id": [_assert_required],
        },
    )

    status = IssueStatus(request_message.status) if request_message.HasField("status") else None
    severity = (
        IssueSeverity(request_message.severity) if request_message.HasField("severity") else None
    )

    updated_issue = _get_tracking_store().update_issue(
        issue_id=issue_id,
        status=status,
        name=request_message.name or None,
        description=request_message.description or None,
        severity=severity,
    )

    response_message = UpdateIssue.Response(issue=updated_issue.to_proto())
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _get_issue(issue_id):
    """
    A request handler for `GET /mlflow/issues/{issue_id}` to get an issue.
    """
    issue = _get_tracking_store().get_issue(issue_id)

    response_message = GetIssue.Response(issue=issue.to_proto())
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _search_issues():
    """
    A request handler for `POST /mlflow/issues/search` to search for issues.
    """
    request_message = _get_request_message(SearchIssues())

    # Build kwargs for search_issues
    search_kwargs = {
        "experiment_id": request_message.experiment_id or None,
        "filter_string": request_message.filter_string or None,
        "page_token": request_message.page_token or None,
    }

    if request_message.HasField("max_results"):
        search_kwargs["max_results"] = request_message.max_results

    if request_message.HasField("include_trace_count"):
        search_kwargs["include_trace_count"] = request_message.include_trace_count

    issues = _get_tracking_store().search_issues(**search_kwargs)

    issue_protos = [issue.to_proto() for issue in issues]
    response_message = SearchIssues.Response(
        issues=issue_protos, next_page_token=issues.token or ""
    )
    return _wrap_response(response_message)


# =============================================================================
# Label Schema Handlers (tracking-store CRUD; see mlflow/genai/label_schemas/)
# =============================================================================


@catch_mlflow_exception
@_disable_if_artifacts_only
def _create_label_schema():
    request_message = _get_request_message(
        CreateLabelSchema(),
        schema={
            "experiment_id": [_assert_required, _assert_string],
            "name": [_assert_required, _assert_string],
        },
    )
    schema_type = LabelSchemaType.from_proto(request_message.type)
    input_obj = _input_from_proto(request_message.input)
    kwargs: dict[str, object] = {
        "experiment_id": request_message.experiment_id,
        "name": request_message.name,
        "type": schema_type,
        "input": input_obj,
    }
    if request_message.HasField("instruction"):
        kwargs["instruction"] = request_message.instruction
    if request_message.HasField("enable_comment"):
        kwargs["enable_comment"] = request_message.enable_comment
    created = _get_tracking_store().create_label_schema(**kwargs)
    return _wrap_response(CreateLabelSchema.Response(label_schema=created.to_proto()))


@catch_mlflow_exception
@_disable_if_artifacts_only
def _get_label_schema():
    request_message = _get_request_message(
        GetLabelSchema(),
        schema={"schema_id": [_assert_required, _assert_string]},
    )
    schema = _get_tracking_store().get_label_schema(request_message.schema_id)
    return _wrap_response(GetLabelSchema.Response(label_schema=schema.to_proto()))


@catch_mlflow_exception
@_disable_if_artifacts_only
def _get_label_schema_by_name():
    request_message = _get_request_message(
        GetLabelSchemaByName(),
        schema={
            "experiment_id": [_assert_required, _assert_string],
            "name": [_assert_required, _assert_string],
        },
    )
    schema = _get_tracking_store().get_label_schema_by_name(
        request_message.experiment_id, request_message.name
    )
    return _wrap_response(GetLabelSchemaByName.Response(label_schema=schema.to_proto()))


@catch_mlflow_exception
@_disable_if_artifacts_only
def _list_label_schemas():
    request_message = _get_request_message(
        ListLabelSchemas(),
        schema={
            "experiment_id": [_assert_required, _assert_string],
            "max_results": [
                _assert_intlike,
                lambda x: _assert_intlike_within_range(
                    int(x),
                    1,
                    SEARCH_MAX_RESULTS_THRESHOLD,
                    message=(f"max_results must be between 1 and {SEARCH_MAX_RESULTS_THRESHOLD}."),
                ),
            ],
        },
    )
    max_results = request_message.max_results if request_message.HasField("max_results") else 100
    page_token = request_message.page_token if request_message.HasField("page_token") else None
    schemas = _get_tracking_store().list_label_schemas(
        request_message.experiment_id, max_results=max_results, page_token=page_token
    )
    response = ListLabelSchemas.Response(
        label_schemas=[s.to_proto() for s in schemas],
        next_page_token=schemas.token or "",
    )
    return _wrap_response(response)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _update_label_schema():
    request_message = _get_request_message(
        UpdateLabelSchema(),
        schema={"schema_id": [_assert_required, _assert_string]},
    )
    kwargs: dict[str, object] = {}
    if request_message.HasField("name"):
        kwargs["name"] = request_message.name
    if request_message.HasField("instruction"):
        kwargs["instruction"] = request_message.instruction
    if request_message.HasField("enable_comment"):
        kwargs["enable_comment"] = request_message.enable_comment
    if request_message.HasField("input"):
        kwargs["input"] = _input_from_proto(request_message.input)
    updated = _get_tracking_store().update_label_schema(request_message.schema_id, **kwargs)
    return _wrap_response(UpdateLabelSchema.Response(label_schema=updated.to_proto()))


@catch_mlflow_exception
@_disable_if_artifacts_only
def _delete_label_schema():
    request_message = _get_request_message(
        DeleteLabelSchema(),
        schema={"schema_id": [_assert_required, _assert_string]},
    )
    _get_tracking_store().delete_label_schema(request_message.schema_id)
    return _wrap_response(DeleteLabelSchema.Response())


# =============================================================================
# Review Queue Handlers (tracking-store CRUD; see mlflow/genai/review_queues/)
# =============================================================================


def _review_queue_max_results_validator(x):
    return _assert_intlike_within_range(
        int(x),
        1,
        SEARCH_MAX_RESULTS_THRESHOLD,
        message=f"max_results must be between 1 and {SEARCH_MAX_RESULTS_THRESHOLD}.",
    )


def _get_request_username():
    """The authenticated request user, stamped on ``flask.g`` by the auth
    plugin's before-request hook. ``None`` when no auth plugin is active (a
    no-auth server), where queue ownership is meaningless.
    """
    return getattr(g, "mlflow_authenticated_user", None)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _create_review_queue():
    request_message = _get_request_message(
        CreateReviewQueue(),
        schema={
            "experiment_id": [_assert_required, _assert_string],
            "name": [_assert_required, _assert_string],
        },
    )
    # `from_proto` rejects the proto2 zero-value (UNSPECIFIED); don't replace
    # it with an `_assert_required` schema entry — that only checks HasField,
    # not enum-value validity, and would change the rejection's error shape.
    kwargs: dict[str, object] = {
        "experiment_id": request_message.experiment_id,
        "name": request_message.name,
        "queue_type": ReviewQueueType.from_proto(request_message.queue_type),
        "users": list(request_message.users),
        "schema_ids": list(request_message.schema_ids),
    }
    # `created_by` is the owner: stamp it from the authenticated user, never the
    # client. Stays unset on a no-auth server (owner is meaningless there).
    username = _get_request_username()
    if username is not None:
        kwargs["created_by"] = username
    created = _get_tracking_store().create_review_queue(**kwargs)
    return _wrap_response(CreateReviewQueue.Response(review_queue=created.to_proto()))


@catch_mlflow_exception
@_disable_if_artifacts_only
def _get_or_create_user_queue():
    request_message = _get_request_message(
        GetOrCreateUserQueue(),
        schema={
            "experiment_id": [_assert_required, _assert_string],
            "user": [_assert_required, _assert_string],
        },
    )
    # A user queue is owned by its user (set in the store), not by any client value.
    queue = _get_tracking_store().get_or_create_user_queue(
        request_message.experiment_id, user=request_message.user
    )
    return _wrap_response(GetOrCreateUserQueue.Response(review_queue=queue.to_proto()))


@catch_mlflow_exception
@_disable_if_artifacts_only
def _get_review_queue():
    request_message = _get_request_message(
        GetReviewQueue(),
        schema={"queue_id": [_assert_required, _assert_string]},
    )
    queue = _get_tracking_store().get_review_queue(request_message.queue_id)
    return _wrap_response(GetReviewQueue.Response(review_queue=queue.to_proto()))


@catch_mlflow_exception
@_disable_if_artifacts_only
def _get_review_queue_by_name():
    request_message = _get_request_message(
        GetReviewQueueByName(),
        schema={
            "experiment_id": [_assert_required, _assert_string],
            "name": [_assert_required, _assert_string],
        },
    )
    queue = _get_tracking_store().get_review_queue_by_name(
        request_message.experiment_id, name=request_message.name
    )
    return _wrap_response(GetReviewQueueByName.Response(review_queue=queue.to_proto()))


@catch_mlflow_exception
@_disable_if_artifacts_only
def _list_review_queues():
    request_message = _get_request_message(
        ListReviewQueues(),
        schema={
            "experiment_id": [_assert_required, _assert_string],
            "max_results": [_assert_intlike, _review_queue_max_results_validator],
        },
    )
    max_results = request_message.max_results if request_message.HasField("max_results") else None
    page_token = request_message.page_token if request_message.HasField("page_token") else None
    user = request_message.user if request_message.HasField("user") else None
    item_id = request_message.item_id if request_message.HasField("item_id") else None
    queues = _get_tracking_store().list_review_queues(
        request_message.experiment_id,
        user=user,
        item_id=item_id,
        max_results=max_results,
        page_token=page_token,
    )
    response = ListReviewQueues.Response(
        review_queues=[q.to_proto() for q in queues],
        next_page_token=queues.token or "",
    )
    return _wrap_response(response)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _update_review_queue():
    request_message = _get_request_message(
        UpdateReviewQueue(),
        schema={"queue_id": [_assert_required, _assert_string]},
    )
    users = list(request_message.users) if request_message.update_users else None
    schema_ids = list(request_message.schema_ids) if request_message.update_schema_ids else None
    # Singular fields use proto2 presence (no update_* flag).
    name = request_message.name if request_message.HasField("name") else None
    new_owner = request_message.new_owner if request_message.HasField("new_owner") else None
    updated = _get_tracking_store().update_review_queue(
        request_message.queue_id,
        users=users,
        schema_ids=schema_ids,
        name=name,
        new_owner=new_owner,
    )
    return _wrap_response(UpdateReviewQueue.Response(review_queue=updated.to_proto()))


@catch_mlflow_exception
@_disable_if_artifacts_only
def _delete_review_queue():
    request_message = _get_request_message(
        DeleteReviewQueue(),
        schema={"queue_id": [_assert_required, _assert_string]},
    )
    _get_tracking_store().delete_review_queue(request_message.queue_id)
    return _wrap_response(DeleteReviewQueue.Response())


@catch_mlflow_exception
@_disable_if_artifacts_only
def _add_items_to_review_queue():
    request_message = _get_request_message(
        AddItemsToReviewQueue(),
        schema={"queue_id": [_assert_required, _assert_string]},
    )
    store = _get_tracking_store()
    # Normalize (strip + de-dup) up front so the existence check below validates
    # the same ids the store persists. The store re-normalizes (idempotent); doing
    # it here too keeps the raw-vs-stored mismatch from rejecting a padded-but-valid
    # id as non-existent.
    item_ids = validate_item_ids_for_attach(list(request_message.item_ids))
    kwargs: dict[str, object] = {"item_ids": item_ids}
    if (
        request_message.HasField("item_type")
        and request_message.item_type != REVIEW_ITEM_TYPE_UNSPECIFIED
    ):
        kwargs["item_type"] = ReviewItemType.from_proto(request_message.item_type)
    # Items are trace references with no DB foreign key, so verify every trace
    # exists in the queue's experiment before attaching. A missing or
    # cross-experiment id would otherwise become a ghost PENDING item the UI
    # can't render, and reviewing it would leak traces across experiments.
    queue = store.get_review_queue(request_message.queue_id)
    experiment_by_trace = {
        info.trace_id: str(info.experiment_id) for info in store.batch_get_trace_infos(item_ids)
    }
    if invalid := [i for i in item_ids if experiment_by_trace.get(i) != str(queue.experiment_id)]:
        raise MlflowException(
            f"Cannot attach trace(s) {invalid} to review queue '{request_message.queue_id}': "
            f"they do not exist in experiment '{queue.experiment_id}'.",
            error_code=RESOURCE_DOES_NOT_EXIST,
        )
    items = store.add_items_to_review_queue(request_message.queue_id, **kwargs)
    response = AddItemsToReviewQueue.Response(items=[i.to_proto() for i in items])
    return _wrap_response(response)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _remove_items_from_review_queue():
    request_message = _get_request_message(
        RemoveItemsFromReviewQueue(),
        schema={"queue_id": [_assert_required, _assert_string]},
    )
    _get_tracking_store().remove_items_from_review_queue(
        request_message.queue_id, item_ids=list(request_message.item_ids)
    )
    return _wrap_response(RemoveItemsFromReviewQueue.Response())


@catch_mlflow_exception
@_disable_if_artifacts_only
def _list_review_queue_items():
    request_message = _get_request_message(
        ListReviewQueueItems(),
        schema={
            "queue_id": [_assert_required, _assert_string],
            "max_results": [_assert_intlike, _review_queue_max_results_validator],
        },
    )
    max_results = request_message.max_results if request_message.HasField("max_results") else None
    page_token = request_message.page_token if request_message.HasField("page_token") else None
    status = None
    if request_message.HasField("status") and request_message.status != REVIEW_STATUS_UNSPECIFIED:
        status = ReviewStatus.from_proto(request_message.status)
    items = _get_tracking_store().list_review_queue_items(
        request_message.queue_id,
        status=status,
        max_results=max_results,
        page_token=page_token,
    )
    response = ListReviewQueueItems.Response(
        items=[i.to_proto() for i in items],
        next_page_token=items.token or "",
    )
    return _wrap_response(response)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _set_review_queue_item_status():
    request_message = _get_request_message(
        SetReviewQueueItemStatus(),
        schema={
            "queue_id": [_assert_required, _assert_string],
            "item_id": [_assert_required, _assert_string],
        },
    )
    # `status` is intentionally not in the input schema above: rejection of an
    # absent/UNSPECIFIED status is delegated to `ReviewStatus.from_proto` (a
    # required-field schema entry would only check HasField, not enum value).
    status = ReviewStatus.from_proto(request_message.status)
    # `completed_by` is attribution: bind it to the authenticated caller, never
    # the client value (which could spoof another user). Reopening to `pending`
    # clears attribution, so it stays unset there. On a no-auth server there's no
    # identity to bind to (a single `default` user), so accept the client value.
    username = _get_request_username()
    if username is not None:
        completed_by = None if status == ReviewStatus.PENDING else username
    else:
        completed_by = (
            request_message.completed_by if request_message.HasField("completed_by") else None
        )
    item = _get_tracking_store().set_review_queue_item_status(
        request_message.queue_id,
        item_id=request_message.item_id,
        status=status,
        completed_by=completed_by,
    )
    return _wrap_response(SetReviewQueueItemStatus.Response(item=item.to_proto()))


@catch_mlflow_exception
@_disable_if_artifacts_only
def _invoke_issue_detection_handler():
    """
    Invoke issue detection on traces asynchronously.

    This is a UI-only AJAX endpoint for running issue detection from the frontend.
    """
    from mlflow.genai.discovery.job import _fetch_provider_credentials, invoke_issue_detection_job
    from mlflow.server.jobs import submit_job

    _validate_content_type(request, ["application/json"])

    request_json = _get_validated_flask_request_json(
        schema={
            "experiment_id": [_assert_required, _assert_string],
            "trace_ids": [_assert_required, _assert_array],
            "categories": [_assert_required, _assert_array],
            "provider": [_assert_required, _assert_string],
            "model": [_assert_string],
            "secret_id": [_assert_string],
            "endpoint_name": [_assert_string],
        }
    )

    experiment_id = request_json.get("experiment_id")
    trace_ids = request_json.get("trace_ids", [])
    categories = request_json.get("categories", [])
    provider = request_json.get("provider")
    model = request_json.get("model")
    secret_id = request_json.get("secret_id")
    endpoint_name = request_json.get("endpoint_name")

    if not endpoint_name and not (provider and model):
        raise MlflowException(
            "Either 'endpoint_name' or both 'provider' and 'model' must be provided"
        )

    # Fetch credentials required for executing the job
    if secret_id:
        store = _get_tracking_store()
        credentials = _fetch_provider_credentials(store, provider, secret_id)
    else:
        credentials = None

    # Create the run upfront so we can return run_id immediately
    model_name = f"gateway:/{endpoint_name}" if endpoint_name else f"{provider}:/{model}"
    tags = {
        MLFLOW_RUN_TYPE: MLFLOW_RUN_TYPE_ISSUE_DETECTION,
        "categories": ",".join(categories),
        "model": model_name,
        "total_traces": len(trace_ids),
    }
    if endpoint_name:
        tags["endpoint_name"] = endpoint_name
    run = mlflow.start_run(
        experiment_id=experiment_id,
        tags=tags,
    )
    run_id = run.info.run_id

    job = submit_job(
        function=invoke_issue_detection_job,
        params={
            "experiment_id": experiment_id,
            "trace_ids": trace_ids,
            "categories": categories,
            "run_id": run_id,
            "model": model_name,
        },
        extra_envs=credentials,
    )
    # Tag the run with job ID for later retrieval
    mlflow.set_tag(MLFLOW_ISSUE_DETECTION_JOB_ID, job.job_id)
    mlflow.end_run(RunStatus.to_string(RunStatus.RUNNING))

    return jsonify({"job_id": job.job_id, "run_id": run_id})


@catch_mlflow_exception
@_disable_if_artifacts_only
def _invoke_genai_evaluate_handler():
    """
    Run mlflow.genai.evaluate(...) against the chosen traces + scorers as an
    async job, attached to a brand-new MLflow eval run.

    This is a UI-only AJAX endpoint that backs the "Run evaluation" feature on
    the Evaluation Runs page.
    """
    from mlflow.genai.evaluation.job import invoke_genai_evaluate_job
    from mlflow.server.jobs import submit_job

    _validate_content_type(request, ["application/json"])

    request_json = _get_validated_flask_request_json(
        schema={
            "experiment_id": [_assert_required, _assert_string],
            "trace_ids": [_assert_required, _assert_array],
            "serialized_scorers": [_assert_required, _assert_array],
        }
    )

    experiment_id = request_json["experiment_id"]
    trace_ids = request_json["trace_ids"]
    serialized_scorers = request_json["serialized_scorers"]

    if not trace_ids:
        raise MlflowException(
            "Please select at least one trace to evaluate.",
            error_code=INVALID_PARAMETER_VALUE,
        )
    if not serialized_scorers:
        raise MlflowException(
            "Please select at least one judge.",
            error_code=INVALID_PARAMETER_VALUE,
        )

    # Create the run upfront so we can return run_id immediately, so the run
    # shows up on /evaluation-runs even before the job has produced artifacts.
    tags = {MLFLOW_RUN_TYPE: MLFLOW_RUN_TYPE_GENAI_EVALUATE}
    client = MlflowClient()
    run = client.create_run(experiment_id=experiment_id, tags=tags)
    run_id = run.info.run_id

    username = request.authorization.username if request.authorization else None

    try:
        job = submit_job(
            function=invoke_genai_evaluate_job,
            params={
                "trace_ids": trace_ids,
                "serialized_scorers": serialized_scorers,
                "run_id": run_id,
                "username": username,
            },
        )
        client.set_tag(run_id, MLFLOW_GENAI_EVALUATE_JOB_ID, job.job_id)
    except Exception:
        client.set_terminated(run_id, RunStatus.to_string(RunStatus.FAILED))
        raise

    return jsonify({"job_id": job.job_id, "run_id": run_id})


@catch_mlflow_exception
@_disable_if_artifacts_only
def _get_job(job_id):
    from mlflow.server.jobs import get_job

    job = get_job(job_id)
    return jsonify({
        "status": str(job.status),
        "result": job.parsed_result,
        "status_details": job.status_details,
    })


@catch_mlflow_exception
@_disable_if_artifacts_only
def _cancel_job(job_id):
    from mlflow.server.jobs import cancel_job

    job = cancel_job(job_id)
    return jsonify({
        "status": str(job.status),
        "result": job.parsed_result,
    })


# Deprecated MLflow Tracing APIs. Kept for backward compatibility but do not use.


@catch_mlflow_exception
@_disable_if_artifacts_only
def _deprecated_start_trace_v2():
    """
    A request handler for `POST /mlflow/traces` to create a new TraceInfo record in tracking store.
    """
    request_message = _get_request_message(
        StartTrace(),
        schema={
            "experiment_id": [_assert_string],
            "timestamp_ms": [_assert_intlike],
            "request_metadata": [_assert_map_key_present],
            "tags": [_assert_map_key_present],
        },
    )
    request_metadata = {e.key: e.value for e in request_message.request_metadata}
    tags = {e.key: e.value for e in request_message.tags}

    trace_info = _get_tracking_store().deprecated_start_trace_v2(
        experiment_id=request_message.experiment_id,
        timestamp_ms=request_message.timestamp_ms,
        request_metadata=request_metadata,
        tags=tags,
    )
    response_message = StartTrace.Response(trace_info=trace_info.to_proto())
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _deprecated_end_trace_v2(request_id):
    """
    A request handler for `PATCH /mlflow/traces/{request_id}` to mark an existing TraceInfo
    record completed in tracking store.
    """
    request_message = _get_request_message(
        EndTrace(),
        schema={
            "timestamp_ms": [_assert_intlike],
            "status": [_assert_string],
            "request_metadata": [_assert_map_key_present],
            "tags": [_assert_map_key_present],
        },
    )
    request_metadata = {e.key: e.value for e in request_message.request_metadata}
    tags = {e.key: e.value for e in request_message.tags}

    trace_info = _get_tracking_store().deprecated_end_trace_v2(
        request_id=request_id,
        timestamp_ms=request_message.timestamp_ms,
        status=TraceStatus.from_proto(request_message.status),
        request_metadata=request_metadata,
        tags=tags,
    )

    if isinstance(trace_info, TraceInfo):
        trace_info = TraceInfoV2.from_v3(trace_info)

    response_message = EndTrace.Response(trace_info=trace_info.to_proto())
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _deprecated_get_trace_info_v2(request_id):
    """
    A request handler for `GET /mlflow/traces/{request_id}/info` to retrieve
    an existing TraceInfo record from tracking store.
    """
    trace_info = _get_tracking_store().get_trace_info(request_id)
    trace_info = TraceInfoV2.from_v3(trace_info)
    response_message = GetTraceInfo.Response(trace_info=trace_info.to_proto())
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _deprecated_search_traces_v2():
    """
    A request handler for `GET /mlflow/traces` to search for TraceInfo records in tracking store.
    """
    request_message = _get_request_message(
        SearchTraces(),
        schema={
            "experiment_ids": [
                _assert_array,
                _assert_item_type_string,
                _assert_required,
            ],
            "filter": [_assert_string],
            "max_results": [
                _assert_intlike,
                lambda x: _assert_less_than_or_equal(int(x), 500),
            ],
            "order_by": [_assert_array, _assert_item_type_string],
            "page_token": [_assert_string],
        },
    )

    traces, token = _get_tracking_store().search_traces(
        experiment_ids=request_message.experiment_ids,
        filter_string=request_message.filter,
        max_results=request_message.max_results,
        order_by=request_message.order_by,
        page_token=request_message.page_token or None,
    )
    traces = [TraceInfoV2.from_v3(t) for t in traces]
    response_message = SearchTraces.Response()
    response_message.traces.extend([e.to_proto() for e in traces])
    if token:
        response_message.next_page_token = token
    return _wrap_response(response_message)


# Logged Models APIs


@catch_mlflow_exception
@_disable_if_artifacts_only
def get_logged_model_artifact_handler(model_id: str):
    artifact_file_path = request.args.get("artifact_file_path")
    if not artifact_file_path:
        raise MlflowException(
            'Request must include the "artifact_file_path" query parameter.',
            error_code=BAD_REQUEST,
        )
    validate_path_is_safe(artifact_file_path)

    logged_model: LoggedModel = _get_tracking_store().get_logged_model(model_id)
    if _is_servable_proxied_run_artifact_root(logged_model.artifact_location):
        artifact_repo = _get_artifact_repo_mlflow_artifacts()
        artifact_path = _get_proxied_run_artifact_destination_path(
            proxied_artifact_root=logged_model.artifact_location,
            relative_path=artifact_file_path,
        )
        artifact_path = _get_workspace_scoped_repo_path_if_enabled(artifact_path)
    else:
        artifact_repo = get_artifact_repository(logged_model.artifact_location)
        artifact_path = artifact_file_path

    return _send_artifact(artifact_repo, artifact_path)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _create_logged_model():
    request_message = _get_request_message(
        CreateLoggedModel(),
        schema={
            "experiment_id": [_assert_string, _assert_required],
            "name": [_assert_string],
            "model_type": [_assert_string],
            "source_run_id": [_assert_string],
            "params": [_assert_array],
            "tags": [_assert_array],
        },
    )

    model = _get_tracking_store().create_logged_model(
        experiment_id=request_message.experiment_id,
        name=request_message.name or None,
        model_type=request_message.model_type,
        source_run_id=request_message.source_run_id,
        params=(
            [LoggedModelParameter.from_proto(param) for param in request_message.params]
            if request_message.params
            else None
        ),
        tags=(
            [LoggedModelTag(key=tag.key, value=tag.value) for tag in request_message.tags]
            if request_message.tags
            else None
        ),
    )
    response_message = CreateLoggedModel.Response(model=model.to_proto())
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _log_logged_model_params(model_id: str):
    request_message = _get_request_message(
        LogLoggedModelParamsRequest(),
        schema={
            "model_id": [_assert_string, _assert_required],
            "params": [_assert_array],
        },
    )
    params = (
        [LoggedModelParameter.from_proto(param) for param in request_message.params]
        if request_message.params
        else []
    )
    _get_tracking_store().log_logged_model_params(model_id, params)
    return _wrap_response(LogLoggedModelParamsRequest.Response())


@catch_mlflow_exception
@_disable_if_artifacts_only
def _get_logged_model(model_id: str):
    allow_deleted = request.args.get("allow_deleted", "false").lower() == "true"
    model = _get_tracking_store().get_logged_model(model_id, allow_deleted=allow_deleted)
    response_message = GetLoggedModel.Response(model=model.to_proto())
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _finalize_logged_model(model_id: str):
    request_message = _get_request_message(
        FinalizeLoggedModel(),
        schema={
            "model_id": [_assert_string, _assert_required],
            "status": [_assert_intlike, _assert_required],
        },
    )
    model = _get_tracking_store().finalize_logged_model(
        request_message.model_id, LoggedModelStatus.from_int(request_message.status)
    )
    response_message = FinalizeLoggedModel.Response(model=model.to_proto())
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _delete_logged_model(model_id: str):
    _get_tracking_store().delete_logged_model(model_id)
    return _wrap_response(DeleteLoggedModel.Response())


@catch_mlflow_exception
@_disable_if_artifacts_only
def _set_logged_model_tags(model_id: str):
    request_message = _get_request_message(
        SetLoggedModelTags(),
        schema={"tags": [_assert_array]},
    )
    tags = [LoggedModelTag(key=tag.key, value=tag.value) for tag in request_message.tags]
    _get_tracking_store().set_logged_model_tags(model_id, tags)
    return _wrap_response(SetLoggedModelTags.Response())


@catch_mlflow_exception
@_disable_if_artifacts_only
def _delete_logged_model_tag(model_id: str, tag_key: str):
    _get_tracking_store().delete_logged_model_tag(model_id, tag_key)
    return _wrap_response(DeleteLoggedModelTag.Response())


@catch_mlflow_exception
@_disable_if_artifacts_only
def _search_logged_models():
    request_message = _get_request_message(
        SearchLoggedModels(),
        schema={
            "experiment_ids": [
                _assert_array,
                _assert_item_type_string,
                _assert_required,
            ],
            "filter": [_assert_string],
            "datasets": [_assert_array],
            "max_results": [_assert_intlike],
            "order_by": [_assert_array],
            "page_token": [_assert_string],
        },
    )
    models = _get_tracking_store().search_logged_models(
        # Convert `RepeatedScalarContainer` objects (experiment_ids and order_by) to `list`
        # to avoid serialization issues
        experiment_ids=list(request_message.experiment_ids),
        filter_string=request_message.filter or None,
        datasets=(
            [
                {
                    "dataset_name": d.dataset_name,
                    "dataset_digest": d.dataset_digest or None,
                }
                for d in request_message.datasets
            ]
            if request_message.datasets
            else None
        ),
        max_results=request_message.max_results or None,
        order_by=(
            [
                {
                    "field_name": ob.field_name,
                    "ascending": ob.ascending,
                    "dataset_name": ob.dataset_name or None,
                    "dataset_digest": ob.dataset_digest or None,
                }
                for ob in request_message.order_by
            ]
            if request_message.order_by
            else None
        ),
        page_token=request_message.page_token or None,
    )
    response_message = SearchLoggedModels.Response()
    response_message.models.extend([e.to_proto() for e in models])
    if models.token:
        response_message.next_page_token = models.token
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _list_logged_model_artifacts(model_id: str):
    request_message = _get_request_message(
        ListLoggedModelArtifacts(),
        schema={"artifact_directory_path": [_assert_string]},
    )
    if request_message.HasField("artifact_directory_path"):
        artifact_path = validate_path_is_safe(request_message.artifact_directory_path)
    else:
        artifact_path = None

    return _list_logged_model_artifacts_impl(model_id, artifact_path)


def _list_logged_model_artifacts_impl(
    model_id: str, artifact_directory_path: str | None
) -> Response:
    response = ListLoggedModelArtifacts.Response()
    logged_model: LoggedModel = _get_tracking_store().get_logged_model(model_id)
    if _is_servable_proxied_run_artifact_root(logged_model.artifact_location):
        artifacts = _list_artifacts_for_proxied_run_artifact_root(
            proxied_artifact_root=logged_model.artifact_location,
            relative_path=artifact_directory_path,
        )
    else:
        artifacts = get_artifact_repository(logged_model.artifact_location).list_artifacts(
            artifact_directory_path
        )

    response.files.extend([a.to_proto() for a in artifacts])
    response.root_uri = logged_model.artifact_location
    return _wrap_response(response)


# =============================================================================
# Scorer Management Handlers
# =============================================================================


@catch_mlflow_exception
@_disable_if_artifacts_only
def _register_scorer():
    request_message = _get_request_message(
        RegisterScorer(),
        schema={
            "experiment_id": [_assert_required, _assert_string],
            "name": [_assert_required, _assert_string],
            "serialized_scorer": [_assert_required, _assert_string],
        },
    )
    # Decorator scorers contain a `call_source` field that is executed via exec() during
    # deserialization. The Python client blocks this via `_check_can_be_registered()`, but
    # that check is client-side only and can be bypassed by calling the REST API directly.
    # Enforce the same restriction here in the server handler so it applies regardless of
    # how the request arrives.
    try:
        serialized_data = json.loads(request_message.serialized_scorer)
    except json.JSONDecodeError as e:
        raise MlflowException.invalid_parameter_value("serialized_scorer must be valid JSON") from e
    if serialized_data.get("call_source") is not None:
        raise MlflowException.invalid_parameter_value(
            DECORATOR_SCORER_REGISTRATION_NOT_SUPPORTED_ERROR
        )
    scorer_version = _get_tracking_store().register_scorer(
        request_message.experiment_id,
        request_message.name,
        request_message.serialized_scorer,
    )
    response_message = RegisterScorer.Response()
    response_message.version = scorer_version.scorer_version
    response_message.scorer_id = scorer_version.scorer_id
    response_message.experiment_id = scorer_version.experiment_id
    response_message.name = scorer_version.scorer_name
    response_message.serialized_scorer = scorer_version._serialized_scorer
    response_message.creation_time = scorer_version.creation_time
    response = Response(mimetype="application/json")
    response.set_data(message_to_json(response_message))
    return response


@catch_mlflow_exception
@_disable_if_artifacts_only
def _list_scorers():
    request_message = _get_request_message(
        ListScorers(),
        schema={"experiment_id": [_assert_string]},
    )
    response_message = ListScorers.Response()
    store = _get_tracking_store()
    if request_message.experiment_id:
        scorers = store.list_scorers(request_message.experiment_id)
    else:
        # Cross-experiment listing: walk the active workspace's experiments
        # via the workspace-aware ``search_experiments`` pagination, then
        # batch the scorer fetch through ``list_scorers_across_experiments``.
        # Auth-side ``filter_list_scorers`` applies per-row RBAC filtering on
        # the response.
        experiment_ids: list[str] = []
        page_token: str | None = None
        while True:
            page = store.search_experiments(
                view_type=ViewType.ACTIVE_ONLY,
                max_results=1000,
                page_token=page_token,
            )
            experiment_ids.extend(e.experiment_id for e in page)
            if not (page_token := page.token):
                break
        scorers = store.list_scorers_across_experiments(experiment_ids)
    response_message.scorers.extend([scorer.to_proto() for scorer in scorers])
    response = Response(mimetype="application/json")
    response.set_data(message_to_json(response_message))
    return response


@catch_mlflow_exception
@_disable_if_artifacts_only
def _list_scorer_versions():
    request_message = _get_request_message(
        ListScorerVersions(),
        schema={
            "experiment_id": [_assert_required, _assert_string],
            "name": [_assert_required, _assert_string],
        },
    )
    response_message = ListScorerVersions.Response()
    scorers = _get_tracking_store().list_scorer_versions(
        request_message.experiment_id, request_message.name
    )
    response_message.scorers.extend([scorer.to_proto() for scorer in scorers])
    response = Response(mimetype="application/json")
    response.set_data(message_to_json(response_message))
    return response


@catch_mlflow_exception
@_disable_if_artifacts_only
def _get_scorer():
    request_message = _get_request_message(
        GetScorer(),
        schema={
            "experiment_id": [_assert_required, _assert_string],
            "name": [_assert_required, _assert_string],
            "version": [_assert_intlike],
        },
    )
    response_message = GetScorer.Response()
    scorer_version = _get_tracking_store().get_scorer(
        request_message.experiment_id,
        request_message.name,
        request_message.version if request_message.HasField("version") else None,
    )
    response_message.scorer.CopyFrom(scorer_version.to_proto())
    response = Response(mimetype="application/json")
    response.set_data(message_to_json(response_message))
    return response


@catch_mlflow_exception
@_disable_if_artifacts_only
def _delete_scorer():
    request_message = _get_request_message(
        DeleteScorer(),
        schema={
            "experiment_id": [_assert_required, _assert_string],
            "name": [_assert_required, _assert_string],
            "version": [_assert_intlike],
        },
    )
    _get_tracking_store().delete_scorer(
        request_message.experiment_id,
        request_message.name,
        request_message.version if request_message.HasField("version") else None,
    )
    response_message = DeleteScorer.Response()
    response = Response(mimetype="application/json")
    response.set_data(message_to_json(response_message))
    return response


@catch_mlflow_exception
@_disable_if_artifacts_only
def _get_online_scoring_configs():
    """
    Get online scoring configurations for a list of scorer IDs.

    Query Parameters:
        scorer_ids: List of scorer IDs to fetch configurations for.

    Returns:
        JSON response containing a list of configurations.
    """
    request_json = _get_validated_flask_request_json(
        flask_request=request,
        schema={
            "scorer_ids": [_assert_required, _assert_array, _assert_item_type_string],
        },
    )

    scorer_ids = request_json["scorer_ids"]
    configs = _get_tracking_store().get_online_scoring_configs(scorer_ids)

    response = Response(mimetype="application/json")
    response.set_data(json.dumps({"configs": [c.to_dict() for c in configs]}))
    return response


@catch_mlflow_exception
@_disable_if_artifacts_only
def _upsert_online_scoring_config():
    """
    Update the online scoring configuration for a registered scorer.

    Request Body (JSON):
        experiment_id: The ID of the Experiment containing the scorer.
        name: The scorer name.
        sample_rate: The sampling rate (0.0 to 1.0).
        filter_string: Optional filter string for trace selection.

    Returns:
        JSON response containing the updated configuration.
    """
    request_json = _get_validated_flask_request_json(
        flask_request=request,
        schema={
            "experiment_id": [_assert_required, _assert_string],
            "name": [_assert_required, _assert_string],
            "sample_rate": [_assert_required],
            "filter_string": [],
        },
    )

    filter_string = request_json.get("filter_string")
    if filter_string is not None and not isinstance(filter_string, str):
        raise MlflowException(
            f"Invalid value {filter_string!r} for parameter 'filter_string' supplied: "
            f"Value was of type '{type(filter_string).__name__}'. "
            "Expected type 'str' or None.",
            error_code=INVALID_PARAMETER_VALUE,
        )

    config = _get_tracking_store().upsert_online_scoring_config(
        experiment_id=request_json["experiment_id"],
        scorer_name=request_json["name"],
        sample_rate=float(request_json["sample_rate"]),
        filter_string=filter_string,
    )

    response = Response(mimetype="application/json")
    response.set_data(json.dumps({"config": config.to_dict()}))
    return response


# =============================================================================
# Secrets Management Handlers
# =============================================================================


@catch_mlflow_exception
@_disable_if_artifacts_only
def _create_gateway_secret():
    request_message = _get_request_message(
        CreateGatewaySecret(),
        schema={
            "secret_name": [_assert_required, _assert_string],
            "secret_value": [_assert_secret_value],
            "provider": [_assert_string],
            "created_by": [_assert_string],
        },
    )
    # Empty map means no auth_config was provided
    auth_config = dict(request_message.auth_config) or None

    secret = _get_tracking_store().create_gateway_secret(
        secret_name=request_message.secret_name,
        secret_value=dict(request_message.secret_value),
        provider=request_message.provider or None,
        auth_config=auth_config,
        created_by=request_message.created_by or None,
    )
    response_message = CreateGatewaySecret.Response()
    response_message.secret.CopyFrom(secret.to_proto())
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _get_gateway_secret_info():
    request_message = _get_request_message(
        GetGatewaySecretInfo(),
        schema={
            "secret_id": [_assert_required, _assert_string],
        },
    )
    secret = _get_tracking_store().get_secret_info(request_message.secret_id)
    response_message = GetGatewaySecretInfo.Response()
    response_message.secret.CopyFrom(secret.to_proto())
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _update_gateway_secret():
    request_message = _get_request_message(
        UpdateGatewaySecret(),
        schema={
            "secret_id": [_assert_required, _assert_string],
            "updated_by": [_assert_string],
        },
    )
    # Empty map means no auth_config was provided
    auth_config = dict(request_message.auth_config) or None

    # Empty map means no update to secret_value
    secret_value = dict(request_message.secret_value) or None

    secret = _get_tracking_store().update_gateway_secret(
        secret_id=request_message.secret_id,
        secret_value=secret_value,
        auth_config=auth_config,
        updated_by=request_message.updated_by or None,
    )
    response_message = UpdateGatewaySecret.Response()
    response_message.secret.CopyFrom(secret.to_proto())
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _delete_gateway_secret():
    request_message = _get_request_message(
        DeleteGatewaySecret(),
        schema={
            "secret_id": [_assert_required, _assert_string],
        },
    )
    _get_tracking_store().delete_gateway_secret(request_message.secret_id)
    response_message = DeleteGatewaySecret.Response()
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _list_gateway_secrets():
    request_message = _get_request_message(
        ListGatewaySecretInfos(),
        schema={
            "provider": [_assert_string],
        },
    )
    secrets = _get_tracking_store().list_secret_infos(
        provider=request_message.provider or None,
    )
    response_message = ListGatewaySecretInfos.Response()
    response_message.secrets.extend([s.to_proto() for s in secrets])
    return _wrap_response(response_message)


# =============================================================================
# Endpoints Management Handlers
# =============================================================================


@catch_mlflow_exception
@_disable_if_artifacts_only
def _create_gateway_endpoint():
    request_message = _get_request_message(
        CreateGatewayEndpoint(),
        schema={
            "name": [_assert_required, _assert_string],
            "created_by": [_assert_string],
            "model_configs": [_assert_required],
            "routing_strategy": [_assert_string],
        },
    )
    if request_message.name and not is_valid_endpoint_name(request_message.name):
        raise MlflowException.invalid_parameter_value(
            f"Invalid endpoint name '{request_message.name}'. "
            "Name can only contain letters, numbers, underscores, hyphens, and dots."
        )
    # Convert proto fallback_config to entity FallbackConfig
    fallback_config = None
    if request_message.HasField("fallback_config"):
        fallback_config = FallbackConfig(
            strategy=FallbackStrategy.from_proto(request_message.fallback_config.strategy)
            if request_message.fallback_config.HasField("strategy")
            else None,
            max_attempts=request_message.fallback_config.max_attempts
            if request_message.fallback_config.HasField("max_attempts")
            else None,
        )

    model_configs = [
        GatewayEndpointModelConfig.from_proto(config) for config in request_message.model_configs
    ]

    # Determine experiment_id and usage_tracking
    experiment_id = (
        request_message.experiment_id if request_message.HasField("experiment_id") else None
    )
    usage_tracking = (
        request_message.usage_tracking if request_message.HasField("usage_tracking") else True
    )

    endpoint = _get_tracking_store().create_gateway_endpoint(
        name=request_message.name or None,
        model_configs=model_configs,
        created_by=request_message.created_by or None,
        routing_strategy=RoutingStrategyEntity.from_proto(request_message.routing_strategy)
        if request_message.HasField("routing_strategy")
        else None,
        fallback_config=fallback_config,
        experiment_id=experiment_id,
        usage_tracking=usage_tracking,
    )
    response_message = CreateGatewayEndpoint.Response()
    response_message.endpoint.CopyFrom(endpoint.to_proto())
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _get_gateway_endpoint():
    request_message = _get_request_message(
        GetGatewayEndpoint(),
        schema={
            "endpoint_id": [_assert_string],
            "name": [_assert_string],
        },
    )
    endpoint = _get_tracking_store().get_gateway_endpoint(
        endpoint_id=request_message.endpoint_id or None,
        name=request_message.name or None,
    )
    response_message = GetGatewayEndpoint.Response()
    response_message.endpoint.CopyFrom(endpoint.to_proto())
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _update_gateway_endpoint():
    request_message = _get_request_message(
        UpdateGatewayEndpoint(),
        schema={
            "endpoint_id": [_assert_required, _assert_string],
            "name": [_assert_string],
            "updated_by": [_assert_string],
            "routing_strategy": [_assert_string],
        },
    )
    if request_message.name and not is_valid_endpoint_name(request_message.name):
        raise MlflowException.invalid_parameter_value(
            f"Invalid endpoint name '{request_message.name}'. "
            "Name can only contain letters, numbers, underscores, hyphens, and dots."
        )
    # Convert proto fallback_config to entity FallbackConfig
    fallback_config = None
    if request_message.HasField("fallback_config"):
        fallback_config = FallbackConfig(
            strategy=FallbackStrategy.from_proto(request_message.fallback_config.strategy)
            if request_message.fallback_config.HasField("strategy")
            else None,
            max_attempts=request_message.fallback_config.max_attempts
            if request_message.fallback_config.HasField("max_attempts")
            else None,
        )

    # Convert proto model_configs to entity GatewayEndpointModelConfig list
    model_configs = None
    if request_message.model_configs:
        model_configs = [
            GatewayEndpointModelConfig.from_proto(config)
            for config in request_message.model_configs
        ]

    # Determine experiment_id and usage_tracking
    experiment_id = (
        request_message.experiment_id if request_message.HasField("experiment_id") else None
    )
    usage_tracking = (
        request_message.usage_tracking if request_message.HasField("usage_tracking") else None
    )

    endpoint = _get_tracking_store().update_gateway_endpoint(
        endpoint_id=request_message.endpoint_id,
        name=request_message.name or None,
        model_configs=model_configs,
        updated_by=request_message.updated_by or None,
        routing_strategy=RoutingStrategyEntity.from_proto(request_message.routing_strategy)
        if request_message.HasField("routing_strategy")
        else None,
        fallback_config=fallback_config,
        experiment_id=experiment_id,
        usage_tracking=usage_tracking,
    )
    response_message = UpdateGatewayEndpoint.Response()
    response_message.endpoint.CopyFrom(endpoint.to_proto())
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _delete_gateway_endpoint():
    request_message = _get_request_message(
        DeleteGatewayEndpoint(),
        schema={
            "endpoint_id": [_assert_required, _assert_string],
        },
    )
    _get_tracking_store().delete_gateway_endpoint(request_message.endpoint_id)
    response_message = DeleteGatewayEndpoint.Response()
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _list_gateway_endpoints():
    request_message = _get_request_message(
        ListGatewayEndpoints(),
        schema={
            "provider": [_assert_string],
        },
    )
    endpoints = _get_tracking_store().list_gateway_endpoints(
        provider=request_message.provider or None,
    )
    response_message = ListGatewayEndpoints.Response()
    response_message.endpoints.extend([e.to_proto() for e in endpoints])
    return _wrap_response(response_message)


# =============================================================================
# Model Definitions Management Handlers
# =============================================================================


@catch_mlflow_exception
@_disable_if_artifacts_only
def _create_gateway_model_definition():
    request_message = _get_request_message(
        CreateGatewayModelDefinition(),
        schema={
            "name": [_assert_required, _assert_string],
            "secret_id": [_assert_required, _assert_string],
            "provider": [_assert_required, _assert_string],
            "model_name": [_assert_required, _assert_string],
            "created_by": [_assert_string],
        },
    )
    model_definition = _get_tracking_store().create_gateway_model_definition(
        name=request_message.name,
        secret_id=request_message.secret_id,
        provider=request_message.provider,
        model_name=request_message.model_name,
        created_by=request_message.created_by or None,
    )
    response_message = CreateGatewayModelDefinition.Response()
    response_message.model_definition.CopyFrom(model_definition.to_proto())
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _get_gateway_model_definition():
    request_message = _get_request_message(
        GetGatewayModelDefinition(),
        schema={
            "model_definition_id": [_assert_required, _assert_string],
        },
    )
    model_definition = _get_tracking_store().get_gateway_model_definition(
        request_message.model_definition_id
    )
    response_message = GetGatewayModelDefinition.Response()
    response_message.model_definition.CopyFrom(model_definition.to_proto())
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _list_gateway_model_definitions():
    request_message = _get_request_message(
        ListGatewayModelDefinitions(),
        schema={
            "provider": [_assert_string],
            "secret_id": [_assert_string],
        },
    )
    model_definitions = _get_tracking_store().list_gateway_model_definitions(
        provider=request_message.provider or None,
        secret_id=request_message.secret_id or None,
    )
    response_message = ListGatewayModelDefinitions.Response()
    response_message.model_definitions.extend([m.to_proto() for m in model_definitions])
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _update_gateway_model_definition():
    request_message = _get_request_message(
        UpdateGatewayModelDefinition(),
        schema={
            "model_definition_id": [_assert_required, _assert_string],
            "name": [_assert_string],
            "secret_id": [_assert_string],
            "model_name": [_assert_string],
            "updated_by": [_assert_string],
            "provider": [_assert_string],
        },
    )
    model_definition = _get_tracking_store().update_gateway_model_definition(
        model_definition_id=request_message.model_definition_id,
        name=request_message.name or None,
        secret_id=request_message.secret_id or None,
        model_name=request_message.model_name or None,
        updated_by=request_message.updated_by or None,
        provider=request_message.provider or None,
    )
    response_message = UpdateGatewayModelDefinition.Response()
    response_message.model_definition.CopyFrom(model_definition.to_proto())
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _delete_gateway_model_definition():
    request_message = _get_request_message(
        DeleteGatewayModelDefinition(),
        schema={
            "model_definition_id": [_assert_required, _assert_string],
        },
    )
    _get_tracking_store().delete_gateway_model_definition(request_message.model_definition_id)
    response_message = DeleteGatewayModelDefinition.Response()
    return _wrap_response(response_message)


# =============================================================================
# Endpoint Model Mappings Management Handlers
# =============================================================================


@catch_mlflow_exception
@_disable_if_artifacts_only
def _attach_model_to_gateway_endpoint():
    request_message = _get_request_message(
        AttachModelToGatewayEndpoint(),
        schema={
            "endpoint_id": [_assert_required, _assert_string],
            "model_config": [_assert_required],
            "created_by": [_assert_string],
        },
    )

    model_config = GatewayEndpointModelConfig.from_proto(request_message.model_config)

    mapping = _get_tracking_store().attach_model_to_endpoint(
        endpoint_id=request_message.endpoint_id,
        model_config=model_config,
        created_by=request_message.created_by or None,
    )
    response_message = AttachModelToGatewayEndpoint.Response()
    response_message.mapping.CopyFrom(mapping.to_proto())
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _detach_model_from_gateway_endpoint():
    request_message = _get_request_message(
        DetachModelFromGatewayEndpoint(),
        schema={
            "endpoint_id": [_assert_required, _assert_string],
            "model_definition_id": [_assert_required, _assert_string],
        },
    )
    _get_tracking_store().detach_model_from_endpoint(
        endpoint_id=request_message.endpoint_id,
        model_definition_id=request_message.model_definition_id,
    )
    response_message = DetachModelFromGatewayEndpoint.Response()
    return _wrap_response(response_message)


# =============================================================================
# Endpoint Bindings Management Handlers
# =============================================================================


@catch_mlflow_exception
@_disable_if_artifacts_only
def _create_gateway_endpoint_binding():
    request_message = _get_request_message(
        CreateGatewayEndpointBinding(),
        schema={
            "endpoint_id": [_assert_required, _assert_string],
            "resource_type": [_assert_required, _assert_string],
            "resource_id": [_assert_required, _assert_string],
            "created_by": [_assert_string],
        },
    )
    binding = _get_tracking_store().create_endpoint_binding(
        endpoint_id=request_message.endpoint_id,
        resource_type=GatewayResourceType(request_message.resource_type),
        resource_id=request_message.resource_id,
        created_by=request_message.created_by or None,
    )
    response_message = CreateGatewayEndpointBinding.Response()
    response_message.binding.CopyFrom(binding.to_proto())
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _delete_gateway_endpoint_binding():
    request_message = _get_request_message(
        DeleteGatewayEndpointBinding(),
        schema={
            "endpoint_id": [_assert_required, _assert_string],
            "resource_type": [_assert_required, _assert_string],
            "resource_id": [_assert_required, _assert_string],
        },
    )
    _get_tracking_store().delete_endpoint_binding(
        endpoint_id=request_message.endpoint_id,
        resource_type=request_message.resource_type,
        resource_id=request_message.resource_id,
    )
    response_message = DeleteGatewayEndpointBinding.Response()
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _list_gateway_endpoint_bindings():
    request_message = _get_request_message(
        ListGatewayEndpointBindings(),
        schema={
            "endpoint_id": [_assert_string],
            "resource_type": [_assert_string],
            "resource_id": [_assert_string],
        },
    )
    bindings = _get_tracking_store().list_endpoint_bindings(
        endpoint_id=request_message.endpoint_id or None,
        resource_type=request_message.resource_type or None,
        resource_id=request_message.resource_id or None,
    )
    response_message = ListGatewayEndpointBindings.Response()
    response_message.bindings.extend([b.to_proto() for b in bindings])
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _set_gateway_endpoint_tag():
    request_message = _get_request_message(
        SetGatewayEndpointTag(),
        schema={
            "endpoint_id": [_assert_required, _assert_string],
            "key": [_assert_required, _assert_string],
            "value": [_assert_string],
        },
    )
    tag = GatewayEndpointTag(request_message.key, request_message.value)
    _get_tracking_store().set_gateway_endpoint_tag(request_message.endpoint_id, tag)
    response_message = SetGatewayEndpointTag.Response()
    response = Response(mimetype="application/json")
    response.set_data(message_to_json(response_message))
    return response


@catch_mlflow_exception
@_disable_if_artifacts_only
def _delete_gateway_endpoint_tag():
    request_message = _get_request_message(
        DeleteGatewayEndpointTag(),
        schema={
            "endpoint_id": [_assert_required, _assert_string],
            "key": [_assert_required, _assert_string],
        },
    )
    _get_tracking_store().delete_gateway_endpoint_tag(
        request_message.endpoint_id, request_message.key
    )
    response_message = DeleteGatewayEndpointTag.Response()
    response = Response(mimetype="application/json")
    response.set_data(message_to_json(response_message))
    return response


# =============================================================================
# Budget Policy Management Handlers
# =============================================================================


@catch_mlflow_exception
@_disable_if_artifacts_only
def _create_budget_policy():
    request_message = _get_request_message(
        CreateGatewayBudgetPolicy(),
        schema={
            "budget_unit": [_assert_required],
            "budget_amount": [_assert_required],
            "duration": [_assert_required],
            "target_scope": [_assert_required],
            "budget_action": [_assert_required],
            "created_by": [_assert_string],
        },
    )
    budget_unit = BudgetUnit.from_proto(request_message.budget_unit)
    if budget_unit is None:
        raise MlflowException(
            message=f"Invalid budget_unit: {request_message.budget_unit}",
            error_code=INVALID_PARAMETER_VALUE,
        )
    duration_unit = BudgetDurationUnit.from_proto(request_message.duration.unit)
    if duration_unit is None:
        raise MlflowException(
            message=f"Invalid duration.unit: {request_message.duration.unit}",
            error_code=INVALID_PARAMETER_VALUE,
        )
    if request_message.duration.value <= 0:
        raise MlflowException(
            message=f"duration.value must be a positive integer, got "
            f"{request_message.duration.value}",
            error_code=INVALID_PARAMETER_VALUE,
        )
    target_scope = BudgetTargetScope.from_proto(request_message.target_scope)
    if target_scope is None:
        raise MlflowException(
            message=f"Invalid target_scope: {request_message.target_scope}",
            error_code=INVALID_PARAMETER_VALUE,
        )
    budget_action = BudgetAction.from_proto(request_message.budget_action)
    if budget_action is None:
        raise MlflowException(
            message=f"Invalid budget_action: {request_message.budget_action}",
            error_code=INVALID_PARAMETER_VALUE,
        )
    store = _get_tracking_store()
    policy = store.create_budget_policy(
        budget_unit=budget_unit,
        budget_amount=request_message.budget_amount,
        duration=BudgetDuration(unit=duration_unit, value=request_message.duration.value),
        target_scope=target_scope,
        budget_action=budget_action,
        created_by=request_message.created_by or None,
    )
    get_budget_tracker().invalidate()
    maybe_refresh_budget_policies(store)
    response_message = CreateGatewayBudgetPolicy.Response()
    response_message.budget_policy.CopyFrom(policy.to_proto())
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _get_budget_policy():
    request_message = _get_request_message(
        GetGatewayBudgetPolicy(),
        schema={
            "budget_policy_id": [_assert_required, _assert_string],
        },
    )
    policy = _get_tracking_store().get_budget_policy(
        budget_policy_id=request_message.budget_policy_id,
    )
    response_message = GetGatewayBudgetPolicy.Response()
    response_message.budget_policy.CopyFrom(policy.to_proto())
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _update_budget_policy():
    request_message = _get_request_message(
        UpdateGatewayBudgetPolicy(),
        schema={
            "budget_policy_id": [_assert_required, _assert_string],
            "updated_by": [_assert_string],
        },
    )
    budget_unit = None
    if request_message.HasField("budget_unit"):
        budget_unit = BudgetUnit.from_proto(request_message.budget_unit)
        if budget_unit is None:
            raise MlflowException(
                message=f"Invalid budget_unit: {request_message.budget_unit}",
                error_code=INVALID_PARAMETER_VALUE,
            )
    duration = None
    if request_message.HasField("duration"):
        duration_unit = BudgetDurationUnit.from_proto(request_message.duration.unit)
        if duration_unit is None:
            raise MlflowException(
                message=f"Invalid duration.unit: {request_message.duration.unit}",
                error_code=INVALID_PARAMETER_VALUE,
            )
        if request_message.duration.value <= 0:
            raise MlflowException(
                message=f"duration.value must be a positive integer, got "
                f"{request_message.duration.value}",
                error_code=INVALID_PARAMETER_VALUE,
            )
        duration = BudgetDuration(unit=duration_unit, value=request_message.duration.value)
    target_scope = None
    if request_message.HasField("target_scope"):
        target_scope = BudgetTargetScope.from_proto(request_message.target_scope)
        if target_scope is None:
            raise MlflowException(
                message=f"Invalid target_scope: {request_message.target_scope}",
                error_code=INVALID_PARAMETER_VALUE,
            )
    budget_action = None
    if request_message.HasField("budget_action"):
        budget_action = BudgetAction.from_proto(request_message.budget_action)
        if budget_action is None:
            raise MlflowException(
                message=f"Invalid budget_action: {request_message.budget_action}",
                error_code=INVALID_PARAMETER_VALUE,
            )
    store = _get_tracking_store()
    policy = store.update_budget_policy(
        budget_policy_id=request_message.budget_policy_id,
        budget_unit=budget_unit,
        budget_amount=request_message.budget_amount
        if request_message.HasField("budget_amount")
        else None,
        duration=duration,
        target_scope=target_scope,
        budget_action=budget_action,
        updated_by=request_message.updated_by or None,
    )
    get_budget_tracker().invalidate()
    maybe_refresh_budget_policies(store)
    response_message = UpdateGatewayBudgetPolicy.Response()
    response_message.budget_policy.CopyFrom(policy.to_proto())
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _delete_budget_policy():
    request_message = _get_request_message(
        DeleteGatewayBudgetPolicy(),
        schema={
            "budget_policy_id": [_assert_required, _assert_string],
        },
    )
    store = _get_tracking_store()
    store.delete_budget_policy(request_message.budget_policy_id)
    get_budget_tracker().invalidate()
    maybe_refresh_budget_policies(store)
    response_message = DeleteGatewayBudgetPolicy.Response()
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _list_budget_policies():
    request_message = _get_request_message(
        ListGatewayBudgetPolicies(),
        schema={
            "max_results": [_assert_intlike],
            "page_token": [_assert_string],
        },
    )
    budget_policies = _get_tracking_store().list_budget_policies(
        max_results=request_message.max_results or SEARCH_MAX_RESULTS_DEFAULT,
        page_token=request_message.page_token or None,
    )
    response_message = ListGatewayBudgetPolicies.Response()
    response_message.budget_policies.extend([p.to_proto() for p in budget_policies])
    if budget_policies.token:
        response_message.next_page_token = budget_policies.token
    return _wrap_response(response_message)


def _get_request_workspace_for_budget_windows():
    workspace = workspace_context.get_request_workspace()
    if not MLFLOW_ENABLE_WORKSPACES.get():
        return workspace

    if not workspace_context.is_request_workspace_resolved():
        raise MlflowException(
            "A request workspace must be provided when workspaces are enabled.",
            BAD_REQUEST,
        )

    return workspace


@catch_mlflow_exception
@_disable_if_artifacts_only
def _list_budget_windows():
    _get_request_message(ListGatewayBudgetWindows())
    workspace = _get_request_workspace_for_budget_windows()
    store = _get_tracking_store()
    maybe_refresh_budget_policies(store)
    windows = get_budget_tracker().get_all_windows()
    if workspace is not None:
        windows = [w for w in windows if _policy_applies(w.policy, workspace)]
    response_message = ListGatewayBudgetWindows.Response()
    for w in windows:
        window_msg = ListGatewayBudgetWindows.BudgetWindow(
            budget_policy_id=w.policy.budget_policy_id,
            window_start_ms=int(w.window_start.timestamp() * 1000),
            window_end_ms=int(w.window_end.timestamp() * 1000),
            current_spend=w.cumulative_spend,
        )
        response_message.windows.append(window_msg)
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _create_gateway_guardrail():
    request_message = _get_request_message(
        CreateGatewayGuardrail(),
        schema={
            "name": [_assert_required, _assert_string],
            "scorer_id": [_assert_required, _assert_string],
            "scorer_version": [_assert_required, _assert_intlike],
            "stage": [_assert_required],
            "action": [_assert_required],
            "action_endpoint_id": [_assert_string],
        },
    )
    from mlflow.entities.gateway_guardrail import GuardrailAction, GuardrailStage

    stage = GuardrailStage.from_proto(request_message.stage)
    if stage is None:
        raise MlflowException(
            message=f"Invalid stage: {request_message.stage}",
            error_code=INVALID_PARAMETER_VALUE,
        )
    action = GuardrailAction.from_proto(request_message.action)
    if action is None:
        raise MlflowException(
            message=f"Invalid action: {request_message.action}",
            error_code=INVALID_PARAMETER_VALUE,
        )
    guardrail = _get_tracking_store().create_gateway_guardrail(
        name=request_message.name,
        scorer_id=request_message.scorer_id,
        scorer_version=request_message.scorer_version,
        stage=stage,
        action=action,
        action_endpoint_id=request_message.action_endpoint_id or None,
        created_by=_get_user(),
    )
    response_message = CreateGatewayGuardrail.Response()
    response_message.guardrail.CopyFrom(guardrail.to_proto())
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _get_gateway_guardrail():
    request_message = _get_request_message(
        GetGatewayGuardrail(),
        schema={"guardrail_id": [_assert_required, _assert_string]},
    )
    guardrail = _get_tracking_store().get_gateway_guardrail(
        guardrail_id=request_message.guardrail_id,
    )
    response_message = GetGatewayGuardrail.Response()
    response_message.guardrail.CopyFrom(guardrail.to_proto())
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _delete_gateway_guardrail():
    request_message = _get_request_message(
        DeleteGatewayGuardrail(),
        schema={"guardrail_id": [_assert_required, _assert_string]},
    )
    _get_tracking_store().delete_gateway_guardrail(request_message.guardrail_id)
    return _wrap_response(DeleteGatewayGuardrail.Response())


@catch_mlflow_exception
@_disable_if_artifacts_only
def _list_gateway_guardrails():
    request_message = _get_request_message(
        ListGatewayGuardrails(),
        schema={
            "max_results": [_assert_intlike],
            "page_token": [_assert_string],
        },
    )
    guardrails = _get_tracking_store().list_gateway_guardrails(
        max_results=request_message.max_results or SEARCH_MAX_RESULTS_DEFAULT,
        page_token=request_message.page_token or None,
    )
    response_message = ListGatewayGuardrails.Response()
    response_message.guardrails.extend([g.to_proto() for g in guardrails])
    if guardrails.token:
        response_message.next_page_token = guardrails.token
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _add_guardrail_to_endpoint():
    request_message = _get_request_message(
        AddGuardrailToEndpoint(),
        schema={
            "endpoint_id": [_assert_required, _assert_string],
            "guardrail_id": [_assert_required, _assert_string],
            "execution_order": [_assert_intlike],
        },
    )
    config = _get_tracking_store().add_guardrail_to_endpoint(
        endpoint_id=request_message.endpoint_id,
        guardrail_id=request_message.guardrail_id,
        execution_order=(
            request_message.execution_order if request_message.HasField("execution_order") else None
        ),
        created_by=_get_user(),
    )
    response_message = AddGuardrailToEndpoint.Response()
    response_message.config.CopyFrom(config.to_proto())
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _remove_guardrail_from_endpoint():
    request_message = _get_request_message(
        RemoveGuardrailFromEndpoint(),
        schema={
            "endpoint_id": [_assert_required, _assert_string],
            "guardrail_id": [_assert_required, _assert_string],
        },
    )
    _get_tracking_store().remove_guardrail_from_endpoint(
        endpoint_id=request_message.endpoint_id,
        guardrail_id=request_message.guardrail_id,
    )
    return _wrap_response(RemoveGuardrailFromEndpoint.Response())


@catch_mlflow_exception
@_disable_if_artifacts_only
def _list_endpoint_guardrail_configs():
    request_message = _get_request_message(
        ListEndpointGuardrailConfigs(),
        schema={"endpoint_id": [_assert_required, _assert_string]},
    )
    configs = _get_tracking_store().list_endpoint_guardrail_configs(
        endpoint_id=request_message.endpoint_id,
    )
    response_message = ListEndpointGuardrailConfigs.Response()
    response_message.configs.extend([c.to_proto() for c in configs])
    return _wrap_response(response_message)


@catch_mlflow_exception
def _update_endpoint_guardrail_config():
    request_message = _get_request_message(
        UpdateEndpointGuardrailConfig(),
        schema={
            "endpoint_id": [_assert_required, _assert_string],
            "guardrail_id": [_assert_required, _assert_string],
        },
    )
    kwargs = {
        "endpoint_id": request_message.endpoint_id,
        "guardrail_id": request_message.guardrail_id,
    }
    if request_message.HasField("execution_order"):
        kwargs["execution_order"] = request_message.execution_order
    config = _get_tracking_store().update_endpoint_guardrail_config(**kwargs)
    response_message = UpdateEndpointGuardrailConfig.Response()
    response_message.config.CopyFrom(config.to_proto())
    return _wrap_response(response_message)


@catch_mlflow_exception
def _get_server_info():
    from mlflow.store.tracking.file_store import FileStore
    from mlflow.store.tracking.sqlalchemy_store import SqlAlchemyStore

    store = _get_tracking_store()
    try:
        trace_archival_config = get_trace_archival_server_config()
    except Exception:
        _logger.warning(
            "Failed to load trace archival config while serving server-info; "
            + "defaulting to disabled.",
            exc_info=True,
        )
        trace_archival_config = None
    trace_archival_enabled = bool(
        trace_archival_config
        and trace_archival_config.enabled
        and _store_supports_trace_archival(store)
    )

    if isinstance(store, FileStore):
        store_type = "FileStore"
    elif isinstance(store, SqlAlchemyStore):
        store_type = "SqlStore"
    else:
        store_type = None
    return jsonify({
        "store_type": store_type,
        "workspaces_enabled": MLFLOW_ENABLE_WORKSPACES.get(),
        "trace_archival_enabled": trace_archival_enabled,
    })


@catch_mlflow_exception
@_disable_if_artifacts_only
def _list_supported_providers():
    try:
        providers = get_all_providers()
        return jsonify({"providers": sorted(providers)})
    except ImportError as e:
        raise MlflowException(str(e), error_code=INVALID_PARAMETER_VALUE)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _list_supported_models():
    try:
        provider_filter = request.args.get("provider")
        models = get_models(provider=provider_filter)
        return jsonify({"models": models})
    except ImportError as e:
        raise MlflowException(str(e), error_code=INVALID_PARAMETER_VALUE)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _get_provider_config():
    try:
        provider = request.args.get("provider")
        config = get_provider_config_response(provider)
        return jsonify(config)
    except (ImportError, ValueError) as e:
        raise MlflowException(str(e), error_code=INVALID_PARAMETER_VALUE)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _get_secrets_config():
    using_default_passphrase = not os.environ.get(CRYPTO_KEK_PASSPHRASE_ENV_VAR)
    return jsonify({
        "secrets_available": True,
        "using_default_passphrase": using_default_passphrase,
    })


@catch_mlflow_exception
@_disable_if_artifacts_only
def _invoke_scorer_handler():
    """
    Invoke a scorer on traces asynchronously.

    This is a UI-only AJAX endpoint for invoking scorers from the frontend.
    """
    _validate_content_type(request, ["application/json"])

    args = request.json
    experiment_id = args.get("experiment_id")
    serialized_scorer = args.get("serialized_scorer")
    trace_ids = args.get("trace_ids", [])
    log_assessments = args.get("log_assessments", False)

    if not experiment_id:
        raise MlflowException(
            "Missing required parameter: experiment_id",
            error_code=INVALID_PARAMETER_VALUE,
        )
    if not serialized_scorer:
        raise MlflowException(
            "Missing required parameter: serialized_scorer",
            error_code=INVALID_PARAMETER_VALUE,
        )
    if not trace_ids:
        raise MlflowException(
            "Please select at least one trace to evaluate.",
            error_code=INVALID_PARAMETER_VALUE,
        )

    from mlflow.genai.scorers.base import Scorer
    from mlflow.genai.scorers.job import get_trace_batches_for_scorer, invoke_scorer_job
    from mlflow.server.jobs import submit_job

    scorer = Scorer.model_validate_json(serialized_scorer)

    tracking_store = _get_tracking_store()
    batches = get_trace_batches_for_scorer(trace_ids, scorer, tracking_store)

    # Extract the authenticated username so that job subprocesses can make
    # gateway requests authorised as the original user (not the admin).
    username = request.authorization.username if request.authorization else None

    jobs = []
    for batch_trace_ids in batches:
        job = submit_job(
            function=invoke_scorer_job,
            params={
                "experiment_id": experiment_id,
                "serialized_scorer": serialized_scorer,
                "trace_ids": batch_trace_ids,
                "log_assessments": log_assessments,
                "username": username,
            },
        )
        jobs.append({"job_id": job.job_id, "trace_ids": batch_trace_ids})

    return jsonify({"jobs": jobs})


def _get_rest_path(base_path, version=2):
    return _add_static_prefix(f"/api/{version}.0{base_path}")


def _get_ajax_path(base_path, version=2):
    return _add_static_prefix(f"/ajax-api/{version}.0{base_path}")


def _add_static_prefix(route: str) -> str:
    if prefix := os.environ.get(STATIC_PREFIX_ENV_VAR):
        return prefix.rstrip("/") + route
    return route


def _get_paths(base_path, version=2):
    """
    A service endpoints base path is typically something like /mlflow/experiment.
    We should register paths like /api/2.0/mlflow/experiment and
    /ajax-api/2.0/mlflow/experiment in the Flask router.
    """
    base_path = _convert_path_parameter_to_flask_format(base_path)
    return [_get_rest_path(base_path, version), _get_ajax_path(base_path, version)]


def _convert_path_parameter_to_flask_format(path):
    """
    Converts path parameter format to Flask compatible format.

    Some protobuf endpoint paths contain parameters like /mlflow/trace/{request_id}.
    This can be interpreted correctly by gRPC framework like Armeria, but Flask does
    not understand it. Instead, we need to specify it with a different format,
    like /mlflow/trace/<request_id>.
    """
    # Handle simple parameters like {trace_id}
    path = re.sub(r"{(\w+)}", r"<\1>", path)

    # Handle Databricks-specific syntax like {assessment.trace_id} -> <trace_id>
    # This is needed because Databricks can extract trace_id from request body,
    # but Flask needs it in the URL path
    return re.sub(r"{assessment\.trace_id}", r"<trace_id>", path)


def get_handler(request_class):
    """
    Args:
        request_class: The type of protobuf message
    """
    return HANDLERS.get(request_class, _not_implemented)


def get_service_endpoints(service, get_handler):
    ret = []
    for service_method in service.DESCRIPTOR.methods:
        endpoints = service_method.GetOptions().Extensions[databricks_pb2.rpc].endpoints
        for endpoint in endpoints:
            for http_path in _get_paths(endpoint.path, version=endpoint.since.major):
                handler = get_handler(service().GetRequestClass(service_method))
                ret.append((http_path, handler, [endpoint.method]))
    return ret


def get_endpoints(get_handler=get_handler):
    """
    Returns:
        List of tuples (path, handler, methods)
    """
    return (
        get_service_endpoints(MlflowService, get_handler)
        + get_internal_online_scoring_endpoints()
        + get_service_endpoints(ModelRegistryService, get_handler)
        + get_service_endpoints(MlflowArtifactsService, get_handler)
        + get_service_endpoints(WebhookService, get_handler)
        + [(_add_static_prefix("/graphql"), _graphql, ["GET", "POST"])]
        # NB: Use _get_paths() so that the endpoint is reachable at both
        # <static-prefix>/api/3.0/mlflow/server-info (for the Python client)
        # and <static-prefix>/ajax-api/3.0/mlflow/server-info (for the frontend).
        + [
            (_path, _get_server_info, ["GET"])
            for _path in _get_paths("/mlflow/server-info", version=3)
        ]
        + get_gateway_endpoints()
        + get_demo_endpoints()
        + get_issues_detection_endpoints()
        + get_genai_evaluate_endpoints()
        + get_job_endpoints()
    )


def get_gateway_endpoints():
    """Returns endpoint tuples for gateway provider/model discovery APIs and scorer invocation."""
    return [
        (
            _get_ajax_path("/mlflow/gateway/supported-providers", version=3),
            _list_supported_providers,
            ["GET"],
        ),
        (
            _get_ajax_path("/mlflow/gateway/supported-models", version=3),
            _list_supported_models,
            ["GET"],
        ),
        (
            _get_ajax_path("/mlflow/gateway/provider-config", version=3),
            _get_provider_config,
            ["GET"],
        ),
        (
            _get_ajax_path("/mlflow/gateway/secrets/config", version=3),
            _get_secrets_config,
            ["GET"],
        ),
        (
            _get_ajax_path("/mlflow/scorer/invoke", version=3),
            _invoke_scorer_handler,
            ["POST"],
        ),
    ]


def get_issues_detection_endpoints():
    return [
        (
            _get_ajax_path("/mlflow/issues/invoke", version=3),
            _invoke_issue_detection_handler,
            ["POST"],
        ),
    ]


def get_genai_evaluate_endpoints():
    return [
        (
            _get_ajax_path("/mlflow/genai/evaluate/invoke", version=3),
            _invoke_genai_evaluate_handler,
            ["POST"],
        ),
    ]


def get_job_endpoints():
    return [
        (
            _get_ajax_path("/mlflow/jobs/cancel/<job_id>", version=3),
            _cancel_job,
            ["PATCH"],
        ),
        (
            _get_ajax_path("/mlflow/jobs/<job_id>", version=3),
            _get_job,
            ["GET"],
        ),
    ]


# Demo APIs

# Serialize demo generation so concurrent requests (e.g. FastAPI running Flask
# handlers in a thread pool) cannot race on the process-wide MLFLOW_WORKSPACE
# env var that WorkspaceContext temporarily sets during generate_all_demos.
_demo_generate_lock = threading.Lock()


def get_demo_endpoints():
    """Returns endpoint tuples for demo data generation and deletion APIs."""
    return [
        (
            _get_ajax_path("/mlflow/demo/generate", version=3),
            _generate_demo,
            ["POST"],
        ),
        (
            _get_ajax_path("/mlflow/demo/delete", version=3),
            _delete_demo,
            ["POST"],
        ),
    ]


@catch_mlflow_exception
@_disable_if_artifacts_only
def _generate_demo():
    """Generate demo data for registered demo generators.

    Accepts an optional JSON body with a ``features`` list to generate only specific
    features (e.g. ``{"features": ["traces", "prompts"]}``). When omitted, all features
    are generated.
    """
    from mlflow.demo import generate_all_demos
    from mlflow.demo.base import DEMO_EXPERIMENT_NAME
    from mlflow.demo.registry import demo_registry

    request_json = request.get_json(silent=True) or {}
    features = request_json.get("features")

    store = _get_tracking_store()
    experiment = store.get_experiment_by_name(DEMO_EXPERIMENT_NAME)

    generator_names = demo_registry.list_generators()
    if features is not None:
        generator_names = [n for n in generator_names if n in features]

    all_exist = False
    if experiment and experiment.lifecycle_stage == "active":
        all_exist = all(demo_registry.get(name)().is_generated() for name in generator_names)

    if experiment and all_exist:
        return jsonify({
            "status": "exists",
            "experiment_id": experiment.experiment_id,
            "features_generated": [],
            "navigation_url": f"/experiments/{experiment.experiment_id}",
        })

    with _demo_generate_lock:
        results = generate_all_demos(features=features)

    experiment = store.get_experiment_by_name(DEMO_EXPERIMENT_NAME)
    experiment_id = experiment.experiment_id if experiment else None
    navigation_url = f"/experiments/{experiment_id}" if experiment_id else "/experiments"

    return jsonify({
        "status": "created",
        "experiment_id": experiment_id,
        "features_generated": [r.feature for r in results],
        "navigation_url": navigation_url,
    })


@catch_mlflow_exception
@_disable_if_artifacts_only
def _delete_demo():
    """Delete demo data for all registered demo generators.

    Performs a full hard delete of the demo experiment and all associated data,
    equivalent to what `mlflow gc` would do. This ensures the demo data is
    completely removed rather than just soft-deleted.
    """
    from mlflow.demo.base import DEMO_EXPERIMENT_NAME
    from mlflow.demo.registry import demo_registry

    deleted_features = []
    for name in demo_registry.list_generators():
        generator = demo_registry.get(name)()
        if generator._data_exists():
            generator.delete_demo()
            deleted_features.append(name)

    store = _get_tracking_store()
    experiment = store.get_experiment_by_name(DEMO_EXPERIMENT_NAME)
    if experiment and experiment.lifecycle_stage == "active":
        store.delete_experiment(experiment.experiment_id)

    return jsonify({
        "status": "deleted",
        "features_deleted": deleted_features,
    })


def get_internal_online_scoring_endpoints():
    """Returns endpoint definitions for internal (non public) online scoring APIs."""
    return [
        (
            _get_ajax_path("/mlflow/scorers/online-configs", version=3),
            _get_online_scoring_configs,
            ["GET"],
        ),
        (
            _get_rest_path("/mlflow/scorers/online-configs", version=3),
            _get_online_scoring_configs,
            ["GET"],
        ),
        (
            _get_ajax_path("/mlflow/scorers/online-config", version=3),
            _upsert_online_scoring_config,
            ["PUT"],
        ),
        (
            _get_rest_path("/mlflow/scorers/online-config", version=3),
            _upsert_online_scoring_config,
            ["PUT"],
        ),
    ]


# Evaluation Dataset APIs


@catch_mlflow_exception
@_disable_if_artifacts_only
def _create_dataset_handler():
    request_message = _get_request_message(
        CreateDataset(),
        schema={
            "name": [_assert_required, _assert_string],
            "experiment_ids": [_assert_array],
            "tags": [_assert_string],
        },
    )

    tags = None
    if hasattr(request_message, "tags") and request_message.tags:
        tags = json.loads(request_message.tags)

    dataset = _get_tracking_store().create_dataset(
        name=request_message.name,
        experiment_ids=list(request_message.experiment_ids)
        if request_message.experiment_ids
        else None,
        tags=tags,
    )

    response_message = CreateDataset.Response()
    response_message.dataset.CopyFrom(dataset.to_proto())
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _get_dataset_handler(dataset_id):
    dataset = _get_tracking_store().get_dataset(dataset_id)

    response_message = GetDataset.Response()
    response_message.dataset.CopyFrom(dataset.to_proto())
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _delete_dataset_handler(dataset_id):
    _get_tracking_store().delete_dataset(dataset_id)

    response_message = DeleteDataset.Response()
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _search_evaluation_datasets_handler():
    request_message = _get_request_message(
        SearchEvaluationDatasets(),
        schema={
            "experiment_ids": [_assert_array],
            "filter_string": [_assert_string],
            "max_results": [_assert_intlike],
            "order_by": [_assert_array],
            "page_token": [_assert_string],
        },
    )

    datasets = _get_tracking_store().search_datasets(
        experiment_ids=list(request_message.experiment_ids)
        if request_message.experiment_ids
        else None,
        filter_string=request_message.filter_string or None,
        max_results=request_message.max_results or None,
        order_by=list(request_message.order_by) if request_message.order_by else None,
        page_token=request_message.page_token or None,
    )

    response_message = SearchEvaluationDatasets.Response()
    response_message.datasets.extend([d.to_proto() for d in datasets])
    if datasets.token:
        response_message.next_page_token = datasets.token

    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _set_dataset_tags_handler(dataset_id):
    request_message = _get_request_message(
        SetDatasetTags(),
        schema={
            "tags": [_assert_required, _assert_string],
        },
    )

    tags = json.loads(request_message.tags)

    _get_tracking_store().set_dataset_tags(
        dataset_id=dataset_id,
        tags=tags,
    )

    response_message = SetDatasetTags.Response()
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _delete_dataset_tag_handler(dataset_id, key):
    _get_tracking_store().delete_dataset_tag(
        dataset_id=dataset_id,
        key=key,
    )

    response_message = DeleteDatasetTag.Response()
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _upsert_dataset_records_handler(dataset_id):
    request_message = _get_request_message(
        UpsertDatasetRecords(),
        schema={
            "records": [_assert_required, _assert_string],
        },
    )

    records = json.loads(request_message.records)

    result = _get_tracking_store().upsert_dataset_records(
        dataset_id=dataset_id,
        records=records,
    )

    response_message = UpsertDatasetRecords.Response()
    response_message.inserted_count = result["inserted"]
    response_message.updated_count = result["updated"]

    return _wrap_response(response_message)


def _get_dataset_experiment_ids_handler(dataset_id):
    """
    Get experiment IDs associated with an evaluation dataset.
    """
    experiment_ids = _get_tracking_store().get_dataset_experiment_ids(dataset_id=dataset_id)

    response_message = GetDatasetExperimentIds.Response()
    response_message.experiment_ids.extend(experiment_ids)

    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _add_dataset_to_experiments_handler(dataset_id):
    request_message = _get_request_message(
        AddDatasetToExperiments(),
        schema={
            "experiment_ids": [_assert_array],
        },
    )

    dataset = _get_tracking_store().add_dataset_to_experiments(
        dataset_id=dataset_id,
        experiment_ids=request_message.experiment_ids,
    )

    response_message = AddDatasetToExperiments.Response()
    response_message.dataset.CopyFrom(dataset.to_proto())
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _remove_dataset_from_experiments_handler(dataset_id):
    request_message = _get_request_message(
        RemoveDatasetFromExperiments(),
        schema={
            "experiment_ids": [_assert_array],
        },
    )

    dataset = _get_tracking_store().remove_dataset_from_experiments(
        dataset_id=dataset_id,
        experiment_ids=request_message.experiment_ids,
    )

    response_message = RemoveDatasetFromExperiments.Response()
    response_message.dataset.CopyFrom(dataset.to_proto())
    return _wrap_response(response_message)


def _get_dataset_records_handler(dataset_id):
    request_message = _get_request_message(
        GetDatasetRecords(),
        schema={
            "max_results": [_assert_intlike],
            "page_token": [_assert_string],
        },
    )

    max_results = request_message.max_results or 1000
    page_token = request_message.page_token or None

    # Use the pagination-aware method
    records, next_page_token = _get_tracking_store()._load_dataset_records(
        dataset_id, max_results=max_results, page_token=page_token
    )

    response_message = GetDatasetRecords.Response()

    records_dicts = [record.to_dict() for record in records]
    response_message.records = json.dumps(records_dicts)

    if next_page_token:
        response_message.next_page_token = next_page_token

    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _delete_dataset_records_handler(dataset_id):
    request_message = _get_request_message(
        DeleteDatasetRecords(),
        schema={
            "dataset_record_ids": [_assert_array],
        },
    )

    deleted_count = _get_tracking_store().delete_dataset_records(
        dataset_id=dataset_id,
        dataset_record_ids=list(request_message.dataset_record_ids),
    )

    response_message = DeleteDatasetRecords.Response()
    response_message.deleted_count = deleted_count
    return _wrap_response(response_message)


# Cache for telemetry config with 3 hour TTL
_telemetry_config_cache = TTLCache(maxsize=1, ttl=10800)


def _get_or_fetch_ui_telemetry_config():
    if (config := _telemetry_config_cache.get("config")) is None:
        config = fetch_ui_telemetry_config()
        _telemetry_config_cache["config"] = config
    return config


@catch_mlflow_exception
def get_ui_telemetry_handler():
    """
    GET handler for /telemetry endpoint.
    Returns the telemetry client configuration by fetching it directly.
    """
    if is_telemetry_disabled():
        return jsonify(FALLBACK_UI_CONFIG)

    config = _get_or_fetch_ui_telemetry_config()

    # UI telemetry should be also disabled if overall telemetry is disabled
    disable_ui_telemetry = config.get("disable_ui_telemetry", True) or config.get(
        "disable_telemetry", True
    )
    response = {
        "disable_ui_telemetry": disable_ui_telemetry,
        "disable_ui_events": config.get("disable_ui_events", []),
        "ui_rollout_percentage": config.get("ui_rollout_percentage", 0),
    }
    return jsonify(response)


@catch_mlflow_exception
def post_ui_telemetry_handler():
    """
    POST handler for /telemetry endpoint.
    Accepts telemetry records and adds them to the telemetry client.
    """
    try:
        if is_telemetry_disabled():
            return jsonify({"status": "disabled"})

        data = request.json.get("records", [])

        if not data:
            return jsonify({"status": "success"})

        if (client := get_telemetry_client()) is None:
            return jsonify({"status": "disabled"})

        # check cached config to see if telemetry is disabled
        # if so, don't process the records. we don't rely on the
        # config from the telemetry client because it is only fetched
        # once, so it won't be updated unless the server is restarted.
        config = _get_or_fetch_ui_telemetry_config()

        # if updated telemetry config is disabled / missing, tell the UI to stop sending records
        if config.get("disable_ui_telemetry", True) or config.get("disable_telemetry", True):
            return jsonify({"status": "disabled"})

        server_installation_id = get_or_create_installation_id()
        records = [
            Record(
                event_name=event["event_name"],
                timestamp_ns=event["timestamp_ns"],
                params=event["params"],
                status=Status.SUCCESS,
                installation_id=event["installation_id"],
                session_id=event["session_id"],
                server_installation_id=server_installation_id,
                duration_ms=0,
            )
            for event in data
        ]

        client.add_records(records)

        return jsonify({"status": "success"})
    except Exception as e:
        _logger.debug(f"Failed to process UI telemetry records: {e}")
        # if we run into unexpected errors, likely something is wrong
        # with the data format. if we return success, the UI will continue
        # to send records. if we return an error, the UI will retry sending
        # records. the safest thing to do is to tell the UI to stop sending
        return jsonify({"status": "disabled"})


def _parse_prompt_uri(prompt_uri: str) -> tuple[str, str]:
    """
    Parse a prompt URI to extract the prompt name and version.

    Args:
        prompt_uri: Prompt URI in the format "prompts:/prompt_name/version"

    Returns:
        A tuple of (prompt_name, version). Returns empty strings if parsing fails.
    """
    try:
        # Format: "prompts:/prompt_name/version"
        if prompt_uri.startswith("prompts:/"):
            parts = prompt_uri.replace("prompts:/", "").split("/")
            if len(parts) >= 2:
                return parts[0], parts[1]
    except Exception:
        pass
    return "", ""


@catch_mlflow_exception
@_disable_if_artifacts_only
def _create_prompt_optimization_job():
    # These imports must be local to avoid circular import with mlflow.server.jobs
    from mlflow.genai.datasets import get_dataset as get_genai_dataset
    from mlflow.genai.optimize.job import OptimizerType, optimize_prompts_job
    from mlflow.server.jobs import submit_job

    request_message = _get_request_message(
        CreatePromptOptimizationJob(),
        schema={
            "experiment_id": [_assert_string],
            "source_prompt_uri": [_assert_string, _assert_required],
            "config": [_assert_required],
            "tags": [_assert_array],
        },
    )

    prompt_uri = request_message.source_prompt_uri or ""
    if not prompt_uri:
        raise MlflowException(
            "source_prompt_uri is required for optimization job",
            error_code=INVALID_PARAMETER_VALUE,
        )

    config = request_message.config
    dataset_id = config.dataset_id or ""

    scorers = list(config.scorers) if config.scorers else []

    optimizer_type = OptimizerType.from_proto(config.optimizer_type)

    experiment_id = (request_message.experiment_id or "").strip()
    if not experiment_id:
        raise MlflowException(
            "experiment_id is required for optimization job",
            error_code=INVALID_PARAMETER_VALUE,
        )

    # Parse optimizer_config_json to dict for the job function
    # Validate before creating run to avoid creating unused runs on validation failure
    optimizer_config = None
    if config.optimizer_config_json:
        try:
            optimizer_config = json.loads(config.optimizer_config_json)
        except json.JSONDecodeError as e:
            raise MlflowException(
                f"Invalid JSON in optimizer_config_json: {e}",
                error_code=INVALID_PARAMETER_VALUE,
            )

    # Create MLflow run upfront so run_id is immediately available
    # The job will resume this run when it starts executing
    tracking_store = _get_tracking_store()
    start_time = int(time.time() * 1000)

    # Parse prompt name and version from URI for more descriptive run name
    prompt_name, prompt_version = _parse_prompt_uri(prompt_uri)
    run_name = f"optimize_prompt_{optimizer_type}_{prompt_name}_{prompt_version}_{start_time}"

    run = tracking_store.create_run(
        experiment_id=experiment_id,
        user_id=_get_user(),
        start_time=start_time,
        tags=[],
        run_name=run_name,
    )
    run_id = run.info.run_id

    # Log optimization config as run parameters
    params_to_log = [
        Param("source_prompt_uri", prompt_uri),
        Param("optimizer_type", optimizer_type),
        Param("dataset_id", dataset_id),
        Param("scorer_names", json.dumps(scorers)),
    ]
    if config.optimizer_config_json:
        params_to_log.append(Param("optimizer_config_json", config.optimizer_config_json))
    tracking_store.log_batch(run_id=run_id, metrics=[], params=params_to_log, tags=[])

    # Link the evaluation dataset to the run for lineage tracking (if dataset_id is provided)
    if dataset_id:
        dataset = get_genai_dataset(dataset_id=dataset_id)
        dataset_input = DatasetInput(
            dataset=dataset._to_mlflow_entity(),
            tags=[InputTag(key="mlflow.data.context", value="optimization")],
        )
        tracking_store.log_inputs(run_id=run_id, datasets=[dataset_input])

    params = {
        "run_id": run_id,
        "experiment_id": experiment_id,
        "prompt_uri": prompt_uri,
        "dataset_id": dataset_id,
        "optimizer_type": optimizer_type,
        "optimizer_config": optimizer_config,
        "scorer_names": scorers,
    }

    job_entity = submit_job(optimize_prompts_job, params)

    response_message = CreatePromptOptimizationJob.Response()
    optimization_job = PromptOptimizationJobProto()
    optimization_job.job_id = job_entity.job_id
    optimization_job.run_id = run_id
    optimization_job.state.status = JobStatus.JOB_STATUS_PENDING
    optimization_job.creation_timestamp_ms = job_entity.creation_time
    optimization_job.experiment_id = experiment_id
    optimization_job.config.CopyFrom(config)
    optimization_job.source_prompt_uri = prompt_uri

    for tag in request_message.tags:
        job_tag = optimization_job.tags.add()
        job_tag.key = tag.key
        job_tag.value = tag.value

    response_message.job.CopyFrom(optimization_job)
    return _wrap_response(response_message)


def _build_prompt_optimization_job_from_entity(job_entity):
    from mlflow.genai.optimize.job import OptimizerType

    optimization_job = PromptOptimizationJobProto()
    optimization_job.job_id = job_entity.job_id
    optimization_job.state.status = job_entity.status.to_proto()
    optimization_job.creation_timestamp_ms = job_entity.creation_time

    params = json.loads(job_entity.params)
    if "experiment_id" in params:
        optimization_job.experiment_id = params["experiment_id"]
    if "prompt_uri" in params:
        optimization_job.source_prompt_uri = params["prompt_uri"]

    if run_id := params.get("run_id"):
        optimization_job.run_id = run_id

    # Populate config from job params
    config = optimization_job.config
    if "optimizer_type" in params:
        try:
            optimizer_type = OptimizerType(params["optimizer_type"])
            config.optimizer_type = optimizer_type.to_proto()
        except (ValueError, KeyError):
            pass
    if params.get("dataset_id"):
        config.dataset_id = params["dataset_id"]
    if "scorer_names" in params:
        try:
            scorer_names = params["scorer_names"]
            if isinstance(scorer_names, str):
                scorer_names = json.loads(scorer_names)
            if isinstance(scorer_names, list):
                config.scorers.extend(scorer_names)
        except (json.JSONDecodeError, TypeError):
            pass
    if params.get("optimizer_config"):
        optimizer_config = params["optimizer_config"]
        if isinstance(optimizer_config, dict):
            config.optimizer_config_json = json.dumps(optimizer_config)
        elif isinstance(optimizer_config, str):
            config.optimizer_config_json = optimizer_config

    # Get optimized_prompt_uri from job result (only available when job succeeds)
    if job_entity.status.name == "SUCCEEDED" and job_entity.parsed_result:
        result = job_entity.parsed_result
        if isinstance(result, dict) and result.get("optimized_prompt_uri"):
            optimization_job.optimized_prompt_uri = result["optimized_prompt_uri"]

    # If job failed, add error message to state
    if job_entity.status.name == "FAILED" and job_entity.parsed_result:
        optimization_job.state.error_message = str(job_entity.parsed_result)

    return optimization_job


@catch_mlflow_exception
@_disable_if_artifacts_only
def _get_prompt_optimization_job(job_id):
    from mlflow.server.jobs import get_job

    job_entity = get_job(job_id)
    optimization_job = _build_prompt_optimization_job_from_entity(job_entity)

    # Fetch MLflow run to get evaluation scores from metrics
    try:
        mlflow_run = _get_tracking_store().get_run(optimization_job.run_id)
        run_metrics = mlflow_run.data.metrics

        # Populate evaluation scores from run metrics
        # Aggregated scores are logged as "initial_eval_score" and "final_eval_score"
        # Per-scorer scores are logged as "initial_eval_score.<scorer_name>" and
        # "final_eval_score.<scorer_name>"
        total_metric_calls = None
        for metric_name, metric_value in run_metrics.items():
            match metric_name.split(".", 1):
                case ["initial_eval_score"]:
                    optimization_job.initial_eval_scores["aggregate"] = metric_value
                case ["final_eval_score"]:
                    optimization_job.final_eval_scores["aggregate"] = metric_value
                case ["initial_eval_score", scorer_name]:
                    optimization_job.initial_eval_scores[scorer_name] = metric_value
                case ["final_eval_score", scorer_name]:
                    optimization_job.final_eval_scores[scorer_name] = metric_value
                case ["total_metric_calls"]:
                    total_metric_calls = metric_value

        if total_metric_calls is not None:
            params = json.loads(job_entity.params)
            optimizer_config = params.get("optimizer_config", {})
            if max_metric_calls := optimizer_config.get("max_metric_calls"):
                progress = round(min(total_metric_calls / max_metric_calls, 1.0), 2)
                optimization_job.state.metadata["progress"] = str(progress)

    except Exception as e:
        _logger.debug("Failed to fetch run details for optimization job %s: %s", job_id, e)

    response_message = GetPromptOptimizationJob.Response()
    response_message.job.CopyFrom(optimization_job)
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _search_prompt_optimization_jobs():
    request_message = _get_request_message(
        SearchPromptOptimizationJobs(),
        schema={
            "experiment_id": [_assert_required, _assert_string],
        },
    )

    job_store = _get_job_store()

    # Search for optimize_prompts jobs in the specified experiment
    jobs = job_store.list_jobs(
        job_name="optimize_prompts",
        params={"experiment_id": request_message.experiment_id},
    )

    response_message = SearchPromptOptimizationJobs.Response()

    for job_entity in jobs:
        optimization_job = _build_prompt_optimization_job_from_entity(job_entity)
        response_message.jobs.append(optimization_job)

    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _cancel_prompt_optimization_job(job_id):
    # This import must be local to avoid circular import with mlflow.server.jobs
    from mlflow.server.jobs import cancel_job

    job_entity = cancel_job(job_id)
    optimization_job = _build_prompt_optimization_job_from_entity(job_entity)
    # Override status to CANCELED since cancel_job may not update the entity status immediately
    optimization_job.state.status = JobStatus.JOB_STATUS_CANCELED

    # Terminate the underlying MLflow run if it exists
    if optimization_job.run_id:
        try:
            _get_tracking_store().update_run_info(
                run_id=optimization_job.run_id,
                run_status=RunStatus.KILLED,
                end_time=get_current_time_millis(),
                run_name=None,
            )
        except Exception:
            # If the run doesn't exist or is already terminated, log warning and continue
            _logger.warning(
                "Failed to terminate MLflow run '%s' when canceling job '%s'",
                optimization_job.run_id,
                job_id,
            )

    response_message = CancelPromptOptimizationJob.Response()
    response_message.job.CopyFrom(optimization_job)
    return _wrap_response(response_message)


@catch_mlflow_exception
@_disable_if_artifacts_only
def _delete_prompt_optimization_job(job_id):
    job_store = _get_job_store()
    job_entity = job_store.get_job(job_id)
    optimization_job = _build_prompt_optimization_job_from_entity(job_entity)
    run_id = optimization_job.run_id

    job_store.delete_jobs(job_ids=[job_id])

    # Delete the associated MLflow run if it exists.
    # Check if run exists before attempting deletion - user may have
    # deleted it manually before the job deletion request.
    if run_id:
        try:
            _get_tracking_store().get_run(run_id)
            _get_tracking_store().delete_run(run_id)
        except MlflowException:
            pass

    response_message = DeletePromptOptimizationJob.Response()
    return _wrap_response(response_message)


HANDLERS = {
    # Tracking Server APIs
    CreateExperiment: _create_experiment,
    GetExperiment: _get_experiment,
    GetExperimentByName: _get_experiment_by_name,
    DeleteExperiment: _delete_experiment,
    RestoreExperiment: _restore_experiment,
    UpdateExperiment: _update_experiment,
    CreateRun: _create_run,
    UpdateRun: _update_run,
    DeleteRun: _delete_run,
    RestoreRun: _restore_run,
    LogParam: _log_param,
    LogMetric: _log_metric,
    SetExperimentTag: _set_experiment_tag,
    DeleteExperimentTag: _delete_experiment_tag,
    SetTag: _set_tag,
    DeleteTag: _delete_tag,
    LogBatch: _log_batch,
    LogModel: _log_model,
    GetRun: _get_run,
    SearchRuns: _search_runs,
    ListArtifacts: _list_artifacts,
    CreatePresignedUploadUrl: _create_presigned_upload_url,
    GetMetricHistory: _get_metric_history,
    GetMetricHistoryBulkInterval: get_metric_history_bulk_interval_handler,
    SearchExperiments: _search_experiments,
    LogInputs: _log_inputs,
    LogOutputs: _log_outputs,
    # Evaluation Dataset APIs
    CreateDataset: _create_dataset_handler,
    GetDataset: _get_dataset_handler,
    DeleteDataset: _delete_dataset_handler,
    SearchEvaluationDatasets: _search_evaluation_datasets_handler,
    SetDatasetTags: _set_dataset_tags_handler,
    DeleteDatasetTag: _delete_dataset_tag_handler,
    UpsertDatasetRecords: _upsert_dataset_records_handler,
    GetDatasetExperimentIds: _get_dataset_experiment_ids_handler,
    GetDatasetRecords: _get_dataset_records_handler,
    DeleteDatasetRecords: _delete_dataset_records_handler,
    AddDatasetToExperiments: _add_dataset_to_experiments_handler,
    RemoveDatasetFromExperiments: _remove_dataset_from_experiments_handler,
    # Model Registry APIs
    CreateRegisteredModel: _create_registered_model,
    GetRegisteredModel: _get_registered_model,
    DeleteRegisteredModel: _delete_registered_model,
    UpdateRegisteredModel: _update_registered_model,
    RenameRegisteredModel: _rename_registered_model,
    SearchRegisteredModels: _search_registered_models,
    GetLatestVersions: _get_latest_versions,
    CreateModelVersion: _create_model_version,
    GetModelVersion: _get_model_version,
    DeleteModelVersion: _delete_model_version,
    UpdateModelVersion: _update_model_version,
    TransitionModelVersionStage: _transition_stage,
    GetModelVersionDownloadUri: _get_model_version_download_uri,
    SearchModelVersions: _search_model_versions,
    SetRegisteredModelTag: _set_registered_model_tag,
    DeleteRegisteredModelTag: _delete_registered_model_tag,
    SetModelVersionTag: _set_model_version_tag,
    DeleteModelVersionTag: _delete_model_version_tag,
    SetRegisteredModelAlias: _set_registered_model_alias,
    DeleteRegisteredModelAlias: _delete_registered_model_alias,
    GetModelVersionByAlias: _get_model_version_by_alias,
    # Webhook APIs
    CreateWebhook: _create_webhook,
    ListWebhooks: _list_webhooks,
    GetWebhook: _get_webhook,
    UpdateWebhook: _update_webhook,
    DeleteWebhook: _delete_webhook,
    TestWebhook: _test_webhook,
    # MLflow Artifacts APIs
    DownloadArtifact: _download_artifact,
    UploadArtifact: _upload_artifact,
    ListArtifactsMlflowArtifacts: _list_artifacts_mlflow_artifacts,
    DeleteArtifact: _delete_artifact_mlflow_artifacts,
    CreateMultipartUpload: _create_multipart_upload_artifact,
    CompleteMultipartUpload: _complete_multipart_upload_artifact,
    AbortMultipartUpload: _abort_multipart_upload_artifact,
    GetPresignedDownloadUrl: _get_presigned_download_url,
    # MLflow Tracing APIs (V3)
    StartTraceV3: _start_trace_v3,
    GetTraceInfoV3: _get_trace_info_v3,
    SearchTracesV3: _search_traces_v3,
    DeleteTracesV3: _delete_traces,
    CalculateTraceFilterCorrelation: _calculate_trace_filter_correlation,
    SetTraceTagV3: _set_trace_tag_v3,
    DeleteTraceTagV3: _delete_trace_tag_v3,
    LinkTracesToRun: _link_traces_to_run,
    LinkPromptsToTrace: _link_prompts_to_trace,
    BatchGetTraces: _batch_get_traces,
    BatchGetTraceInfos: _batch_get_trace_infos,
    GetTrace: _get_trace,
    QueryTraceMetrics: _query_trace_metrics,
    # Assessment APIs
    CreateAssessment: _create_assessment,
    GetAssessmentRequest: _get_assessment,
    UpdateAssessment: _update_assessment,
    DeleteAssessment: _delete_assessment,
    # Issue APIs
    CreateIssue: _create_issue,
    UpdateIssue: _update_issue,
    GetIssue: _get_issue,
    SearchIssues: _search_issues,
    # Label Schema APIs
    CreateLabelSchema: _create_label_schema,
    GetLabelSchema: _get_label_schema,
    GetLabelSchemaByName: _get_label_schema_by_name,
    ListLabelSchemas: _list_label_schemas,
    UpdateLabelSchema: _update_label_schema,
    DeleteLabelSchema: _delete_label_schema,
    CreateReviewQueue: _create_review_queue,
    GetOrCreateUserQueue: _get_or_create_user_queue,
    GetReviewQueue: _get_review_queue,
    GetReviewQueueByName: _get_review_queue_by_name,
    ListReviewQueues: _list_review_queues,
    UpdateReviewQueue: _update_review_queue,
    DeleteReviewQueue: _delete_review_queue,
    AddItemsToReviewQueue: _add_items_to_review_queue,
    RemoveItemsFromReviewQueue: _remove_items_from_review_queue,
    ListReviewQueueItems: _list_review_queue_items,
    SetReviewQueueItemStatus: _set_review_queue_item_status,
    # Legacy MLflow Tracing V2 APIs. Kept for backward compatibility but do not use.
    StartTrace: _deprecated_start_trace_v2,
    EndTrace: _deprecated_end_trace_v2,
    GetTraceInfo: _deprecated_get_trace_info_v2,
    SearchTraces: _deprecated_search_traces_v2,
    DeleteTraces: _delete_traces,
    SetTraceTag: _set_trace_tag,
    DeleteTraceTag: _delete_trace_tag,
    # Logged Models APIs
    CreateLoggedModel: _create_logged_model,
    GetLoggedModel: _get_logged_model,
    FinalizeLoggedModel: _finalize_logged_model,
    DeleteLoggedModel: _delete_logged_model,
    SetLoggedModelTags: _set_logged_model_tags,
    DeleteLoggedModelTag: _delete_logged_model_tag,
    SearchLoggedModels: _search_logged_models,
    ListLoggedModelArtifacts: _list_logged_model_artifacts,
    LogLoggedModelParamsRequest: _log_logged_model_params,
    # Scorer APIs
    RegisterScorer: _register_scorer,
    ListScorers: _list_scorers,
    ListScorerVersions: _list_scorer_versions,
    GetScorer: _get_scorer,
    DeleteScorer: _delete_scorer,
    # Secrets APIs
    CreateGatewaySecret: _create_gateway_secret,
    GetGatewaySecretInfo: _get_gateway_secret_info,
    UpdateGatewaySecret: _update_gateway_secret,
    DeleteGatewaySecret: _delete_gateway_secret,
    ListGatewaySecretInfos: _list_gateway_secrets,
    # Endpoints APIs
    CreateGatewayEndpoint: _create_gateway_endpoint,
    GetGatewayEndpoint: _get_gateway_endpoint,
    UpdateGatewayEndpoint: _update_gateway_endpoint,
    DeleteGatewayEndpoint: _delete_gateway_endpoint,
    ListGatewayEndpoints: _list_gateway_endpoints,
    # Model Definitions APIs
    CreateGatewayModelDefinition: _create_gateway_model_definition,
    GetGatewayModelDefinition: _get_gateway_model_definition,
    ListGatewayModelDefinitions: _list_gateway_model_definitions,
    UpdateGatewayModelDefinition: _update_gateway_model_definition,
    DeleteGatewayModelDefinition: _delete_gateway_model_definition,
    # Endpoint Model Mappings APIs
    AttachModelToGatewayEndpoint: _attach_model_to_gateway_endpoint,
    DetachModelFromGatewayEndpoint: _detach_model_from_gateway_endpoint,
    # Endpoint Bindings APIs
    CreateGatewayEndpointBinding: _create_gateway_endpoint_binding,
    DeleteGatewayEndpointBinding: _delete_gateway_endpoint_binding,
    ListGatewayEndpointBindings: _list_gateway_endpoint_bindings,
    # Endpoint Tags APIs
    SetGatewayEndpointTag: _set_gateway_endpoint_tag,
    DeleteGatewayEndpointTag: _delete_gateway_endpoint_tag,
    # Budget Policy APIs
    CreateGatewayBudgetPolicy: _create_budget_policy,
    GetGatewayBudgetPolicy: _get_budget_policy,
    UpdateGatewayBudgetPolicy: _update_budget_policy,
    DeleteGatewayBudgetPolicy: _delete_budget_policy,
    ListGatewayBudgetPolicies: _list_budget_policies,
    ListGatewayBudgetWindows: _list_budget_windows,
    # Guardrail APIs
    CreateGatewayGuardrail: _create_gateway_guardrail,
    GetGatewayGuardrail: _get_gateway_guardrail,
    DeleteGatewayGuardrail: _delete_gateway_guardrail,
    ListGatewayGuardrails: _list_gateway_guardrails,
    AddGuardrailToEndpoint: _add_guardrail_to_endpoint,
    RemoveGuardrailFromEndpoint: _remove_guardrail_from_endpoint,
    ListEndpointGuardrailConfigs: _list_endpoint_guardrail_configs,
    UpdateEndpointGuardrailConfig: _update_endpoint_guardrail_config,
    # Prompt Optimization APIs
    CreatePromptOptimizationJob: _create_prompt_optimization_job,
    GetPromptOptimizationJob: _get_prompt_optimization_job,
    SearchPromptOptimizationJobs: _search_prompt_optimization_jobs,
    CancelPromptOptimizationJob: _cancel_prompt_optimization_job,
    DeletePromptOptimizationJob: _delete_prompt_optimization_job,
    # Workspace APIs
    ListWorkspaces: _list_workspaces_handler,
    CreateWorkspace: _create_workspace_handler,
    GetWorkspace: _get_workspace_handler,
    UpdateWorkspace: _update_workspace_handler,
    DeleteWorkspace: _delete_workspace_handler,
}