mlflow--mlflow/tests/server/test_gateway_api.py

import json
from pathlib import Path
from typing import Any
from unittest import mock
from unittest.mock import AsyncMock, MagicMock, patch

import pytest
from fastapi import FastAPI, HTTPException
from fastapi.responses import StreamingResponse
from starlette.testclient import TestClient

import mlflow
from mlflow.entities import (
    FallbackConfig,
    FallbackStrategy,
    GatewayEndpointModelConfig,
    GatewayModelLinkageType,
    RoutingStrategy,
    SpanType,
)
from mlflow.entities.gateway_guardrail import GuardrailAction, GuardrailStage
from mlflow.entities.trace_state import TraceState
from mlflow.exceptions import MlflowException
from mlflow.gateway.config import (
    EndpointType,
    GatewayRequestType,
    GeminiConfig,
    LiteLLMConfig,
    MistralConfig,
    OpenAIAPIType,
    OpenAIConfig,
)
from mlflow.gateway.constants import MLFLOW_GATEWAY_DURATION_HEADER, MLFLOW_GATEWAY_OVERHEAD_HEADER
from mlflow.gateway.guardrails import _SANITIZE_BYPASS_HEADER, JudgeGuardrail
from mlflow.gateway.providers.anthropic import AnthropicProvider
from mlflow.gateway.providers.base import (
    FallbackProvider,
    TrafficRouteProvider,
)
from mlflow.gateway.providers.databricks import DatabricksConfig, DatabricksProvider
from mlflow.gateway.providers.gemini import GeminiProvider
from mlflow.gateway.providers.litellm import LiteLLMProvider
from mlflow.gateway.providers.mistral import MistralProvider
from mlflow.gateway.providers.openai import OpenAIProvider
from mlflow.gateway.providers.utils import provider_call_duration_ms
from mlflow.gateway.schemas import chat, embeddings
from mlflow.server.fastapi_app import add_gateway_timing_middleware
from mlflow.server.gateway_api import (
    _build_endpoint_config,
    _create_provider_from_endpoint_name,
    anthropic_passthrough_messages,
    chat_completions,
    gateway_router,
    gemini_passthrough_generate_content,
    gemini_passthrough_stream_generate_content,
    invocations,
    openai_passthrough_chat,
    openai_passthrough_embeddings,
    openai_passthrough_responses,
    openai_passthrough_responses_compact,
)
from mlflow.store.tracking.gateway.entities import GatewayEndpointConfig, GatewayModelConfig
from mlflow.store.tracking.sqlalchemy_store import SqlAlchemyStore
from mlflow.tracing.client import TracingClient
from mlflow.tracing.constant import (
    SpanAttributeKey,
    TokenUsageKey,
    TraceMetadataKey,
)

pytestmark = pytest.mark.notrackingurimock

TEST_PASSPHRASE = "test-passphrase-for-gateway-api-tests"


@pytest.fixture(autouse=True)
def set_kek_passphrase(monkeypatch):
    monkeypatch.setenv("MLFLOW_CRYPTO_KEK_PASSPHRASE", TEST_PASSPHRASE)


@pytest.fixture
def store(tmp_path: Path, db_uri: str):
    artifact_uri = tmp_path / "artifacts"
    artifact_uri.mkdir(exist_ok=True)
    mlflow.set_tracking_uri(db_uri)
    yield SqlAlchemyStore(db_uri, artifact_uri.as_uri())
    mlflow.set_tracking_uri(None)


def create_mock_request(
    cached_body: dict[str, Any] | None = None,
    username: str | None = None,
    user_id: int | str | None = None,
) -> MagicMock:
    """Create a mock request with proper state attributes for gateway tests."""
    mock_request = MagicMock()
    mock_request.state.cached_body = cached_body
    mock_request.state.username = username
    mock_request.state.user_id = user_id
    return mock_request


def _make_model_config(provider="openai", model_name="gpt-4o"):
    return GatewayModelConfig(
        model_definition_id="md-test",
        provider=provider,
        model_name=model_name,
        secret_value={"api_key": "sk-test"},
    )


def test_build_endpoint_config_rejects_provider_not_in_allowed_list(monkeypatch):
    monkeypatch.setenv("MLFLOW_GATEWAY_ALLOWED_PROVIDERS", "anthropic")
    with pytest.raises(MlflowException, match="not allowed"):
        _build_endpoint_config("test-ep", _make_model_config("openai"), EndpointType.LLM_V1_CHAT)


def test_build_endpoint_config_allows_provider_in_allowed_list(monkeypatch):
    monkeypatch.setenv("MLFLOW_GATEWAY_ALLOWED_PROVIDERS", "openai")
    config = _build_endpoint_config(
        "test-ep", _make_model_config("openai"), EndpointType.LLM_V1_CHAT
    )
    assert config.name == "test-ep"
    assert isinstance(config.model.config, OpenAIConfig)


def test_build_endpoint_config_allows_provider_when_no_filter():
    config = _build_endpoint_config(
        "test-ep", _make_model_config("openai"), EndpointType.LLM_V1_CHAT
    )
    assert config.name == "test-ep"


def test_create_provider_from_endpoint_name_openai(store: SqlAlchemyStore):
    # Create test data
    secret = store.create_gateway_secret(
        secret_name="openai-key",
        secret_value={"api_key": "sk-test-123"},
        provider="openai",
    )
    model_def = store.create_gateway_model_definition(
        name="gpt-model",
        secret_id=secret.secret_id,
        provider="openai",
        model_name="gpt-4o",
    )
    endpoint = store.create_gateway_endpoint(
        name="test-openai-endpoint",
        model_configs=[
            GatewayEndpointModelConfig(
                model_definition_id=model_def.model_definition_id,
                linkage_type=GatewayModelLinkageType.PRIMARY,
                weight=1.0,
            ),
        ],
    )

    provider, _ = _create_provider_from_endpoint_name(
        store, endpoint.name, EndpointType.LLM_V1_CHAT
    )

    assert isinstance(provider, OpenAIProvider)
    assert isinstance(provider.config.model.config, OpenAIConfig)
    assert provider.config.model.config.openai_api_key == "sk-test-123"


def test_create_provider_from_endpoint_name_azure_openai(store: SqlAlchemyStore):
    # Test Azure OpenAI configuration
    secret = store.create_gateway_secret(
        secret_name="azure-openai-key",
        secret_value={"api_key": "azure-api-key-test"},
        provider="openai",
        auth_config={
            "api_type": "azure",
            "api_base": "https://my-resource.openai.azure.com",
            "api_version": "2024-02-01",
        },
    )
    model_def = store.create_gateway_model_definition(
        name="azure-gpt-model",
        secret_id=secret.secret_id,
        provider="azure",
        model_name="gpt-4",
    )
    endpoint = store.create_gateway_endpoint(
        name="test-azure-endpoint",
        model_configs=[
            GatewayEndpointModelConfig(
                model_definition_id=model_def.model_definition_id,
                linkage_type=GatewayModelLinkageType.PRIMARY,
                weight=1.0,
            ),
        ],
    )

    provider, _ = _create_provider_from_endpoint_name(
        store, endpoint.name, EndpointType.LLM_V1_CHAT
    )

    assert isinstance(provider, OpenAIProvider)
    assert isinstance(provider.config.model.config, OpenAIConfig)
    assert provider.config.model.config.openai_api_type == OpenAIAPIType.AZURE
    assert provider.config.model.config.openai_api_base == "https://my-resource.openai.azure.com"
    assert provider.config.model.config.openai_deployment_name == "gpt-4"
    assert provider.config.model.config.openai_api_version == "2024-02-01"
    assert provider.config.model.config.openai_api_key == "azure-api-key-test"


def test_create_provider_from_endpoint_name_azure_openai_with_azuread(store: SqlAlchemyStore):
    # Test Azure OpenAI with AzureAD authentication
    secret = store.create_gateway_secret(
        secret_name="azuread-openai-key",
        secret_value={"api_key": "azuread-api-key-test"},
        provider="openai",
        auth_config={
            "api_type": "azuread",
            "api_base": "https://my-resource-ad.openai.azure.com",
            "deployment_name": "gpt-4-deployment-ad",
            "api_version": "2024-02-01",
        },
    )
    model_def = store.create_gateway_model_definition(
        name="azuread-gpt-model",
        secret_id=secret.secret_id,
        provider="openai",
        model_name="gpt-4",
    )
    endpoint = store.create_gateway_endpoint(
        name="test-azuread-endpoint",
        model_configs=[
            GatewayEndpointModelConfig(
                model_definition_id=model_def.model_definition_id,
                linkage_type=GatewayModelLinkageType.PRIMARY,
                weight=1.0,
            ),
        ],
    )

    provider, _ = _create_provider_from_endpoint_name(
        store, endpoint.name, EndpointType.LLM_V1_CHAT
    )

    assert isinstance(provider, OpenAIProvider)
    assert isinstance(provider.config.model.config, OpenAIConfig)
    assert provider.config.model.config.openai_api_type == OpenAIAPIType.AZUREAD
    assert provider.config.model.config.openai_api_base == "https://my-resource-ad.openai.azure.com"
    assert provider.config.model.config.openai_deployment_name == "gpt-4-deployment-ad"
    assert provider.config.model.config.openai_api_version == "2024-02-01"
    assert provider.config.model.config.openai_api_key == "azuread-api-key-test"


def test_create_provider_from_endpoint_name_anthropic(store: SqlAlchemyStore):
    secret = store.create_gateway_secret(
        secret_name="anthropic-key",
        secret_value={"api_key": "sk-ant-test"},
        provider="anthropic",
    )
    model_def = store.create_gateway_model_definition(
        name="claude-model",
        secret_id=secret.secret_id,
        provider="anthropic",
        model_name="claude-3-sonnet",
    )
    endpoint = store.create_gateway_endpoint(
        name="test-anthropic-endpoint",
        model_configs=[
            GatewayEndpointModelConfig(
                model_definition_id=model_def.model_definition_id,
                linkage_type=GatewayModelLinkageType.PRIMARY,
                weight=1.0,
            ),
        ],
    )

    provider, _ = _create_provider_from_endpoint_name(
        store, endpoint.name, EndpointType.LLM_V1_CHAT
    )

    assert isinstance(provider, AnthropicProvider)
    assert provider.config.model.config.anthropic_api_key == "sk-ant-test"
    assert provider.base_url == "https://api.anthropic.com/v1"


def test_create_provider_from_endpoint_name_anthropic_with_api_base(store: SqlAlchemyStore):
    secret = store.create_gateway_secret(
        secret_name="anthropic-proxy-key",
        secret_value={"api_key": "sk-ant-proxy-test"},
        provider="anthropic",
        auth_config={"api_base": "http://localhost:6655/anthropic/v1"},
    )
    model_def = store.create_gateway_model_definition(
        name="claude-proxy-model",
        secret_id=secret.secret_id,
        provider="anthropic",
        model_name="claude-3-7-sonnet",
    )
    endpoint = store.create_gateway_endpoint(
        name="test-anthropic-proxy-endpoint",
        model_configs=[
            GatewayEndpointModelConfig(
                model_definition_id=model_def.model_definition_id,
                linkage_type=GatewayModelLinkageType.PRIMARY,
                weight=1.0,
            ),
        ],
    )

    provider, _ = _create_provider_from_endpoint_name(
        store, endpoint.name, EndpointType.LLM_V1_CHAT
    )

    assert isinstance(provider, AnthropicProvider)
    assert provider.config.model.config.anthropic_api_key == "sk-ant-proxy-test"
    assert provider.config.model.config.anthropic_api_base == "http://localhost:6655/anthropic/v1"
    assert provider.base_url == "http://localhost:6655/anthropic/v1"


def test_create_provider_from_endpoint_name_mistral(store: SqlAlchemyStore):
    # Test Mistral provider
    secret = store.create_gateway_secret(
        secret_name="mistral-key",
        secret_value={"api_key": "mistral-test-key"},
        provider="mistral",
    )
    model_def = store.create_gateway_model_definition(
        name="mistral-model",
        secret_id=secret.secret_id,
        provider="mistral",
        model_name="mistral-large-latest",
    )
    endpoint = store.create_gateway_endpoint(
        name="test-mistral-endpoint",
        model_configs=[
            GatewayEndpointModelConfig(
                model_definition_id=model_def.model_definition_id,
                linkage_type=GatewayModelLinkageType.PRIMARY,
                weight=1.0,
            ),
        ],
    )

    provider, _ = _create_provider_from_endpoint_name(
        store, endpoint.name, EndpointType.LLM_V1_CHAT
    )

    assert isinstance(provider, MistralProvider)
    assert isinstance(provider.config.model.config, MistralConfig)
    assert provider.config.model.config.mistral_api_key == "mistral-test-key"


def test_create_provider_from_endpoint_name_gemini(store: SqlAlchemyStore):
    # Test Gemini provider
    secret = store.create_gateway_secret(
        secret_name="gemini-key",
        secret_value={"api_key": "gemini-test-key"},
        provider="gemini",
    )
    model_def = store.create_gateway_model_definition(
        name="gemini-model",
        secret_id=secret.secret_id,
        provider="gemini",
        model_name="gemini-1.5-pro",
    )
    endpoint = store.create_gateway_endpoint(
        name="test-gemini-endpoint",
        model_configs=[
            GatewayEndpointModelConfig(
                model_definition_id=model_def.model_definition_id,
                linkage_type=GatewayModelLinkageType.PRIMARY,
                weight=1.0,
            ),
        ],
    )

    provider, _ = _create_provider_from_endpoint_name(
        store, endpoint.name, EndpointType.LLM_V1_CHAT
    )

    assert isinstance(provider, GeminiProvider)
    assert isinstance(provider.config.model.config, GeminiConfig)
    assert provider.config.model.config.gemini_api_key == "gemini-test-key"


def test_create_provider_from_endpoint_name_litellm(store: SqlAlchemyStore):
    secret = store.create_gateway_secret(
        secret_name="litellm-key",
        secret_value={"api_key": "litellm-test-key"},
        provider="litellm",
    )
    model_def = store.create_gateway_model_definition(
        name="litellm-model",
        secret_id=secret.secret_id,
        provider="litellm",
        model_name="claude-3-5-sonnet-20241022",
    )
    endpoint = store.create_gateway_endpoint(
        name="test-litellm-endpoint",
        model_configs=[
            GatewayEndpointModelConfig(
                model_definition_id=model_def.model_definition_id,
                linkage_type=GatewayModelLinkageType.PRIMARY,
                weight=1.0,
            ),
        ],
    )

    provider, _ = _create_provider_from_endpoint_name(
        store, endpoint.name, EndpointType.LLM_V1_CHAT
    )

    assert isinstance(provider, LiteLLMProvider)
    assert isinstance(provider.config.model.config, LiteLLMConfig)
    assert provider.config.model.config.litellm_auth_config["api_key"] == "litellm-test-key"
    assert provider.config.model.config.litellm_provider == "litellm"
    # get_provider_name() returns the actual provider name for tracing/metrics
    assert provider.get_provider_name() == "litellm"


def test_create_provider_from_endpoint_name_litellm_with_api_base(store: SqlAlchemyStore):
    secret = store.create_gateway_secret(
        secret_name="litellm-custom-key",
        secret_value={"api_key": "litellm-custom-key"},
        provider="litellm",
        auth_config={"api_base": "https://custom-api.example.com"},
    )
    model_def = store.create_gateway_model_definition(
        name="litellm-custom-model",
        secret_id=secret.secret_id,
        provider="litellm",
        model_name="custom-model",
    )
    endpoint = store.create_gateway_endpoint(
        name="test-litellm-custom-endpoint",
        model_configs=[
            GatewayEndpointModelConfig(
                model_definition_id=model_def.model_definition_id,
                linkage_type=GatewayModelLinkageType.PRIMARY,
                weight=1.0,
            ),
        ],
    )

    provider, _ = _create_provider_from_endpoint_name(
        store, endpoint.name, EndpointType.LLM_V1_CHAT
    )

    assert isinstance(provider, LiteLLMProvider)
    assert isinstance(provider.config.model.config, LiteLLMConfig)
    assert provider.config.model.config.litellm_auth_config["api_key"] == "litellm-custom-key"
    assert (
        provider.config.model.config.litellm_auth_config["api_base"]
        == "https://custom-api.example.com"
    )
    assert provider.config.model.config.litellm_provider == "litellm"


@pytest.mark.parametrize(
    "input_url",
    [
        "https://my-workspace.databricks.com",
        "https://my-workspace.databricks.com/serving-endpoints",
    ],
)
def test_create_provider_from_endpoint_name_databricks_normalizes_base_url(
    store: SqlAlchemyStore, input_url: str
):
    secret = store.create_gateway_secret(
        secret_name="databricks-key",
        secret_value={"api_key": "databricks-token-123"},
        provider="databricks",
        auth_config={"api_base": input_url},
    )
    model_def = store.create_gateway_model_definition(
        name="databricks-model",
        secret_id=secret.secret_id,
        provider="databricks",
        model_name="databricks-dbrx-instruct",
    )
    endpoint = store.create_gateway_endpoint(
        name="test-databricks-endpoint",
        model_configs=[
            GatewayEndpointModelConfig(
                model_definition_id=model_def.model_definition_id,
                linkage_type=GatewayModelLinkageType.PRIMARY,
                weight=1.0,
            ),
        ],
    )

    provider, _ = _create_provider_from_endpoint_name(
        store, endpoint.name, EndpointType.LLM_V1_CHAT
    )

    assert isinstance(provider, DatabricksProvider)
    assert isinstance(provider.config.model.config, DatabricksConfig)
    # Verify the base URL was normalized to include /serving-endpoints
    assert provider._api_base == "https://my-workspace.databricks.com/serving-endpoints"


def test_api_key_not_read_from_file(store: SqlAlchemyStore, tmp_path: Path, monkeypatch):
    monkeypatch.delenv("MLFLOW_GATEWAY_RESOLVE_API_KEY_FROM_FILE", raising=False)
    # Create a file whose path will be used as the "api_key" value
    secret_file = tmp_path / "secret.txt"
    secret_file.write_text("file-content-should-not-appear")

    secret = store.create_gateway_secret(
        secret_name="lfi-test-key",
        # Use the file path as the api_key — the gateway must NOT read the file
        secret_value={"api_key": str(secret_file)},
        provider="openai",
    )
    model_def = store.create_gateway_model_definition(
        name="lfi-test-model",
        secret_id=secret.secret_id,
        provider="openai",
        model_name="gpt-4o",
    )
    endpoint = store.create_gateway_endpoint(
        name="lfi-test-endpoint",
        model_configs=[
            GatewayEndpointModelConfig(
                model_definition_id=model_def.model_definition_id,
                linkage_type=GatewayModelLinkageType.PRIMARY,
                weight=1.0,
            ),
        ],
    )

    provider, _ = _create_provider_from_endpoint_name(
        store, endpoint.name, EndpointType.LLM_V1_CHAT
    )

    # The key must be the literal file path string, NOT the file contents
    assert provider.config.model.config.openai_api_key == str(secret_file)
    assert provider.config.model.config.openai_api_key != "file-content-should-not-appear"


def test_create_provider_from_endpoint_name_nonexistent_endpoint(store: SqlAlchemyStore):
    with pytest.raises(MlflowException, match="not found"):
        _create_provider_from_endpoint_name(store, "nonexistent-id", EndpointType.LLM_V1_CHAT)


@pytest.mark.asyncio
async def test_invocations_handler_chat(store: SqlAlchemyStore):
    # Create test data
    secret = store.create_gateway_secret(
        secret_name="chat-key",
        secret_value={"api_key": "sk-test-chat"},
        provider="openai",
    )
    model_def = store.create_gateway_model_definition(
        name="chat-model",
        secret_id=secret.secret_id,
        provider="openai",
        model_name="gpt-4",
    )
    endpoint = store.create_gateway_endpoint(
        name="chat-endpoint",
        model_configs=[
            GatewayEndpointModelConfig(
                model_definition_id=model_def.model_definition_id,
                linkage_type=GatewayModelLinkageType.PRIMARY,
                weight=1.0,
            ),
        ],
    )

    # Mock the provider's chat method
    mock_response = chat.ResponsePayload(
        id="test-id",
        object="chat.completion",
        created=1234567890,
        model="gpt-4",
        choices=[
            chat.Choice(
                index=0,
                message=chat.ResponseMessage(role="assistant", content="Hello!"),
                finish_reason="stop",
            )
        ],
        usage=chat.ChatUsage(prompt_tokens=10, completion_tokens=5, total_tokens=15),
    )

    # Create a mock request with chat payload
    mock_request = create_mock_request()
    mock_request.json = AsyncMock(
        return_value={
            "messages": [{"role": "user", "content": "Hi"}],
            "temperature": 0.7,
            "stream": False,
        }
    )

    # Patch the provider creation to return a mocked provider
    with patch(
        "mlflow.server.gateway_api._create_provider_from_endpoint_name"
    ) as mock_create_provider:
        mock_provider = MagicMock()
        mock_provider.chat = AsyncMock(return_value=mock_response)
        mock_endpoint_config = GatewayEndpointConfig(
            endpoint_id=endpoint.endpoint_id, endpoint_name=endpoint.name, models=[]
        )
        mock_create_provider.return_value = (mock_provider, mock_endpoint_config)

        # Call the handler
        response = await invocations(endpoint.name, mock_request)

        # Verify
        assert response.id == "test-id"
        assert response.choices[0].message.content == "Hello!"
        assert mock_provider.chat.called


@pytest.mark.asyncio
async def test_invocations_handler_embeddings(store: SqlAlchemyStore):
    # Create test data
    secret = store.create_gateway_secret(
        secret_name="embed-key",
        secret_value={"api_key": "sk-test-embed"},
        provider="openai",
    )
    model_def = store.create_gateway_model_definition(
        name="embed-model",
        secret_id=secret.secret_id,
        provider="openai",
        model_name="text-embedding-ada-002",
    )
    endpoint = store.create_gateway_endpoint(
        name="embed-endpoint",
        model_configs=[
            GatewayEndpointModelConfig(
                model_definition_id=model_def.model_definition_id,
                linkage_type=GatewayModelLinkageType.PRIMARY,
                weight=1.0,
            ),
        ],
    )

    # Mock the provider's embeddings method
    mock_response = embeddings.ResponsePayload(
        object="list",
        data=[embeddings.EmbeddingObject(embedding=[0.1, 0.2, 0.3], index=0)],
        model="text-embedding-ada-002",
        usage=embeddings.EmbeddingsUsage(prompt_tokens=5, total_tokens=5),
    )

    # Create a mock request with embeddings payload
    mock_request = create_mock_request()
    mock_request.json = AsyncMock(return_value={"input": "test text"})

    # Patch the provider creation to return a mocked provider
    with patch(
        "mlflow.server.gateway_api._create_provider_from_endpoint_name"
    ) as mock_create_provider:
        mock_provider = MagicMock()
        mock_provider.embeddings = AsyncMock(return_value=mock_response)
        mock_endpoint_config = GatewayEndpointConfig(
            endpoint_id=endpoint.endpoint_id, endpoint_name=endpoint.name, models=[]
        )
        mock_create_provider.return_value = (mock_provider, mock_endpoint_config)

        # Call the handler
        response = await invocations(endpoint.name, mock_request)

        # Verify
        assert response.object == "list"
        assert len(response.data) == 1
        assert response.data[0].embedding == [0.1, 0.2, 0.3]
        assert mock_provider.embeddings.called


def test_gateway_router_initialization():
    assert gateway_router is not None
    assert gateway_router.prefix == "/gateway"


@pytest.mark.asyncio
async def test_invocations_handler_invalid_json(store: SqlAlchemyStore):
    secret = store.create_gateway_secret(
        secret_name="test-key",
        secret_value={"api_key": "sk-test"},
        provider="openai",
    )
    model_def = store.create_gateway_model_definition(
        name="test-model",
        secret_id=secret.secret_id,
        provider="openai",
        model_name="gpt-4",
    )
    endpoint = store.create_gateway_endpoint(
        name="test-endpoint",
        model_configs=[
            GatewayEndpointModelConfig(
                model_definition_id=model_def.model_definition_id,
                linkage_type=GatewayModelLinkageType.PRIMARY,
                weight=1.0,
            ),
        ],
    )

    # Mock request that raises exception when parsing JSON
    mock_request = create_mock_request()
    mock_request.json = AsyncMock(side_effect=ValueError("Invalid JSON"))

    with pytest.raises(HTTPException, match="Invalid JSON payload: Invalid JSON") as exc_info:
        await invocations(endpoint.name, mock_request)

    assert exc_info.value.status_code == 400


@pytest.mark.asyncio
async def test_invocations_handler_missing_fields(store: SqlAlchemyStore):
    secret = store.create_gateway_secret(
        secret_name="test-key",
        secret_value={"api_key": "sk-test"},
        provider="openai",
    )
    model_def = store.create_gateway_model_definition(
        name="test-model",
        secret_id=secret.secret_id,
        provider="openai",
        model_name="gpt-4",
    )
    endpoint = store.create_gateway_endpoint(
        name="test-endpoint",
        model_configs=[
            GatewayEndpointModelConfig(
                model_definition_id=model_def.model_definition_id,
                linkage_type=GatewayModelLinkageType.PRIMARY,
                weight=1.0,
            ),
        ],
    )

    # Create request with neither messages nor input
    mock_request = create_mock_request()
    mock_request.json = AsyncMock(return_value={"temperature": 0.7})

    with patch(
        "mlflow.server.gateway_api._create_provider_from_endpoint_name"
    ) as mock_create_provider:
        mock_provider = MagicMock()
        mock_endpoint_config = GatewayEndpointConfig(
            endpoint_id=endpoint.endpoint_id, endpoint_name=endpoint.name, models=[]
        )
        mock_create_provider.return_value = (mock_provider, mock_endpoint_config)

        with pytest.raises(
            HTTPException, match="Invalid request: payload format must be either chat or embeddings"
        ) as exc_info:
            await invocations(endpoint.name, mock_request)

        assert exc_info.value.status_code == 400


@pytest.mark.asyncio
async def test_invocations_handler_invalid_chat_payload(store: SqlAlchemyStore):
    secret = store.create_gateway_secret(
        secret_name="test-key",
        secret_value={"api_key": "sk-test"},
        provider="openai",
    )
    model_def = store.create_gateway_model_definition(
        name="test-model",
        secret_id=secret.secret_id,
        provider="openai",
        model_name="gpt-4",
    )
    endpoint = store.create_gateway_endpoint(
        name="test-endpoint",
        model_configs=[
            GatewayEndpointModelConfig(
                model_definition_id=model_def.model_definition_id,
                linkage_type=GatewayModelLinkageType.PRIMARY,
                weight=1.0,
            ),
        ],
    )

    # Create request with invalid messages structure
    mock_request = create_mock_request()
    mock_request.json = AsyncMock(
        return_value={
            "messages": "not a list",  # Should be a list
            "stream": False,
        }
    )

    with patch(
        "mlflow.server.gateway_api._create_provider_from_endpoint_name"
    ) as mock_create_provider:
        mock_provider = MagicMock()
        mock_endpoint_config = GatewayEndpointConfig(
            endpoint_id=endpoint.endpoint_id, endpoint_name=endpoint.name, models=[]
        )
        mock_create_provider.return_value = (mock_provider, mock_endpoint_config)

        with pytest.raises(HTTPException, match="Invalid chat payload") as exc_info:
            await invocations(endpoint.name, mock_request)

        assert exc_info.value.status_code == 400


@pytest.mark.asyncio
async def test_invocations_handler_invalid_embeddings_payload(store: SqlAlchemyStore):
    secret = store.create_gateway_secret(
        secret_name="test-key",
        secret_value={"api_key": "sk-test"},
        provider="openai",
    )
    model_def = store.create_gateway_model_definition(
        name="test-model",
        secret_id=secret.secret_id,
        provider="openai",
        model_name="text-embedding-ada-002",
    )
    endpoint = store.create_gateway_endpoint(
        name="test-endpoint",
        model_configs=[
            GatewayEndpointModelConfig(
                model_definition_id=model_def.model_definition_id,
                linkage_type=GatewayModelLinkageType.PRIMARY,
                weight=1.0,
            ),
        ],
    )

    # Create request with invalid input structure
    mock_request = create_mock_request()
    mock_request.json = AsyncMock(
        return_value={
            "input": 123,  # Should be string or list of strings
        }
    )

    with patch(
        "mlflow.server.gateway_api._create_provider_from_endpoint_name"
    ) as mock_create_provider:
        mock_provider = MagicMock()
        mock_endpoint_config = GatewayEndpointConfig(
            endpoint_id=endpoint.endpoint_id, endpoint_name=endpoint.name, models=[]
        )
        mock_create_provider.return_value = (mock_provider, mock_endpoint_config)

        with pytest.raises(HTTPException, match="Invalid embeddings payload") as exc_info:
            await invocations(endpoint.name, mock_request)

        assert exc_info.value.status_code == 400


@pytest.mark.asyncio
async def test_invocations_handler_streaming(store: SqlAlchemyStore):
    secret = store.create_gateway_secret(
        secret_name="test-key",
        secret_value={"api_key": "sk-test"},
        provider="openai",
    )
    model_def = store.create_gateway_model_definition(
        name="test-model",
        secret_id=secret.secret_id,
        provider="openai",
        model_name="gpt-4",
    )
    endpoint = store.create_gateway_endpoint(
        name="test-endpoint",
        model_configs=[
            GatewayEndpointModelConfig(
                model_definition_id=model_def.model_definition_id,
                linkage_type=GatewayModelLinkageType.PRIMARY,
                weight=1.0,
            ),
        ],
    )

    # Create streaming request
    mock_request = create_mock_request()
    mock_request.json = AsyncMock(
        return_value={
            "messages": [{"role": "user", "content": "Hi"}],
            "stream": True,
        }
    )

    # Mock streaming chunks
    async def mock_stream():
        yield chat.StreamResponsePayload(
            id="test-id",
            object="chat.completion.chunk",
            created=1234567890,
            model="gpt-4",
            choices=[
                chat.StreamChoice(
                    index=0,
                    delta=chat.StreamDelta(role="assistant", content="Hello"),
                    finish_reason=None,
                )
            ],
        )

    with patch(
        "mlflow.server.gateway_api._create_provider_from_endpoint_name"
    ) as mock_create_provider:
        mock_provider = MagicMock()
        mock_provider.chat_stream = MagicMock(return_value=mock_stream())
        mock_endpoint_config = GatewayEndpointConfig(
            endpoint_id=endpoint.endpoint_id, endpoint_name=endpoint.name, models=[]
        )
        mock_create_provider.return_value = (mock_provider, mock_endpoint_config)

        response = await invocations(endpoint.name, mock_request)

        assert isinstance(response, StreamingResponse)
        assert response.media_type == "text/event-stream"
        # chat_stream is inside a lazy async generator; consume the body to trigger execution
        async for _ in response.body_iterator:
            pass
        assert mock_provider.chat_stream.called


def test_create_provider_from_endpoint_name_no_models(store: SqlAlchemyStore):
    # Create a minimal endpoint to get an endpoint_name
    secret = store.create_gateway_secret(
        secret_name="test-key",
        secret_value={"api_key": "sk-test"},
        provider="openai",
    )
    model_def = store.create_gateway_model_definition(
        name="test-model",
        secret_id=secret.secret_id,
        provider="openai",
        model_name="gpt-4",
    )
    endpoint = store.create_gateway_endpoint(
        name="test-endpoint",
        model_configs=[
            GatewayEndpointModelConfig(
                model_definition_id=model_def.model_definition_id,
                linkage_type=GatewayModelLinkageType.PRIMARY,
                weight=1.0,
            ),
        ],
    )

    # Mock get_endpoint_config to return an empty models list
    with patch(
        "mlflow.server.gateway_api.get_endpoint_config",
        return_value=GatewayEndpointConfig(
            endpoint_id=endpoint.endpoint_id, endpoint_name="test-endpoint", models=[]
        ),
    ):
        with pytest.raises(MlflowException, match="has no PRIMARY models configured"):
            _create_provider_from_endpoint_name(store, endpoint.name, EndpointType.LLM_V1_CHAT)


# =============================================================================
# OpenAI-compatible chat completions endpoint tests
# =============================================================================


@pytest.mark.asyncio
async def test_chat_completions_endpoint(store: SqlAlchemyStore):
    secret = store.create_gateway_secret(
        secret_name="openai-compat-key",
        secret_value={"api_key": "sk-test"},
        provider="openai",
    )
    model_def = store.create_gateway_model_definition(
        name="openai-compat-model",
        secret_id=secret.secret_id,
        provider="openai",
        model_name="gpt-4",
    )
    store.create_gateway_endpoint(
        name="my-chat-endpoint",
        model_configs=[
            GatewayEndpointModelConfig(
                model_definition_id=model_def.model_definition_id,
                linkage_type=GatewayModelLinkageType.PRIMARY,
                weight=1.0,
            ),
        ],
    )

    # Mock the provider's chat method
    mock_response = chat.ResponsePayload(
        id="test-id",
        object="chat.completion",
        created=1234567890,
        model="gpt-4",
        choices=[
            chat.Choice(
                index=0,
                message=chat.ResponseMessage(role="assistant", content="Hello from OpenAI!"),
                finish_reason="stop",
            )
        ],
        usage=chat.ChatUsage(prompt_tokens=10, completion_tokens=5, total_tokens=15),
    )

    # Create a mock request with OpenAI-compatible format
    mock_request = create_mock_request()
    mock_request.json = AsyncMock(
        return_value={
            "model": "my-chat-endpoint",  # Endpoint name via model parameter
            "messages": [{"role": "user", "content": "Hi"}],
            "temperature": 0.7,
            "stream": False,
        }
    )

    # Patch the provider creation to return a mocked provider
    with (
        patch(
            "mlflow.server.gateway_api._create_provider_from_endpoint_name"
        ) as mock_create_provider,
        patch("mlflow.server.gateway_api.load_guardrails", return_value=[]),
    ):
        mock_provider = MagicMock()
        mock_provider.chat = AsyncMock(return_value=mock_response)
        mock_endpoint_config = GatewayEndpointConfig(
            endpoint_id="test-endpoint-id", endpoint_name="my-chat-endpoint", models=[]
        )
        mock_create_provider.return_value = (mock_provider, mock_endpoint_config)

        # Call the handler
        response = await chat_completions(mock_request)

        # Verify
        assert response.id == "test-id"
        assert response.choices[0].message.content == "Hello from OpenAI!"
        assert mock_provider.chat.called


def test_response_timing_headers(store: SqlAlchemyStore):
    app = FastAPI()
    app.include_router(gateway_router)
    add_gateway_timing_middleware(app)

    mock_response = chat.ResponsePayload(
        id="test-id",
        object="chat.completion",
        created=1234567890,
        model="gpt-4",
        choices=[
            chat.Choice(
                index=0,
                message=chat.ResponseMessage(role="assistant", content="Hello!"),
                finish_reason="stop",
            )
        ],
        usage=chat.ChatUsage(prompt_tokens=10, completion_tokens=5, total_tokens=15),
    )
    mock_endpoint_config = GatewayEndpointConfig(
        endpoint_id="test-endpoint-id", endpoint_name="my-endpoint", models=[]
    )

    async def _mock_chat_with_provider_timing(payload):
        # Simulate a real provider call by setting the ContextVar as send_request would.
        provider_call_duration_ms.set(50.0)
        return mock_response

    with (
        patch("mlflow.server.gateway_api._get_store", return_value=store),
        patch("mlflow.server.gateway_api.get_request_workspace", return_value=None),
        patch("mlflow.server.gateway_api.check_budget_limit"),
        patch("mlflow.server.gateway_api.load_guardrails", return_value=[]),
        patch(
            "mlflow.server.gateway_api._create_provider_from_endpoint_name"
        ) as mock_create_provider,
    ):
        mock_provider = MagicMock()
        mock_provider.chat = _mock_chat_with_provider_timing
        mock_create_provider.return_value = (mock_provider, mock_endpoint_config)

        client = TestClient(app)
        response = client.post(
            "/gateway/mlflow/v1/chat/completions",
            json={"model": "my-endpoint", "messages": [{"role": "user", "content": "Hi"}]},
        )

    assert response.status_code == 200
    duration = int(response.headers[MLFLOW_GATEWAY_DURATION_HEADER])
    overhead = int(response.headers[MLFLOW_GATEWAY_OVERHEAD_HEADER])
    assert duration >= 0
    assert 0 <= overhead <= duration


def test_response_timing_headers_streaming(store: SqlAlchemyStore):
    app = FastAPI()
    app.include_router(gateway_router)
    add_gateway_timing_middleware(app)

    mock_endpoint_config = GatewayEndpointConfig(
        endpoint_id="test-endpoint-id", endpoint_name="my-endpoint", models=[]
    )

    async def _mock_chat_stream(payload):
        yield chat.StreamResponsePayload(
            id="test-id",
            object="chat.completion.chunk",
            created=1234567890,
            model="gpt-4",
            choices=[
                chat.StreamChoice(
                    index=0,
                    delta=chat.StreamDelta(role="assistant", content="Hello"),
                    finish_reason=None,
                )
            ],
        )

    with (
        patch("mlflow.server.gateway_api._get_store", return_value=store),
        patch("mlflow.server.gateway_api.get_request_workspace", return_value=None),
        patch("mlflow.server.gateway_api.check_budget_limit"),
        patch("mlflow.server.gateway_api.load_guardrails", return_value=[]),
        patch(
            "mlflow.server.gateway_api._create_provider_from_endpoint_name"
        ) as mock_create_provider,
    ):
        mock_provider = MagicMock()
        mock_provider.chat_stream = MagicMock(return_value=_mock_chat_stream(None))
        mock_create_provider.return_value = (mock_provider, mock_endpoint_config)

        client = TestClient(app)
        response = client.post(
            "/gateway/mlflow/v1/chat/completions",
            json={
                "model": "my-endpoint",
                "messages": [{"role": "user", "content": "Hi"}],
                "stream": True,
            },
        )

    assert response.status_code == 200
    assert MLFLOW_GATEWAY_DURATION_HEADER in response.headers
    assert int(response.headers[MLFLOW_GATEWAY_DURATION_HEADER]) >= 0
    # Overhead header is omitted for streaming since provider_call_duration_ms is not set.
    assert MLFLOW_GATEWAY_OVERHEAD_HEADER not in response.headers


def test_response_timing_headers_error(store: SqlAlchemyStore):
    app = FastAPI()
    app.include_router(gateway_router)
    add_gateway_timing_middleware(app)

    mock_endpoint_config = GatewayEndpointConfig(
        endpoint_id="test-endpoint-id", endpoint_name="my-error-endpoint", models=[]
    )

    async def _mock_chat_raises(payload):
        provider_call_duration_ms.set(30.0)
        raise HTTPException(status_code=502, detail="Upstream provider error")

    with (
        patch("mlflow.server.gateway_api._get_store", return_value=store),
        patch("mlflow.server.gateway_api.get_request_workspace", return_value=None),
        patch("mlflow.server.gateway_api.check_budget_limit"),
        patch("mlflow.server.gateway_api.load_guardrails", return_value=[]),
        patch(
            "mlflow.server.gateway_api._create_provider_from_endpoint_name"
        ) as mock_create_provider,
    ):
        mock_provider = MagicMock()
        mock_provider.chat = _mock_chat_raises
        mock_create_provider.return_value = (mock_provider, mock_endpoint_config)

        client = TestClient(app)
        response = client.post(
            "/gateway/mlflow/v1/chat/completions",
            json={"model": "my-error-endpoint", "messages": [{"role": "user", "content": "Hi"}]},
        )

    assert response.status_code == 502
    assert MLFLOW_GATEWAY_DURATION_HEADER in response.headers
    duration = int(response.headers[MLFLOW_GATEWAY_DURATION_HEADER])
    assert duration >= 0
    assert MLFLOW_GATEWAY_OVERHEAD_HEADER in response.headers
    overhead = int(response.headers[MLFLOW_GATEWAY_OVERHEAD_HEADER])
    assert 0 <= overhead <= duration


@pytest.mark.asyncio
async def test_chat_completions_endpoint_streaming(store: SqlAlchemyStore):
    secret = store.create_gateway_secret(
        secret_name="stream-key",
        secret_value={"api_key": "sk-test"},
        provider="openai",
    )
    model_def = store.create_gateway_model_definition(
        name="stream-model",
        secret_id=secret.secret_id,
        provider="openai",
        model_name="gpt-4",
    )
    store.create_gateway_endpoint(
        name="stream-endpoint",
        model_configs=[
            GatewayEndpointModelConfig(
                model_definition_id=model_def.model_definition_id,
                linkage_type=GatewayModelLinkageType.PRIMARY,
                weight=1.0,
            ),
        ],
    )

    # Create streaming request
    mock_request = create_mock_request()
    mock_request.json = AsyncMock(
        return_value={
            "model": "stream-endpoint",
            "messages": [{"role": "user", "content": "Hi"}],
            "stream": True,
        }
    )

    # Mock streaming chunks
    async def mock_stream():
        yield chat.StreamResponsePayload(
            id="test-id",
            object="chat.completion.chunk",
            created=1234567890,
            model="gpt-4",
            choices=[
                chat.StreamChoice(
                    index=0,
                    delta=chat.StreamDelta(role="assistant", content="Hello"),
                    finish_reason=None,
                )
            ],
        )

    with (
        patch(
            "mlflow.server.gateway_api._create_provider_from_endpoint_name"
        ) as mock_create_provider,
        patch("mlflow.server.gateway_api.load_guardrails", return_value=[]),
    ):
        mock_provider = MagicMock()
        mock_provider.chat_stream = MagicMock(return_value=mock_stream())
        mock_endpoint_config = GatewayEndpointConfig(
            endpoint_id="test-endpoint-id", endpoint_name="stream-endpoint", models=[]
        )
        mock_create_provider.return_value = (mock_provider, mock_endpoint_config)

        response = await chat_completions(mock_request)

        assert isinstance(response, StreamingResponse)
        assert response.media_type == "text/event-stream"
        # chat_stream is inside a lazy async generator; consume the body to trigger execution
        async for _ in response.body_iterator:
            pass
        assert mock_provider.chat_stream.called


@pytest.mark.asyncio
async def test_chat_completions_endpoint_missing_model_parameter(store: SqlAlchemyStore):
    # Create request without model parameter
    mock_request = create_mock_request()
    mock_request.json = AsyncMock(
        return_value={
            "messages": [{"role": "user", "content": "Hi"}],
        }
    )

    with pytest.raises(HTTPException, match="Missing required 'model' parameter") as exc_info:
        await chat_completions(mock_request)

    assert exc_info.value.status_code == 400


@pytest.mark.asyncio
async def test_chat_completions_endpoint_missing_messages(store: SqlAlchemyStore):
    # Create test endpoint first so we can test payload validation
    secret = store.create_gateway_secret(
        secret_name="chat-missing-msg-key",
        secret_value={"api_key": "sk-test-key"},
        provider="openai",
    )
    model_def = store.create_gateway_model_definition(
        name="gpt-missing-msg-model",
        secret_id=secret.secret_id,
        provider="openai",
        model_name="gpt-4",
    )
    store.create_gateway_endpoint(
        name="my-endpoint",
        model_configs=[
            GatewayEndpointModelConfig(
                model_definition_id=model_def.model_definition_id,
                linkage_type=GatewayModelLinkageType.PRIMARY,
                weight=1.0,
            ),
        ],
    )

    # Create request without messages
    mock_request = create_mock_request()
    mock_request.json = AsyncMock(
        return_value={
            "model": "my-endpoint",
            "temperature": 0.7,
        }
    )

    with pytest.raises(HTTPException, match="Invalid chat payload") as exc_info:
        await chat_completions(mock_request)

    assert exc_info.value.status_code == 400


@pytest.mark.asyncio
async def test_chat_completions_endpoint_invalid_json(store: SqlAlchemyStore):
    mock_request = create_mock_request()
    mock_request.json = AsyncMock(side_effect=ValueError("Invalid JSON"))

    with pytest.raises(HTTPException, match="Invalid JSON payload: Invalid JSON") as exc_info:
        await chat_completions(mock_request)

    assert exc_info.value.status_code == 400


@pytest.mark.asyncio
async def test_openai_passthrough_chat(store: SqlAlchemyStore):
    secret = store.create_gateway_secret(
        secret_name="openai-passthrough-key",
        secret_value={"api_key": "sk-test-passthrough"},
        provider="openai",
    )
    model_def = store.create_gateway_model_definition(
        name="openai-passthrough-model",
        secret_id=secret.secret_id,
        provider="openai",
        model_name="gpt-4o",
    )
    store.create_gateway_endpoint(
        name="openai-passthrough-endpoint",
        model_configs=[
            GatewayEndpointModelConfig(
                model_definition_id=model_def.model_definition_id,
                linkage_type=GatewayModelLinkageType.PRIMARY,
                weight=1.0,
            ),
        ],
    )

    # Mock OpenAI API response
    mock_response = {
        "id": "chatcmpl-123",
        "object": "chat.completion",
        "created": 1234567890,
        "model": "gpt-4o",
        "choices": [
            {
                "index": 0,
                "message": {"role": "assistant", "content": "Hello from passthrough!"},
                "finish_reason": "stop",
            }
        ],
        "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
    }

    # Create mock request
    mock_request = create_mock_request()
    mock_request.json = AsyncMock(
        return_value={
            "model": "openai-passthrough-endpoint",
            "messages": [{"role": "user", "content": "Hello"}],
        }
    )

    # Mock send_request directly
    with mock.patch(
        "mlflow.gateway.providers.openai.send_request", return_value=mock_response
    ) as mock_send:
        response = await openai_passthrough_chat(mock_request)

        # Verify send_request was called
        assert mock_send.called
        call_kwargs = mock_send.call_args[1]
        assert call_kwargs["path"] == "chat/completions"
        assert call_kwargs["payload"]["model"] == "gpt-4o"
        assert call_kwargs["payload"]["messages"] == [{"role": "user", "content": "Hello"}]

        # Verify response is raw OpenAI format
        assert response["id"] == "chatcmpl-123"
        assert response["model"] == "gpt-4o"
        assert response["choices"][0]["message"]["content"] == "Hello from passthrough!"


@pytest.mark.asyncio
async def test_openai_passthrough_embeddings(store: SqlAlchemyStore):
    secret = store.create_gateway_secret(
        secret_name="openai-embed-passthrough-key",
        secret_value={"api_key": "sk-test-embed"},
        provider="openai",
    )
    model_def = store.create_gateway_model_definition(
        name="openai-embed-passthrough-model",
        secret_id=secret.secret_id,
        provider="openai",
        model_name="text-embedding-3-small",
    )
    store.create_gateway_endpoint(
        name="openai-embed-passthrough-endpoint",
        model_configs=[
            GatewayEndpointModelConfig(
                model_definition_id=model_def.model_definition_id,
                linkage_type=GatewayModelLinkageType.PRIMARY,
                weight=1.0,
            ),
        ],
    )

    # Mock OpenAI API response
    mock_response = {
        "object": "list",
        "data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3]}],
        "model": "text-embedding-3-small",
        "usage": {"prompt_tokens": 5, "total_tokens": 5},
    }

    # Create mock request
    mock_request = create_mock_request()
    mock_request.json = AsyncMock(
        return_value={
            "model": "openai-embed-passthrough-endpoint",
            "input": "Test input",
        }
    )

    # Mock send_request directly
    with mock.patch(
        "mlflow.gateway.providers.openai.send_request", return_value=mock_response
    ) as mock_send:
        response = await openai_passthrough_embeddings(mock_request)

        # Verify send_request was called
        assert mock_send.called
        call_kwargs = mock_send.call_args[1]
        assert call_kwargs["path"] == "embeddings"
        assert call_kwargs["payload"]["model"] == "text-embedding-3-small"
        assert call_kwargs["payload"]["input"] == "Test input"

        # Verify response is raw OpenAI format
        assert response["model"] == "text-embedding-3-small"
        assert response["data"][0]["embedding"] == [0.1, 0.2, 0.3]


@pytest.mark.asyncio
async def test_openai_passthrough_responses(store: SqlAlchemyStore):
    secret = store.create_gateway_secret(
        secret_name="openai-responses-key",
        secret_value={"api_key": "sk-test-responses"},
        provider="openai",
    )
    model_def = store.create_gateway_model_definition(
        name="openai-responses-model",
        secret_id=secret.secret_id,
        provider="openai",
        model_name="gpt-4o",
    )
    store.create_gateway_endpoint(
        name="openai-responses-endpoint",
        model_configs=[
            GatewayEndpointModelConfig(
                model_definition_id=model_def.model_definition_id,
                linkage_type=GatewayModelLinkageType.PRIMARY,
                weight=1.0,
            ),
        ],
    )

    # Mock OpenAI Responses API response (using correct Responses API schema)
    mock_response = {
        "id": "resp-123",
        "object": "response",
        "created": 1234567890,
        "model": "gpt-4o",
        "status": "completed",
        "output": [
            {
                "role": "assistant",
                "content": [{"type": "output_text", "text": "Response from Responses API"}],
            }
        ],
        "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
    }

    # Create mock request
    mock_request = create_mock_request()
    mock_request.json = AsyncMock(
        return_value={
            "model": "openai-responses-endpoint",
            "input": [{"role": "user", "content": "Hello"}],
            "instructions": "You are a helpful assistant",
            "response_format": {"type": "text"},
        }
    )

    # Mock send_request directly
    with mock.patch(
        "mlflow.gateway.providers.openai.send_request", return_value=mock_response
    ) as mock_send:
        response = await openai_passthrough_responses(mock_request)

        # Verify send_request was called
        assert mock_send.called
        call_kwargs = mock_send.call_args[1]
        assert call_kwargs["path"] == "responses"
        assert call_kwargs["payload"]["model"] == "gpt-4o"
        assert call_kwargs["payload"]["input"] == [{"role": "user", "content": "Hello"}]
        assert call_kwargs["payload"]["instructions"] == "You are a helpful assistant"
        assert call_kwargs["payload"]["response_format"] == {"type": "text"}

        # Verify response is raw OpenAI Responses API format
        assert response["id"] == "resp-123"
        assert response["object"] == "response"
        assert response["status"] == "completed"
        assert response["output"][0]["content"][0]["text"] == "Response from Responses API"


@pytest.mark.asyncio
async def test_openai_passthrough_responses_compact(store: SqlAlchemyStore):
    secret = store.create_gateway_secret(
        secret_name="openai-responses-compact-key",
        secret_value={"api_key": "sk-test-responses-compact"},
        provider="openai",
    )
    model_def = store.create_gateway_model_definition(
        name="openai-responses-compact-model",
        secret_id=secret.secret_id,
        provider="openai",
        model_name="gpt-4o",
    )
    store.create_gateway_endpoint(
        name="openai-responses-compact-endpoint",
        model_configs=[
            GatewayEndpointModelConfig(
                model_definition_id=model_def.model_definition_id,
                linkage_type=GatewayModelLinkageType.PRIMARY,
                weight=1.0,
            ),
        ],
    )

    # Mock OpenAI Responses-compact response: same shape as /responses output
    mock_response = {
        "id": "resp-compact-123",
        "object": "response",
        "created": 1234567890,
        "model": "gpt-4o",
        "status": "completed",
        "output": [
            {
                "role": "assistant",
                "content": [{"type": "output_text", "text": "Compacted summary"}],
            }
        ],
        "usage": {"input_tokens": 50, "output_tokens": 10, "total_tokens": 60},
    }

    # Compaction request — typical body carries `previous_response_id` and `model`
    mock_request = create_mock_request()
    mock_request.json = AsyncMock(
        return_value={
            "model": "openai-responses-compact-endpoint",
            "previous_response_id": "resp_abc123",
        }
    )

    with mock.patch(
        "mlflow.gateway.providers.openai.send_request", return_value=mock_response
    ) as mock_send:
        response = await openai_passthrough_responses_compact(mock_request)

        # Verify send_request was called with the /compact upstream path
        assert mock_send.called
        call_kwargs = mock_send.call_args[1]
        assert call_kwargs["path"] == "responses/compact"
        assert call_kwargs["payload"]["model"] == "gpt-4o"
        assert call_kwargs["payload"]["previous_response_id"] == "resp_abc123"

        # Verify response is the raw OpenAI Responses-compact format
        assert response["id"] == "resp-compact-123"
        assert response["object"] == "response"
        assert response["status"] == "completed"
        assert response["output"][0]["content"][0]["text"] == "Compacted summary"


@pytest.mark.asyncio
async def test_openai_passthrough_responses_compact_rejects_stream():
    """``/responses/compact`` is unary upstream; the handler must reject a
    client-supplied ``stream=true`` with HTTP 400 before invoking the provider
    (whose passthrough machinery treats all non-embeddings actions as
    stream-capable and would otherwise open an SSE stream against an upstream
    endpoint that does not support it).
    """
    mock_request = create_mock_request()
    mock_request.json = AsyncMock(
        return_value={
            "model": "openai-responses-compact-endpoint",
            "previous_response_id": "resp_abc123",
            "stream": True,
        }
    )

    # send_request should never be called — the handler rejects before
    # reaching the provider.
    with (
        mock.patch("mlflow.gateway.providers.openai.send_request") as mock_send,
        pytest.raises(HTTPException, match="stream=true is not supported") as exc_info,
    ):
        await openai_passthrough_responses_compact(mock_request)

    assert exc_info.value.status_code == 400
    assert not mock_send.called


@pytest.mark.asyncio
async def test_openai_passthrough_chat_streaming(store: SqlAlchemyStore):
    secret = store.create_gateway_secret(
        secret_name="openai-stream-passthrough-key",
        secret_value={"api_key": "sk-test-stream"},
        provider="openai",
    )
    model_def = store.create_gateway_model_definition(
        name="openai-stream-passthrough-model",
        secret_id=secret.secret_id,
        provider="openai",
        model_name="gpt-4o",
    )
    store.create_gateway_endpoint(
        name="openai-stream-passthrough-endpoint",
        model_configs=[
            GatewayEndpointModelConfig(
                model_definition_id=model_def.model_definition_id,
                linkage_type=GatewayModelLinkageType.PRIMARY,
                weight=1.0,
            ),
        ],
    )

    # Create mock request with streaming enabled
    mock_request = create_mock_request()
    mock_request.json = AsyncMock(
        return_value={
            "model": "openai-stream-passthrough-endpoint",
            "messages": [{"role": "user", "content": "Hello"}],
            "stream": True,
        }
    )

    # Mock streaming response chunks
    mock_stream_chunks = [
        b'data: {"id":"chatcmpl-123","object":"chat.completion.chunk","created":1234567890,"model":"gpt-4o","choices":[{"index":0,"delta":{"role":"assistant","content":"Hello"},"finish_reason":null}]}\n\n',  # noqa: E501
        b'data: {"id":"chatcmpl-123","object":"chat.completion.chunk","created":1234567890,"model":"gpt-4o","choices":[{"index":0,"delta":{"content":" world"},"finish_reason":null}]}\n\n',  # noqa: E501
        b'data: {"id":"chatcmpl-123","object":"chat.completion.chunk","created":1234567890,"model":"gpt-4o","choices":[{"index":0,"delta":{},"finish_reason":"stop"}]}\n\n',  # noqa: E501
    ]

    async def mock_stream_generator():
        for chunk in mock_stream_chunks:
            yield chunk

    with mock.patch(
        "mlflow.gateway.providers.openai.send_stream_request",
        return_value=mock_stream_generator(),
    ) as mock_send_stream:
        response = await openai_passthrough_chat(mock_request)

        assert isinstance(response, StreamingResponse)
        assert response.media_type == "text/event-stream"

        chunks = [chunk async for chunk in response.body_iterator]

        assert mock_send_stream.called
        assert len(chunks) == 3
        assert b"Hello" in chunks[0]
        assert b"world" in chunks[1]
        assert b"stop" in chunks[2]


@pytest.mark.asyncio
async def test_openai_passthrough_responses_streaming(store: SqlAlchemyStore):
    secret = store.create_gateway_secret(
        secret_name="openai-responses-stream-key",
        secret_value={"api_key": "sk-test-responses-stream"},
        provider="openai",
    )
    model_def = store.create_gateway_model_definition(
        name="openai-responses-stream-model",
        secret_id=secret.secret_id,
        provider="openai",
        model_name="gpt-4o",
    )
    store.create_gateway_endpoint(
        name="openai-responses-stream-endpoint",
        model_configs=[
            GatewayEndpointModelConfig(
                model_definition_id=model_def.model_definition_id,
                linkage_type=GatewayModelLinkageType.PRIMARY,
                weight=1.0,
            ),
        ],
    )

    # Create mock request with streaming enabled
    mock_request = create_mock_request()
    mock_request.json = AsyncMock(
        return_value={
            "model": "openai-responses-stream-endpoint",
            "input": [{"type": "text", "text": "Hello"}],
            "instructions": "You are a helpful assistant",
            "stream": True,
        }
    )

    # Mock streaming response chunks for Responses API
    mock_stream_chunks = [
        b'data: {"type":"response.created","response":{"id":"resp_1","object":"response","created_at":1741290958,"status":"in_progress","error":null,"incomplete_details":null,"instructions":"You are a helpful assistant.","max_output_tokens":null,"model":"gpt-4.1-2025-04-14","output":[],"parallel_tool_calls":true,"previous_response_id":null,"reasoning":{"effort":null,"summary":null},"store":true,"temperature":1.0,"text":{"format":{"type":"text"}},"tool_choice":"auto","tools":[],"top_p":1.0,"truncation":"disabled","usage":null,"user":null,"metadata":{}}}\n\n',  # noqa: E501
        b'data: {"type":"response.output_item.added","output_index":0,"item":{"id":"msg_1","type":"message","status":"in_progress","role":"assistant","content":[]}}\n\n',  # noqa: E501
        b'data: {"type":"response.content_part.added","item_id":"msg_1","output_index":0,"content_index":0,"part":{"type":"output_text","text":"","annotations":[]}}\n\n',  # noqa: E501
        b'data: {"type":"response.output_text.delta","item_id":"msg_1","output_index":0,"content_index":0,"delta":"Hi"}\n\n',  # noqa: E501
        b'data: {"type":"response.output_text.done","item_id":"msg_1","output_index":0,"content_index":0,"text":"Hi there! How can I assist you today?"}\n\n',  # noqa: E501
        b'data: {"type":"response.content_part.done","item_id":"msg_1","output_index":0,"content_index":0,"part":{"type":"output_text","text":"Hi there! How can I assist you today?","annotations":[]}}\n\n',  # noqa: E501
        b'data: {"type":"response.output_item.done","output_index":0,"item":{"id":"msg_1","type":"message","status":"completed","role":"assistant","content":[{"type":"output_text","text":"Hi there! How can I assist you today?","annotations":[]}]}}\n\n',  # noqa: E501
        b'data: {"type":"response.completed","response":{"id":"resp_1","object":"response","created_at":1741290958,"status":"completed","error":null,"incomplete_details":null,"instructions":"You are a helpful assistant.","max_output_tokens":null,"model":"gpt-4.1-2025-04-14","output":[{"id":"msg_1","type":"message","status":"completed","role":"assistant","content":[{"type":"output_text","text":"Hi there! How can I assist you today?","annotations":[]}]}],"parallel_tool_calls":true,"previous_response_id":null,"reasoning":{"effort":null,"summary":null},"store":true,"temperature":1.0,"text":{"format":{"type":"text"}},"tool_choice":"auto","tools":[],"top_p":1.0,"truncation":"disabled","usage":{"input_tokens":37,"output_tokens":11,"output_tokens_details":{"reasoning_tokens":0},"total_tokens":48},"user":null,"metadata":{}}}\n\n',  # noqa: E501
    ]

    async def mock_stream_generator():
        for chunk in mock_stream_chunks:
            yield chunk

    with mock.patch(
        "mlflow.gateway.providers.openai.send_stream_request",
        return_value=mock_stream_generator(),
    ) as mock_send_stream:
        response = await openai_passthrough_responses(mock_request)

        assert isinstance(response, StreamingResponse)
        assert response.media_type == "text/event-stream"

        chunks = [chunk async for chunk in response.body_iterator]

        assert mock_send_stream.called
        assert len(chunks) == 8
        assert b"response.created" in chunks[0]
        assert b"response.output_item.added" in chunks[1]
        assert b"response.content_part.added" in chunks[2]
        assert b"response.output_text.delta" in chunks[3]
        assert b"response.output_text.done" in chunks[4]
        assert b"response.content_part.done" in chunks[5]
        assert b"response.output_item.done" in chunks[6]
        assert b"response.completed" in chunks[7]


# =============================================================================
# Anthropic Messages API passthrough endpoint tests
# =============================================================================


@pytest.mark.asyncio
async def test_anthropic_passthrough_messages(store: SqlAlchemyStore):
    secret = store.create_gateway_secret(
        secret_name="anthropic-passthrough-key",
        secret_value={"api_key": "sk-ant-test"},
        provider="anthropic",
    )
    model_def = store.create_gateway_model_definition(
        name="anthropic-passthrough-model",
        secret_id=secret.secret_id,
        provider="anthropic",
        model_name="claude-3-5-sonnet-20241022",
    )
    store.create_gateway_endpoint(
        name="anthropic-passthrough-endpoint",
        model_configs=[
            GatewayEndpointModelConfig(
                model_definition_id=model_def.model_definition_id,
                linkage_type=GatewayModelLinkageType.PRIMARY,
                weight=1.0,
            ),
        ],
    )

    mock_request = create_mock_request()
    mock_request.json = AsyncMock(
        return_value={
            "model": "anthropic-passthrough-endpoint",
            "messages": [{"role": "user", "content": "Hello"}],
            "max_tokens": 1024,
        }
    )

    mock_response = {
        "id": "msg_01XFDUDYJgAACzvnptvVoYEL",
        "type": "message",
        "role": "assistant",
        "content": [{"type": "text", "text": "Hello! How can I assist you today?"}],
        "model": "claude-3-5-sonnet-20241022",
        "stop_reason": "end_turn",
        "stop_sequence": None,
        "usage": {"input_tokens": 10, "output_tokens": 20},
    }

    with mock.patch(
        "mlflow.gateway.providers.anthropic.send_request", return_value=mock_response
    ) as mock_send:
        response = await anthropic_passthrough_messages(mock_request)

        assert mock_send.called
        call_args = mock_send.call_args
        assert call_args[1]["path"] == "messages"
        assert call_args[1]["payload"]["model"] == "claude-3-5-sonnet-20241022"
        assert call_args[1]["payload"]["messages"] == [{"role": "user", "content": "Hello"}]
        assert call_args[1]["payload"]["max_tokens"] == 1024

        assert response["id"] == "msg_01XFDUDYJgAACzvnptvVoYEL"
        assert response["model"] == "claude-3-5-sonnet-20241022"
        assert response["content"][0]["text"] == "Hello! How can I assist you today?"


@pytest.mark.asyncio
async def test_anthropic_passthrough_messages_streaming(store: SqlAlchemyStore):
    secret = store.create_gateway_secret(
        secret_name="anthropic-stream-passthrough-key",
        secret_value={"api_key": "sk-ant-test-stream"},
        provider="anthropic",
    )
    model_def = store.create_gateway_model_definition(
        name="anthropic-stream-passthrough-model",
        secret_id=secret.secret_id,
        provider="anthropic",
        model_name="claude-3-5-sonnet-20241022",
    )
    store.create_gateway_endpoint(
        name="anthropic-stream-passthrough-endpoint",
        model_configs=[
            GatewayEndpointModelConfig(
                model_definition_id=model_def.model_definition_id,
                linkage_type=GatewayModelLinkageType.PRIMARY,
                weight=1.0,
            ),
        ],
    )

    mock_request = create_mock_request()
    mock_request.json = AsyncMock(
        return_value={
            "model": "anthropic-stream-passthrough-endpoint",
            "messages": [{"role": "user", "content": "Hello"}],
            "max_tokens": 1024,
            "stream": True,
        }
    )

    mock_stream_chunks = [
        b'event: message_start\ndata: {"type":"message_start","message":{"id":"msg_01XFDUDYJgAACzvnptvVoYEL","type":"message","role":"assistant","content":[],"model":"claude-3-5-sonnet-20241022","stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":10,"output_tokens":0}}}\n\n',  # noqa: E501
        b'event: content_block_start\ndata: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}\n\n',  # noqa: E501
        b'event: content_block_delta\ndata: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hello"}}\n\n',  # noqa: E501
        b'event: content_block_delta\ndata: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"!"}}\n\n',  # noqa: E501
        b'event: content_block_stop\ndata: {"type":"content_block_stop","index":0}\n\n',
        b'event: message_delta\ndata: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":20}}\n\n',  # noqa: E501
        b'event: message_stop\ndata: {"type":"message_stop"}\n\n',
    ]

    async def mock_stream_generator():
        for chunk in mock_stream_chunks:
            yield chunk

    with mock.patch(
        "mlflow.gateway.providers.anthropic.send_stream_request",
        return_value=mock_stream_generator(),
    ) as mock_send_stream:
        response = await anthropic_passthrough_messages(mock_request)

        assert isinstance(response, StreamingResponse)
        assert response.media_type == "text/event-stream"

        chunks = [chunk async for chunk in response.body_iterator]

        assert mock_send_stream.called

        assert len(chunks) == 7
        assert b"message_start" in chunks[0]
        assert b"content_block_start" in chunks[1]
        assert b"content_block_delta" in chunks[2]
        assert b"Hello" in chunks[2]
        assert b"content_block_delta" in chunks[3]
        assert b"!" in chunks[3]
        assert b"content_block_stop" in chunks[4]
        assert b"message_delta" in chunks[5]
        assert b"message_stop" in chunks[6]


# =============================================================================
# Gemini generateContent/streamGenerateContent passthrough endpoint tests
# =============================================================================


@pytest.mark.asyncio
async def test_gemini_passthrough_generate_content(store: SqlAlchemyStore):
    secret = store.create_gateway_secret(
        secret_name="gemini-passthrough-key",
        secret_value={"api_key": "test-key"},
        provider="gemini",
    )
    model_def = store.create_gateway_model_definition(
        name="gemini-passthrough-model",
        secret_id=secret.secret_id,
        provider="gemini",
        model_name="gemini-2.0-flash",
    )
    store.create_gateway_endpoint(
        name="gemini-passthrough-endpoint",
        model_configs=[
            GatewayEndpointModelConfig(
                model_definition_id=model_def.model_definition_id,
                linkage_type=GatewayModelLinkageType.PRIMARY,
                weight=1.0,
            ),
        ],
    )

    mock_request = create_mock_request()
    mock_request.json = AsyncMock(
        return_value={
            "contents": [
                {
                    "role": "user",
                    "parts": [{"text": "Hello"}],
                }
            ]
        }
    )

    mock_response = {
        "candidates": [
            {
                "content": {
                    "parts": [{"text": "Hello! How can I assist you today?"}],
                    "role": "model",
                },
                "finishReason": "STOP",
            }
        ],
        "usageMetadata": {
            "promptTokenCount": 5,
            "candidatesTokenCount": 10,
            "totalTokenCount": 15,
        },
    }

    with mock.patch(
        "mlflow.gateway.providers.gemini.send_request", return_value=mock_response
    ) as mock_send:
        response = await gemini_passthrough_generate_content(
            "gemini-passthrough-endpoint", mock_request
        )

        assert mock_send.called
        call_args = mock_send.call_args
        assert call_args[1]["path"] == "gemini-2.0-flash:generateContent"
        assert call_args[1]["payload"]["contents"] == [
            {"role": "user", "parts": [{"text": "Hello"}]}
        ]

        assert (
            response["candidates"][0]["content"]["parts"][0]["text"]
            == "Hello! How can I assist you today?"
        )
        assert response["usageMetadata"]["totalTokenCount"] == 15


@pytest.mark.asyncio
async def test_gemini_passthrough_stream_generate_content(store: SqlAlchemyStore):
    secret = store.create_gateway_secret(
        secret_name="gemini-stream-passthrough-key",
        secret_value={"api_key": "test-stream-key"},
        provider="gemini",
    )
    model_def = store.create_gateway_model_definition(
        name="gemini-stream-passthrough-model",
        secret_id=secret.secret_id,
        provider="gemini",
        model_name="gemini-2.0-flash",
    )
    store.create_gateway_endpoint(
        name="gemini-stream-passthrough-endpoint",
        model_configs=[
            GatewayEndpointModelConfig(
                model_definition_id=model_def.model_definition_id,
                linkage_type=GatewayModelLinkageType.PRIMARY,
                weight=1.0,
            ),
        ],
    )

    mock_request = create_mock_request()
    mock_request.json = AsyncMock(
        return_value={
            "contents": [
                {
                    "role": "user",
                    "parts": [{"text": "Hello"}],
                }
            ]
        }
    )

    mock_stream_chunks = [
        b'data: {"candidates":[{"content":{"parts":[{"text":"Hello"}],"role":"model"}}]}\n\n',
        b'data: {"candidates":[{"content":{"parts":[{"text":"!"}],"role":"model"}}]}\n\n',
        b'data: {"candidates":[{"content":{"parts":[{"text":" How can I help you?"}],"role":"model"},"finishReason":"STOP"}]}\n\n',  # noqa: E501
    ]

    async def mock_stream_generator():
        for chunk in mock_stream_chunks:
            yield chunk

    with mock.patch(
        "mlflow.gateway.providers.gemini.send_stream_request",
        return_value=mock_stream_generator(),
    ) as mock_send_stream:
        response = await gemini_passthrough_stream_generate_content(
            "gemini-stream-passthrough-endpoint", mock_request
        )

        assert isinstance(response, StreamingResponse)
        assert response.media_type == "text/event-stream"

        chunks = [chunk async for chunk in response.body_iterator]

        assert mock_send_stream.called

        assert len(chunks) == 3
        assert b"Hello" in chunks[0]
        assert b"!" in chunks[1]
        assert b"How can I help you?" in chunks[2]
        assert b"STOP" in chunks[2]


def test_create_fallback_provider_single_model(store: SqlAlchemyStore):
    secret = store.create_gateway_secret(
        secret_name="openai-fallback-key",
        secret_value={"api_key": "sk-test-key"},
        provider="openai",
    )
    model_def = store.create_gateway_model_definition(
        name="gpt-fallback-model",
        secret_id=secret.secret_id,
        provider="openai",
        model_name="gpt-4",
    )
    endpoint = store.create_gateway_endpoint(
        name="test-fallback-single-endpoint",
        model_configs=[
            GatewayEndpointModelConfig(
                model_definition_id=model_def.model_definition_id,
                linkage_type=GatewayModelLinkageType.PRIMARY,
                weight=1.0,
            ),
            GatewayEndpointModelConfig(
                model_definition_id=model_def.model_definition_id,
                linkage_type=GatewayModelLinkageType.FALLBACK,
                weight=1.0,
                fallback_order=0,
            ),
        ],
        routing_strategy=RoutingStrategy.REQUEST_BASED_TRAFFIC_SPLIT,
        fallback_config=FallbackConfig(
            strategy=FallbackStrategy.SEQUENTIAL,
            max_attempts=1,
        ),
    )

    provider, _ = _create_provider_from_endpoint_name(
        store, endpoint.name, EndpointType.LLM_V1_CHAT
    )

    assert isinstance(provider, FallbackProvider)
    assert len(provider._providers) == 2
    assert isinstance(provider._providers[0], TrafficRouteProvider)
    assert isinstance(provider._providers[1], OpenAIProvider)
    assert provider._max_attempts == 2


def test_create_fallback_provider_multiple_models(store: SqlAlchemyStore):
    secret1 = store.create_gateway_secret(
        secret_name="openai-primary-key",
        secret_value={"api_key": "sk-primary-key"},
        provider="openai",
    )
    model_def1 = store.create_gateway_model_definition(
        name="gpt-primary-model",
        secret_id=secret1.secret_id,
        provider="openai",
        model_name="gpt-4",
    )

    secret2 = store.create_gateway_secret(
        secret_name="anthropic-fallback-key",
        secret_value={"api_key": "sk-ant-fallback"},
        provider="anthropic",
    )
    model_def2 = store.create_gateway_model_definition(
        name="claude-fallback-model",
        secret_id=secret2.secret_id,
        provider="anthropic",
        model_name="claude-3-sonnet",
    )

    endpoint = store.create_gateway_endpoint(
        name="test-fallback-multi-endpoint",
        model_configs=[
            GatewayEndpointModelConfig(
                model_definition_id=model_def1.model_definition_id,
                linkage_type=GatewayModelLinkageType.PRIMARY,
                weight=1.0,
            ),
            GatewayEndpointModelConfig(
                model_definition_id=model_def2.model_definition_id,
                linkage_type=GatewayModelLinkageType.PRIMARY,
                weight=1.0,
            ),
            GatewayEndpointModelConfig(
                model_definition_id=model_def1.model_definition_id,
                linkage_type=GatewayModelLinkageType.FALLBACK,
                weight=1.0,
                fallback_order=0,
            ),
            GatewayEndpointModelConfig(
                model_definition_id=model_def2.model_definition_id,
                linkage_type=GatewayModelLinkageType.FALLBACK,
                weight=1.0,
                fallback_order=1,
            ),
        ],
        routing_strategy=RoutingStrategy.REQUEST_BASED_TRAFFIC_SPLIT,
        fallback_config=FallbackConfig(
            strategy=FallbackStrategy.SEQUENTIAL,
            max_attempts=2,
        ),
    )

    provider, _ = _create_provider_from_endpoint_name(
        store, endpoint.name, EndpointType.LLM_V1_CHAT
    )

    assert isinstance(provider, FallbackProvider)
    assert len(provider._providers) == 3
    primary = provider._providers[0]
    assert isinstance(primary, TrafficRouteProvider)
    assert isinstance(primary._providers[0], OpenAIProvider)
    assert isinstance(primary._providers[1], AnthropicProvider)
    assert isinstance(provider._providers[1], OpenAIProvider)
    assert isinstance(provider._providers[2], AnthropicProvider)
    assert provider._max_attempts == 3


def test_create_fallback_provider_max_attempts_exceeds_providers(store: SqlAlchemyStore):
    secret = store.create_gateway_secret(
        secret_name="openai-fallback-key",
        secret_value={"api_key": "sk-test-key"},
        provider="openai",
    )
    model_def = store.create_gateway_model_definition(
        name="gpt-fallback-model",
        secret_id=secret.secret_id,
        provider="openai",
        model_name="gpt-4",
    )
    endpoint = store.create_gateway_endpoint(
        name="test-fallback-max-attempts-endpoint",
        model_configs=[
            GatewayEndpointModelConfig(
                model_definition_id=model_def.model_definition_id,
                linkage_type=GatewayModelLinkageType.PRIMARY,
                weight=1.0,
            ),
            GatewayEndpointModelConfig(
                model_definition_id=model_def.model_definition_id,
                linkage_type=GatewayModelLinkageType.FALLBACK,
                weight=1.0,
                fallback_order=0,
            ),
        ],
        routing_strategy=RoutingStrategy.REQUEST_BASED_TRAFFIC_SPLIT,
        fallback_config=FallbackConfig(
            strategy=FallbackStrategy.SEQUENTIAL,
            max_attempts=10,
        ),
    )

    provider, _ = _create_provider_from_endpoint_name(
        store, endpoint.name, EndpointType.LLM_V1_CHAT
    )

    # FallbackProvider is the outer provider, individual providers inside are wrapped
    assert isinstance(provider, FallbackProvider)
    assert provider._max_attempts == 2


def test_create_fallback_provider_no_max_attempts(store: SqlAlchemyStore):
    secret1 = store.create_gateway_secret(
        secret_name="openai-primary-key",
        secret_value={"api_key": "sk-primary-key"},
        provider="openai",
    )
    model_def1 = store.create_gateway_model_definition(
        name="gpt-primary-model",
        secret_id=secret1.secret_id,
        provider="openai",
        model_name="gpt-4",
    )

    secret2 = store.create_gateway_secret(
        secret_name="anthropic-fallback-key",
        secret_value={"api_key": "sk-ant-fallback"},
        provider="anthropic",
    )
    model_def2 = store.create_gateway_model_definition(
        name="claude-fallback-model",
        secret_id=secret2.secret_id,
        provider="anthropic",
        model_name="claude-3-sonnet",
    )

    endpoint = store.create_gateway_endpoint(
        name="test-fallback-no-max-endpoint",
        model_configs=[
            GatewayEndpointModelConfig(
                model_definition_id=model_def1.model_definition_id,
                linkage_type=GatewayModelLinkageType.PRIMARY,
                weight=1.0,
            ),
            GatewayEndpointModelConfig(
                model_definition_id=model_def2.model_definition_id,
                linkage_type=GatewayModelLinkageType.PRIMARY,
                weight=1.0,
            ),
            GatewayEndpointModelConfig(
                model_definition_id=model_def1.model_definition_id,
                linkage_type=GatewayModelLinkageType.FALLBACK,
                weight=1.0,
                fallback_order=0,
            ),
            GatewayEndpointModelConfig(
                model_definition_id=model_def2.model_definition_id,
                linkage_type=GatewayModelLinkageType.FALLBACK,
                weight=1.0,
                fallback_order=1,
            ),
        ],
        routing_strategy=RoutingStrategy.REQUEST_BASED_TRAFFIC_SPLIT,
        fallback_config=FallbackConfig(
            strategy=FallbackStrategy.SEQUENTIAL,
            max_attempts=None,
        ),
    )

    provider, _ = _create_provider_from_endpoint_name(
        store, endpoint.name, EndpointType.LLM_V1_CHAT
    )

    # FallbackProvider is the outer provider, individual providers inside are wrapped
    assert isinstance(provider, FallbackProvider)
    assert len(provider._providers) == 3
    assert provider._max_attempts == 3


def test_create_provider_default_routing_single_model(store: SqlAlchemyStore):
    secret = store.create_gateway_secret(
        secret_name="openai-default-key",
        secret_value={"api_key": "sk-test-key"},
        provider="openai",
    )
    model_def = store.create_gateway_model_definition(
        name="gpt-default-model",
        secret_id=secret.secret_id,
        provider="openai",
        model_name="gpt-4",
    )
    endpoint = store.create_gateway_endpoint(
        name="test-default-routing-endpoint",
        model_configs=[
            GatewayEndpointModelConfig(
                model_definition_id=model_def.model_definition_id,
                linkage_type=GatewayModelLinkageType.PRIMARY,
                weight=1.0,
            ),
        ],
    )

    provider, _ = _create_provider_from_endpoint_name(
        store, endpoint.name, EndpointType.LLM_V1_CHAT
    )

    assert isinstance(provider, OpenAIProvider)
    assert not isinstance(provider, FallbackProvider)


# =============================================================================
# Gateway Tracing Tests
# =============================================================================


async def _call_invocations(endpoint_name: str, request, payload: dict[str, Any]):
    # invocations doesn't use "model" field - endpoint is in URL
    payload_without_model = {k: v for k, v in payload.items() if k != "model"}
    request.json = AsyncMock(return_value=payload_without_model)
    return await invocations(endpoint_name, request)


async def _call_chat_completions(endpoint_name: str, request, payload: dict[str, Any]):
    # chat_completions uses "model" field to specify endpoint
    request.json = AsyncMock(return_value=payload)
    return await chat_completions(request)


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "handler", [_call_invocations, _call_chat_completions], ids=["invocations", "chat_completions"]
)
async def test_gateway_creates_trace_with_usage(store: SqlAlchemyStore, handler):
    endpoint_name = "tracing-test-endpoint"

    # Create experiment for tracing
    experiment_id = store.create_experiment(f"gateway/{endpoint_name}")

    # Create endpoint with usage tracking enabled
    secret = store.create_gateway_secret(
        secret_name="tracing-test-key",
        secret_value={"api_key": "sk-test-tracing"},
        provider="openai",
    )
    model_def = store.create_gateway_model_definition(
        name="tracing-test-model",
        secret_id=secret.secret_id,
        provider="openai",
        model_name="gpt-4",
    )
    endpoint = store.create_gateway_endpoint(
        name=endpoint_name,
        model_configs=[
            GatewayEndpointModelConfig(
                model_definition_id=model_def.model_definition_id,
                linkage_type=GatewayModelLinkageType.PRIMARY,
                weight=1.0,
            ),
        ],
        usage_tracking=True,
        experiment_id=experiment_id,
    )

    mock_request = create_mock_request()
    payload = {
        "model": endpoint_name,
        "messages": [{"role": "user", "content": "Hi"}],
        "stream": False,
    }

    # Mock the OpenAI send_request to return our mock response
    with mock.patch(
        "mlflow.gateway.providers.openai.send_request",
        return_value={
            "id": "test-id",
            "object": "chat.completion",
            "created": 1234567890,
            "model": "gpt-4",
            "choices": [
                {
                    "index": 0,
                    "message": {"role": "assistant", "content": "Hello!"},
                    "finish_reason": "stop",
                }
            ],
            "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
        },
    ):
        response = await handler(endpoint_name, mock_request, payload)

        assert response.id == "test-id"
        assert response.choices[0].message.content == "Hello!"

    # Verify trace was created
    traces = TracingClient().search_traces(locations=[experiment_id])
    assert len(traces) == 1
    trace = traces[0]
    assert trace.info.state == TraceState.OK

    # Verify gateway metadata is present in trace
    assert (
        trace.info.request_metadata.get(TraceMetadataKey.GATEWAY_ENDPOINT_ID)
        == endpoint.endpoint_id
    )
    assert (
        trace.info.request_metadata.get(TraceMetadataKey.GATEWAY_REQUEST_TYPE)
        == GatewayRequestType.UNIFIED_CHAT
    )

    # Verify span has provider information (provider name is lowercased for cost lookup alignment)
    span_names = {span.name for span in trace.data.spans}
    assert "provider/openai/gpt-4" in span_names

    # Find the provider span and check attributes
    provider_span = next(
        (span for span in trace.data.spans if span.name == "provider/openai/gpt-4"), None
    )
    assert provider_span is not None
    assert provider_span.attributes.get(SpanAttributeKey.MODEL_PROVIDER) == "openai"
    assert provider_span.attributes.get(SpanAttributeKey.MODEL) == "gpt-4"

    # Verify token usage is captured on the provider span
    token_usage = provider_span.attributes.get(SpanAttributeKey.CHAT_USAGE)
    assert token_usage is not None
    assert token_usage[TokenUsageKey.INPUT_TOKENS] == 10
    assert token_usage[TokenUsageKey.OUTPUT_TOKENS] == 5
    assert token_usage[TokenUsageKey.TOTAL_TOKENS] == 15

    # Verify trace metadata has aggregated token usage (auto-generated from span attributes)
    trace_token_usage = json.loads(trace.info.trace_metadata.get(TraceMetadataKey.TOKEN_USAGE))
    assert trace_token_usage[TokenUsageKey.INPUT_TOKENS] == 10
    assert trace_token_usage[TokenUsageKey.OUTPUT_TOKENS] == 5
    assert trace_token_usage[TokenUsageKey.TOTAL_TOKENS] == 15


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "handler", [_call_invocations, _call_chat_completions], ids=["invocations", "chat_completions"]
)
async def test_gateway_streaming_creates_trace(store: SqlAlchemyStore, handler):
    endpoint_name = "stream-tracing-test-endpoint"

    # Create experiment for tracing
    experiment_id = store.create_experiment(f"gateway/{endpoint_name}")

    # Create endpoint with usage tracking enabled
    secret = store.create_gateway_secret(
        secret_name="stream-tracing-test-key",
        secret_value={"api_key": "sk-test-stream-tracing"},
        provider="openai",
    )
    model_def = store.create_gateway_model_definition(
        name="stream-tracing-test-model",
        secret_id=secret.secret_id,
        provider="openai",
        model_name="gpt-4",
    )
    endpoint = store.create_gateway_endpoint(
        name=endpoint_name,
        model_configs=[
            GatewayEndpointModelConfig(
                model_definition_id=model_def.model_definition_id,
                linkage_type=GatewayModelLinkageType.PRIMARY,
                weight=1.0,
            ),
        ],
        usage_tracking=True,
        experiment_id=experiment_id,
    )

    mock_request = create_mock_request()
    payload = {
        "model": endpoint_name,
        "messages": [{"role": "user", "content": "Hi"}],
        "stream": True,
    }

    # Mock streaming response chunks with usage in the final chunk
    mock_stream_chunks = [
        b'data: {"id":"chatcmpl-123","object":"chat.completion.chunk","created":1234567890,"model":"gpt-4","choices":[{"index":0,"delta":{"role":"assistant","content":"Hello"},"finish_reason":null}]}\n\n',  # noqa: E501
        b'data: {"id":"chatcmpl-123","object":"chat.completion.chunk","created":1234567890,"model":"gpt-4","choices":[{"index":0,"delta":{"content":"!"},"finish_reason":null}]}\n\n',  # noqa: E501
        b'data: {"id":"chatcmpl-123","object":"chat.completion.chunk","created":1234567890,"model":"gpt-4","choices":[{"index":0,"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":10,"completion_tokens":5,"total_tokens":15}}\n\n',  # noqa: E501
        b"data: [DONE]\n\n",
    ]

    async def mock_stream_generator():
        for chunk in mock_stream_chunks:
            yield chunk

    with mock.patch(
        "mlflow.gateway.providers.openai.send_stream_request",
        return_value=mock_stream_generator(),
    ):
        response = await handler(endpoint_name, mock_request, payload)

        # Verify streaming response is returned
        assert isinstance(response, StreamingResponse)
        # Consume the response
        chunks = [chunk async for chunk in response.body_iterator]
        assert len(chunks) > 0

    # Verify trace was created for the gateway invocation
    traces = TracingClient().search_traces(locations=[experiment_id])
    assert len(traces) == 1
    trace = traces[0]
    assert trace.info.state == TraceState.OK

    # Verify gateway metadata is present in trace
    assert (
        trace.info.request_metadata.get(TraceMetadataKey.GATEWAY_ENDPOINT_ID)
        == endpoint.endpoint_id
    )
    assert (
        trace.info.request_metadata.get(TraceMetadataKey.GATEWAY_REQUEST_TYPE)
        == GatewayRequestType.UNIFIED_CHAT
    )

    # Verify gateway span exists with correct attributes
    gateway_span = next(
        (span for span in trace.data.spans if span.name == f"gateway/{endpoint_name}"), None
    )
    assert gateway_span is not None
    assert gateway_span.attributes.get("endpoint_name") == endpoint_name

    # Verify that streaming output is aggregated into a ChatCompletion-like response
    output = gateway_span.outputs
    assert output is not None
    assert output["object"] == "chat.completion"
    assert output["id"] == "chatcmpl-123"
    assert output["model"] == "gpt-4"
    assert len(output["choices"]) == 1
    assert output["choices"][0]["index"] == 0
    assert output["choices"][0]["message"]["role"] == "assistant"
    assert output["choices"][0]["message"]["content"] == "Hello!"
    assert output["choices"][0]["finish_reason"] == "stop"
    assert output["usage"]["prompt_tokens"] == 10
    assert output["usage"]["completion_tokens"] == 5
    assert output["usage"]["total_tokens"] == 15


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "handler", [_call_invocations, _call_chat_completions], ids=["invocations", "chat_completions"]
)
async def test_gateway_trace_includes_user_metadata(store: SqlAlchemyStore, handler):
    endpoint_name = "user-metadata-tracing-endpoint"

    # Create experiment for tracing
    experiment_id = store.create_experiment(f"gateway/{endpoint_name}")

    # Create endpoint with usage tracking enabled
    secret = store.create_gateway_secret(
        secret_name="user-metadata-tracing-key",
        secret_value={"api_key": "sk-test-user-metadata"},
        provider="openai",
    )
    model_def = store.create_gateway_model_definition(
        name="user-metadata-tracing-model",
        secret_id=secret.secret_id,
        provider="openai",
        model_name="gpt-4",
    )
    endpoint = store.create_gateway_endpoint(
        name=endpoint_name,
        model_configs=[
            GatewayEndpointModelConfig(
                model_definition_id=model_def.model_definition_id,
                linkage_type=GatewayModelLinkageType.PRIMARY,
                weight=1.0,
            ),
        ],
        usage_tracking=True,
        experiment_id=experiment_id,
    )

    # Create mock request with user metadata set (as auth middleware would do)
    mock_request = create_mock_request(username="test_user", user_id=42)
    payload = {
        "model": endpoint_name,
        "messages": [{"role": "user", "content": "Hi"}],
        "stream": False,
    }

    # Mock the OpenAI send_request to return our mock response
    with mock.patch(
        "mlflow.gateway.providers.openai.send_request",
        return_value={
            "id": "test-id",
            "object": "chat.completion",
            "created": 1234567890,
            "model": "gpt-4",
            "choices": [
                {
                    "index": 0,
                    "message": {"role": "assistant", "content": "Hello!"},
                    "finish_reason": "stop",
                }
            ],
            "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
        },
    ):
        response = await handler(endpoint_name, mock_request, payload)

        assert response.id == "test-id"
        assert response.choices[0].message.content == "Hello!"

    # Verify trace was created
    traces = TracingClient().search_traces(locations=[experiment_id])
    assert len(traces) == 1
    trace = traces[0]
    assert trace.info.state == TraceState.OK

    # Verify user metadata is present in trace info
    assert trace.info.request_metadata.get(TraceMetadataKey.AUTH_USERNAME) == "test_user"
    assert trace.info.request_metadata.get(TraceMetadataKey.AUTH_USER_ID) == "42"

    # Verify gateway metadata is present alongside user metadata
    assert (
        trace.info.request_metadata.get(TraceMetadataKey.GATEWAY_ENDPOINT_ID)
        == endpoint.endpoint_id
    )
    assert (
        trace.info.request_metadata.get(TraceMetadataKey.GATEWAY_REQUEST_TYPE)
        == GatewayRequestType.UNIFIED_CHAT
    )

    # Verify span attributes still include endpoint info
    gateway_span = next(
        (span for span in trace.data.spans if span.name == f"gateway/{endpoint_name}"), None
    )
    assert gateway_span is not None
    assert gateway_span.attributes.get("endpoint_name") == endpoint_name


# =============================================================================
# Passthrough Token Usage Tracking Tests
# =============================================================================


@pytest.mark.asyncio
async def test_openai_passthrough_chat_token_usage_tracking(store: SqlAlchemyStore):
    endpoint_name = "openai-passthrough-usage-endpoint"

    # Create experiment for tracing
    experiment_id = store.create_experiment(f"gateway/{endpoint_name}")

    secret = store.create_gateway_secret(
        secret_name="openai-passthrough-usage-key",
        secret_value={"api_key": "sk-test-usage"},
        provider="openai",
    )
    model_def = store.create_gateway_model_definition(
        name="openai-passthrough-usage-model",
        secret_id=secret.secret_id,
        provider="openai",
        model_name="gpt-4o",
    )
    endpoint = store.create_gateway_endpoint(
        name=endpoint_name,
        model_configs=[
            GatewayEndpointModelConfig(
                model_definition_id=model_def.model_definition_id,
                linkage_type=GatewayModelLinkageType.PRIMARY,
                weight=1.0,
            ),
        ],
        usage_tracking=True,
        experiment_id=experiment_id,
    )

    mock_request = create_mock_request()
    mock_request.json = AsyncMock(
        return_value={
            "model": endpoint_name,
            "messages": [{"role": "user", "content": "Hello"}],
        }
    )

    mock_response = {
        "id": "chatcmpl-123",
        "object": "chat.completion",
        "created": 1234567890,
        "model": "gpt-4o",
        "choices": [
            {
                "index": 0,
                "message": {"role": "assistant", "content": "Hello!"},
                "finish_reason": "stop",
            }
        ],
        "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
    }

    with mock.patch("mlflow.gateway.providers.openai.send_request", return_value=mock_response):
        response = await openai_passthrough_chat(mock_request)
        assert response["usage"]["prompt_tokens"] == 10
        assert response["usage"]["completion_tokens"] == 5
        assert response["usage"]["total_tokens"] == 15

    # Verify trace was created with token usage
    traces = TracingClient().search_traces(locations=[experiment_id])
    assert len(traces) == 1
    trace = traces[0]
    assert trace.info.state == TraceState.OK

    # Verify gateway metadata
    assert (
        trace.info.request_metadata.get(TraceMetadataKey.GATEWAY_ENDPOINT_ID)
        == endpoint.endpoint_id
    )
    assert (
        trace.info.request_metadata.get(TraceMetadataKey.GATEWAY_REQUEST_TYPE)
        == GatewayRequestType.PASSTHROUGH_MODEL_OPENAI_CHAT
    )

    # Find the passthrough span and check token usage attributes
    passthrough_span = next(
        (span for span in trace.data.spans if "action" in span.attributes), None
    )
    assert passthrough_span is not None
    assert passthrough_span.attributes.get("action") == "openai_chat"

    token_usage = passthrough_span.attributes.get(SpanAttributeKey.CHAT_USAGE)
    assert token_usage is not None
    assert token_usage[TokenUsageKey.INPUT_TOKENS] == 10
    assert token_usage[TokenUsageKey.OUTPUT_TOKENS] == 5
    assert token_usage[TokenUsageKey.TOTAL_TOKENS] == 15


@pytest.mark.asyncio
async def test_openai_passthrough_embeddings_token_usage_tracking(store: SqlAlchemyStore):
    endpoint_name = "openai-embed-usage-endpoint"

    # Create experiment for tracing
    experiment_id = store.create_experiment(f"gateway/{endpoint_name}")

    secret = store.create_gateway_secret(
        secret_name="openai-embed-usage-key",
        secret_value={"api_key": "sk-test-embed-usage"},
        provider="openai",
    )
    model_def = store.create_gateway_model_definition(
        name="openai-embed-usage-model",
        secret_id=secret.secret_id,
        provider="openai",
        model_name="text-embedding-3-small",
    )
    endpoint = store.create_gateway_endpoint(
        name=endpoint_name,
        model_configs=[
            GatewayEndpointModelConfig(
                model_definition_id=model_def.model_definition_id,
                linkage_type=GatewayModelLinkageType.PRIMARY,
                weight=1.0,
            ),
        ],
        usage_tracking=True,
        experiment_id=experiment_id,
    )

    mock_request = create_mock_request()
    mock_request.json = AsyncMock(
        return_value={
            "model": endpoint_name,
            "input": "Test text for embedding",
        }
    )

    mock_response = {
        "object": "list",
        "data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3]}],
        "model": "text-embedding-3-small",
        "usage": {"prompt_tokens": 5, "total_tokens": 5},
    }

    with mock.patch("mlflow.gateway.providers.openai.send_request", return_value=mock_response):
        response = await openai_passthrough_embeddings(mock_request)
        assert response["usage"]["prompt_tokens"] == 5
        assert response["usage"]["total_tokens"] == 5

    # Verify trace was created with token usage
    traces = TracingClient().search_traces(locations=[experiment_id])
    assert len(traces) == 1
    trace = traces[0]

    # Verify gateway metadata
    assert (
        trace.info.request_metadata.get(TraceMetadataKey.GATEWAY_ENDPOINT_ID)
        == endpoint.endpoint_id
    )
    assert (
        trace.info.request_metadata.get(TraceMetadataKey.GATEWAY_REQUEST_TYPE)
        == GatewayRequestType.PASSTHROUGH_MODEL_OPENAI_EMBEDDINGS
    )

    # Find the passthrough span and check token usage attributes
    passthrough_span = next(
        (span for span in trace.data.spans if "action" in span.attributes), None
    )
    assert passthrough_span is not None
    assert passthrough_span.attributes.get("action") == "openai_embeddings"

    token_usage = passthrough_span.attributes.get(SpanAttributeKey.CHAT_USAGE)
    assert token_usage is not None
    assert token_usage[TokenUsageKey.INPUT_TOKENS] == 5
    assert token_usage[TokenUsageKey.TOTAL_TOKENS] == 5


@pytest.mark.asyncio
async def test_openai_passthrough_responses_token_usage_tracking(store: SqlAlchemyStore):
    endpoint_name = "openai-responses-usage-endpoint"

    # Create experiment for tracing
    experiment_id = store.create_experiment(f"gateway/{endpoint_name}")

    secret = store.create_gateway_secret(
        secret_name="openai-responses-usage-key",
        secret_value={"api_key": "sk-test-responses-usage"},
        provider="openai",
    )
    model_def = store.create_gateway_model_definition(
        name="openai-responses-usage-model",
        secret_id=secret.secret_id,
        provider="openai",
        model_name="gpt-4o",
    )
    endpoint = store.create_gateway_endpoint(
        name=endpoint_name,
        model_configs=[
            GatewayEndpointModelConfig(
                model_definition_id=model_def.model_definition_id,
                linkage_type=GatewayModelLinkageType.PRIMARY,
                weight=1.0,
            ),
        ],
        usage_tracking=True,
        experiment_id=experiment_id,
    )

    mock_request = create_mock_request()
    mock_request.json = AsyncMock(
        return_value={
            "model": endpoint_name,
            "input": [{"role": "user", "content": "Hello"}],
        }
    )

    mock_response = {
        "id": "resp-123",
        "object": "response",
        "created": 1234567890,
        "model": "gpt-4o",
        "status": "completed",
        "output": [
            {
                "role": "assistant",
                "content": [{"type": "output_text", "text": "Hello!"}],
            }
        ],
        "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15},
    }

    with mock.patch("mlflow.gateway.providers.openai.send_request", return_value=mock_response):
        response = await openai_passthrough_responses(mock_request)
        assert response["usage"]["input_tokens"] == 10
        assert response["usage"]["output_tokens"] == 5
        assert response["usage"]["total_tokens"] == 15

    # Verify trace was created with token usage
    traces = TracingClient().search_traces(locations=[experiment_id])
    assert len(traces) == 1
    trace = traces[0]
    assert trace.info.state == TraceState.OK

    # Verify gateway metadata
    assert (
        trace.info.request_metadata.get(TraceMetadataKey.GATEWAY_ENDPOINT_ID)
        == endpoint.endpoint_id
    )
    assert (
        trace.info.request_metadata.get(TraceMetadataKey.GATEWAY_REQUEST_TYPE)
        == GatewayRequestType.PASSTHROUGH_MODEL_OPENAI_RESPONSES
    )

    # Find the passthrough span and check token usage attributes
    passthrough_span = next(
        (span for span in trace.data.spans if "action" in span.attributes), None
    )
    assert passthrough_span is not None
    assert passthrough_span.attributes.get("action") == "openai_responses"

    token_usage = passthrough_span.attributes.get(SpanAttributeKey.CHAT_USAGE)
    assert token_usage is not None
    assert token_usage[TokenUsageKey.INPUT_TOKENS] == 10
    assert token_usage[TokenUsageKey.OUTPUT_TOKENS] == 5
    assert token_usage[TokenUsageKey.TOTAL_TOKENS] == 15


@pytest.mark.asyncio
async def test_anthropic_passthrough_messages_token_usage_tracking(store: SqlAlchemyStore):
    endpoint_name = "anthropic-usage-endpoint"

    # Create experiment for tracing
    experiment_id = store.create_experiment(f"gateway/{endpoint_name}")

    secret = store.create_gateway_secret(
        secret_name="anthropic-usage-key",
        secret_value={"api_key": "sk-ant-usage"},
        provider="anthropic",
    )
    model_def = store.create_gateway_model_definition(
        name="anthropic-usage-model",
        secret_id=secret.secret_id,
        provider="anthropic",
        model_name="claude-3-5-sonnet-20241022",
    )
    endpoint = store.create_gateway_endpoint(
        name=endpoint_name,
        model_configs=[
            GatewayEndpointModelConfig(
                model_definition_id=model_def.model_definition_id,
                linkage_type=GatewayModelLinkageType.PRIMARY,
                weight=1.0,
            ),
        ],
        usage_tracking=True,
        experiment_id=experiment_id,
    )

    mock_request = create_mock_request()
    mock_request.json = AsyncMock(
        return_value={
            "model": endpoint_name,
            "messages": [{"role": "user", "content": "Hello"}],
            "max_tokens": 1024,
        }
    )

    mock_response = {
        "id": "msg_01XFDUDYJgAACzvnptvVoYEL",
        "type": "message",
        "role": "assistant",
        "content": [{"type": "text", "text": "Hello!"}],
        "model": "claude-3-5-sonnet-20241022",
        "stop_reason": "end_turn",
        "stop_sequence": None,
        "usage": {"input_tokens": 12, "output_tokens": 8},
    }

    with mock.patch("mlflow.gateway.providers.anthropic.send_request", return_value=mock_response):
        response = await anthropic_passthrough_messages(mock_request)
        assert response["usage"]["input_tokens"] == 12
        assert response["usage"]["output_tokens"] == 8

    # Verify trace was created with token usage
    traces = TracingClient().search_traces(locations=[experiment_id])
    assert len(traces) == 1
    trace = traces[0]

    # Verify gateway metadata
    assert (
        trace.info.request_metadata.get(TraceMetadataKey.GATEWAY_ENDPOINT_ID)
        == endpoint.endpoint_id
    )
    assert (
        trace.info.request_metadata.get(TraceMetadataKey.GATEWAY_REQUEST_TYPE)
        == GatewayRequestType.PASSTHROUGH_MODEL_ANTHROPIC_MESSAGES
    )

    # Find the passthrough span and check token usage attributes
    passthrough_span = next(
        (span for span in trace.data.spans if "action" in span.attributes), None
    )
    assert passthrough_span is not None
    assert passthrough_span.attributes.get("action") == "anthropic_messages"

    token_usage = passthrough_span.attributes.get(SpanAttributeKey.CHAT_USAGE)
    assert token_usage is not None
    assert token_usage[TokenUsageKey.INPUT_TOKENS] == 12
    assert token_usage[TokenUsageKey.OUTPUT_TOKENS] == 8
    # Anthropic doesn't provide total_tokens, so we calculate it
    assert token_usage[TokenUsageKey.TOTAL_TOKENS] == 20


@pytest.mark.asyncio
async def test_gemini_passthrough_generate_content_token_usage_tracking(store: SqlAlchemyStore):
    endpoint_name = "gemini-usage-endpoint"

    # Create experiment for tracing
    experiment_id = store.create_experiment(f"gateway/{endpoint_name}")

    secret = store.create_gateway_secret(
        secret_name="gemini-usage-key",
        secret_value={"api_key": "test-gemini-usage-key"},
        provider="gemini",
    )
    model_def = store.create_gateway_model_definition(
        name="gemini-usage-model",
        secret_id=secret.secret_id,
        provider="gemini",
        model_name="gemini-2.0-flash",
    )
    endpoint = store.create_gateway_endpoint(
        name=endpoint_name,
        model_configs=[
            GatewayEndpointModelConfig(
                model_definition_id=model_def.model_definition_id,
                linkage_type=GatewayModelLinkageType.PRIMARY,
                weight=1.0,
            ),
        ],
        usage_tracking=True,
        experiment_id=experiment_id,
    )

    mock_request = create_mock_request()
    mock_request.json = AsyncMock(
        return_value={
            "contents": [
                {
                    "role": "user",
                    "parts": [{"text": "Hello"}],
                }
            ]
        }
    )

    mock_response = {
        "candidates": [
            {
                "content": {
                    "parts": [{"text": "Hello! How can I help?"}],
                    "role": "model",
                },
                "finishReason": "STOP",
            }
        ],
        "usageMetadata": {
            "promptTokenCount": 7,
            "candidatesTokenCount": 9,
            "totalTokenCount": 16,
        },
    }

    with mock.patch("mlflow.gateway.providers.gemini.send_request", return_value=mock_response):
        response = await gemini_passthrough_generate_content(endpoint_name, mock_request)
        assert response["usageMetadata"]["promptTokenCount"] == 7
        assert response["usageMetadata"]["candidatesTokenCount"] == 9
        assert response["usageMetadata"]["totalTokenCount"] == 16

    # Verify trace was created with token usage
    traces = TracingClient().search_traces(locations=[experiment_id])
    assert len(traces) == 1
    trace = traces[0]

    # Verify gateway metadata
    assert (
        trace.info.request_metadata.get(TraceMetadataKey.GATEWAY_ENDPOINT_ID)
        == endpoint.endpoint_id
    )
    assert (
        trace.info.request_metadata.get(TraceMetadataKey.GATEWAY_REQUEST_TYPE)
        == GatewayRequestType.PASSTHROUGH_MODEL_GEMINI_GENERATE_CONTENT
    )

    # Find the passthrough span and check token usage attributes
    passthrough_span = next(
        (span for span in trace.data.spans if "action" in span.attributes), None
    )
    assert passthrough_span is not None
    assert passthrough_span.attributes.get("action") == "gemini_generate_content"

    token_usage = passthrough_span.attributes.get(SpanAttributeKey.CHAT_USAGE)
    assert token_usage is not None
    assert token_usage[TokenUsageKey.INPUT_TOKENS] == 7
    assert token_usage[TokenUsageKey.OUTPUT_TOKENS] == 9
    assert token_usage[TokenUsageKey.TOTAL_TOKENS] == 16


@pytest.mark.asyncio
async def test_openai_passthrough_streaming_captures_chunks(store: SqlAlchemyStore):
    endpoint_name = "openai-passthrough-streaming-chunks"

    experiment_id = store.create_experiment(f"gateway/{endpoint_name}")

    secret = store.create_gateway_secret(
        secret_name="openai-stream-chunks-key",
        secret_value={"api_key": "sk-test-stream-chunks"},
        provider="openai",
    )
    model_def = store.create_gateway_model_definition(
        name="openai-stream-chunks-model",
        secret_id=secret.secret_id,
        provider="openai",
        model_name="gpt-4o",
    )
    endpoint = store.create_gateway_endpoint(
        name=endpoint_name,
        model_configs=[
            GatewayEndpointModelConfig(
                model_definition_id=model_def.model_definition_id,
                linkage_type=GatewayModelLinkageType.PRIMARY,
                weight=1.0,
            ),
        ],
        usage_tracking=True,
        experiment_id=experiment_id,
    )

    mock_request = create_mock_request()
    mock_request.json = AsyncMock(
        return_value={
            "model": endpoint_name,
            "messages": [{"role": "user", "content": "Hello"}],
            "stream": True,
        }
    )
    mock_request.headers = {}

    mock_stream_chunks = [
        b'data: {"id":"chatcmpl-123","choices":[{"delta":{"content":"Hi"}}]}\n\n',
        b'data: {"id":"chatcmpl-123","choices":[{"delta":{"content":"!"}}]}\n\n',
        b"data: [DONE]\n\n",
    ]

    async def mock_stream_generator():
        for chunk in mock_stream_chunks:
            yield chunk

    with mock.patch(
        "mlflow.gateway.providers.openai.send_stream_request",
        return_value=mock_stream_generator(),
    ):
        response = await openai_passthrough_chat(mock_request)
        assert isinstance(response, StreamingResponse)
        chunks = [chunk async for chunk in response.body_iterator]
        assert len(chunks) == len(mock_stream_chunks)

    # Verify trace was created
    traces = TracingClient().search_traces(locations=[experiment_id])
    assert len(traces) == 1
    trace = traces[0]
    assert trace.info.state == TraceState.OK

    # Verify gateway metadata
    assert (
        trace.info.request_metadata.get(TraceMetadataKey.GATEWAY_ENDPOINT_ID)
        == endpoint.endpoint_id
    )
    assert (
        trace.info.request_metadata.get(TraceMetadataKey.GATEWAY_REQUEST_TYPE)
        == GatewayRequestType.PASSTHROUGH_MODEL_OPENAI_CHAT
    )

    gateway_span = next(
        (span for span in trace.data.spans if span.name == f"gateway/{endpoint_name}"), None
    )
    assert gateway_span is not None

    # Verify streaming chunks are captured in outputs (raw SSE bytes decoded to strings)
    assert gateway_span.outputs is not None
    assert len(gateway_span.outputs) == len(mock_stream_chunks)
    # Verify the outputs contain actual SSE data (not async generator object repr)
    assert "data:" in gateway_span.outputs[0]
    assert "chatcmpl-123" in gateway_span.outputs[0]


# ─── Guardrail end-to-end scenarios ──────────────────────────────────────────


class _SimpleScorer:
    """Minimal scorer that returns 'yes' or 'no' and tracks call count."""

    def __init__(self, *, passing: bool = True) -> None:
        self.call_count = 0
        self._passing = passing

    def __call__(self, **kwargs) -> str:
        self.call_count += 1
        return "yes" if self._passing else "no"


def _make_guardrail_judge(stage, action=GuardrailAction.VALIDATION, *, passing=True):
    scorer = _SimpleScorer(passing=passing)
    return JudgeGuardrail(
        scorer=scorer,
        stage=GuardrailStage(stage),
        action=GuardrailAction(action),
        name=f"test-{stage.lower()}",
    )


def _make_guardrail_chat_response(content: str = "Hello!") -> chat.ResponsePayload:
    return chat.ResponsePayload(
        id="resp-id",
        object="chat.completion",
        created=1234567890,
        model="gpt-4",
        choices=[
            chat.Choice(
                index=0,
                message=chat.ResponseMessage(role="assistant", content=content),
                finish_reason="stop",
            )
        ],
        usage=chat.ChatUsage(prompt_tokens=5, completion_tokens=5, total_tokens=10),
    )


def _make_guardrail_mock_request(body: dict[str, Any], headers: dict[str, str] | None = None):
    req = MagicMock()
    req.state.cached_body = None
    req.state.username = None
    req.state.user_id = None
    req.json = AsyncMock(return_value=body)
    req.headers = headers or {}
    req.base_url = "http://localhost:5000/"
    return req


_GUARDRAIL_SERIALIZED_SCORER = json.dumps({"name": "safety", "builtin_scorer_class": "Safety"})


def _setup_db_guardrail(
    store: SqlAlchemyStore,
    endpoint_name: str,
    stage: str,
    action: str,
    action_endpoint_name: str | None = None,
    execution_order: int | None = None,
    name: str | None = None,
):
    """Create scorer + guardrail in DB and attach it to the endpoint."""
    guardrail_name = name or f"guardrail-{endpoint_name}-{stage}"
    experiment_id = store.create_experiment(f"exp-{guardrail_name}")
    scorer_ver = store.register_scorer(
        experiment_id, f"scorer-{guardrail_name}", _GUARDRAIL_SERIALIZED_SCORER
    )

    action_endpoint_id = None
    if action_endpoint_name:
        action_endpoint_id = store.get_gateway_endpoint(name=action_endpoint_name).endpoint_id

    guardrail = store.create_gateway_guardrail(
        name=guardrail_name,
        scorer_id=scorer_ver.scorer_id,
        scorer_version=scorer_ver.scorer_version,
        stage=GuardrailStage(stage),
        action=GuardrailAction(action),
        action_endpoint_id=action_endpoint_id,
    )
    endpoint = store.get_gateway_endpoint(name=endpoint_name)
    store.add_guardrail_to_endpoint(
        endpoint.endpoint_id, guardrail.guardrail_id, execution_order=execution_order
    )
    return guardrail, scorer_ver


def _setup_guardrail_endpoint(store: SqlAlchemyStore, name: str):
    secret = store.create_gateway_secret(
        secret_name=f"key-{name}",
        secret_value={"api_key": "sk-test"},
        provider="openai",
    )
    model_def = store.create_gateway_model_definition(
        name=f"model-{name}",
        secret_id=secret.secret_id,
        provider="openai",
        model_name="gpt-4",
    )
    return store.create_gateway_endpoint(
        name=name,
        model_configs=[
            GatewayEndpointModelConfig(
                model_definition_id=model_def.model_definition_id,
                linkage_type=GatewayModelLinkageType.PRIMARY,
                weight=1.0,
            )
        ],
    )


@pytest.mark.asyncio
async def test_invocations_bypass_header_skips_guardrails(store: SqlAlchemyStore):
    endpoint = _setup_guardrail_endpoint(store, "ep-bypass")
    mock_response = _make_guardrail_chat_response("Bypass response")
    mock_request = _make_guardrail_mock_request(
        {"messages": [{"role": "user", "content": "hello"}]},
        headers={_SANITIZE_BYPASS_HEADER: "1"},
    )

    with (
        patch("mlflow.server.gateway_api._create_provider_from_endpoint_name") as mock_create,
        patch("mlflow.server.gateway_api.load_guardrails") as mock_load,
    ):
        mock_provider = MagicMock()
        mock_provider.chat = AsyncMock(return_value=mock_response)
        mock_create.return_value = (
            mock_provider,
            GatewayEndpointConfig(
                endpoint_id=endpoint.endpoint_id, endpoint_name=endpoint.name, models=[]
            ),
        )
        response = await invocations(endpoint.name, mock_request)
        mock_load.assert_not_called()

    assert response.choices[0].message.content == "Bypass response"


@pytest.mark.asyncio
async def test_invocations_bypass_header_wrong_value_runs_guardrails(store: SqlAlchemyStore):
    endpoint = _setup_guardrail_endpoint(store, "ep-bypass-wrong-val")
    _setup_db_guardrail(store, "ep-bypass-wrong-val", "BEFORE", "VALIDATION")

    mock_request = _make_guardrail_mock_request(
        {"messages": [{"role": "user", "content": "hello"}]},
        headers={_SANITIZE_BYPASS_HEADER: "true"},  # wrong value — must not bypass
    )

    blocking_scorer = _SimpleScorer(passing=False)

    with (
        patch("mlflow.genai.scorers.base.Scorer.model_validate", return_value=blocking_scorer),
        patch("mlflow.gateway.providers.openai.OpenAIProvider.chat", AsyncMock()),
    ):
        with pytest.raises(HTTPException, match="400"):
            await invocations(endpoint.name, mock_request)

    assert blocking_scorer.call_count == 1


@pytest.mark.asyncio
async def test_real_db_pre_llm_guardrail_passes(store: SqlAlchemyStore):
    endpoint = _setup_guardrail_endpoint(store, "real-ep-pre-llm-pass")
    _setup_db_guardrail(store, "real-ep-pre-llm-pass", "BEFORE", "VALIDATION")

    mock_response = _make_guardrail_chat_response("Safe response")
    mock_request = _make_guardrail_mock_request({
        "messages": [{"role": "user", "content": "hello"}]
    })

    passing_scorer = _SimpleScorer(passing=True)

    with (
        patch("mlflow.genai.scorers.base.Scorer.model_validate", return_value=passing_scorer),
        patch(
            "mlflow.gateway.providers.openai.OpenAIProvider.chat",
            AsyncMock(return_value=mock_response),
        ),
    ):
        response = await invocations(endpoint.name, mock_request)

    assert response.choices[0].message.content == "Safe response"
    assert passing_scorer.call_count == 1


@pytest.mark.asyncio
async def test_real_db_pre_llm_guardrail_blocks(store: SqlAlchemyStore):
    endpoint = _setup_guardrail_endpoint(store, "real-ep-pre-llm-block")
    _setup_db_guardrail(store, "real-ep-pre-llm-block", "BEFORE", "VALIDATION")

    mock_request = _make_guardrail_mock_request({
        "messages": [{"role": "user", "content": "bad input"}]
    })

    blocking_scorer = _SimpleScorer(passing=False)

    with (
        patch("mlflow.genai.scorers.base.Scorer.model_validate", return_value=blocking_scorer),
        patch(
            "mlflow.gateway.providers.openai.OpenAIProvider.chat",
            AsyncMock(),
        ) as mock_chat,
    ):
        with pytest.raises(HTTPException, match="400"):
            await invocations(endpoint.name, mock_request)

    assert not mock_chat.called
    assert blocking_scorer.call_count == 1


@pytest.mark.asyncio
async def test_real_db_post_llm_guardrail_blocks(store: SqlAlchemyStore):
    endpoint = _setup_guardrail_endpoint(store, "real-ep-post-llm-block")
    _setup_db_guardrail(store, "real-ep-post-llm-block", "AFTER", "VALIDATION")

    mock_response = _make_guardrail_chat_response("Unsafe output")
    mock_request = _make_guardrail_mock_request({
        "messages": [{"role": "user", "content": "hello"}]
    })

    blocking_scorer = _SimpleScorer(passing=False)

    with (
        patch("mlflow.genai.scorers.base.Scorer.model_validate", return_value=blocking_scorer),
        patch(
            "mlflow.gateway.providers.openai.OpenAIProvider.chat",
            AsyncMock(return_value=mock_response),
        ) as mock_chat,
    ):
        with pytest.raises(HTTPException, match="400"):
            await invocations(endpoint.name, mock_request)

    assert mock_chat.called
    assert blocking_scorer.call_count == 1


@pytest.mark.asyncio
async def test_invocations_before_sanitize_rewrites_request(store: SqlAlchemyStore):
    endpoint = _setup_guardrail_endpoint(store, "ep-sanitize-before")
    sanitizer = _setup_guardrail_endpoint(store, "ep-sanitizer")
    _setup_db_guardrail(
        store, "ep-sanitize-before", "BEFORE", "SANITIZATION", action_endpoint_name=sanitizer.name
    )

    sanitized_body = {"messages": [{"role": "user", "content": "cleaned input"}]}
    mock_response = _make_guardrail_chat_response("Response to cleaned input")
    mock_request = _make_guardrail_mock_request({
        "messages": [{"role": "user", "content": "bad input"}]
    })

    failing_scorer = _SimpleScorer(passing=False)
    captured_payloads: list[Any] = []

    async def fake_chat(payload):
        captured_payloads.append(payload)
        return mock_response

    with (
        patch("mlflow.genai.scorers.base.Scorer.model_validate", return_value=failing_scorer),
        patch(
            "mlflow.gateway.guardrails.send_request",
            AsyncMock(
                return_value={"choices": [{"message": {"content": json.dumps(sanitized_body)}}]}
            ),
        ),
        patch("mlflow.gateway.providers.openai.OpenAIProvider.chat", side_effect=fake_chat),
    ):
        response = await invocations(endpoint.name, mock_request)

    assert response.choices[0].message.content == "Response to cleaned input"
    assert failing_scorer.call_count == 1
    assert captured_payloads[0].messages[0].content == "cleaned input"


@pytest.mark.asyncio
async def test_invocations_after_sanitize_rewrites_response(store: SqlAlchemyStore):
    endpoint = _setup_guardrail_endpoint(store, "ep-sanitize-after")
    sanitizer = _setup_guardrail_endpoint(store, "ep-sanitizer-after")
    _setup_db_guardrail(
        store, "ep-sanitize-after", "AFTER", "SANITIZATION", action_endpoint_name=sanitizer.name
    )

    sanitized_response = {
        "id": "resp-sanitized",
        "object": "chat.completion",
        "created": 1234567890,
        "model": "gpt-4",
        "choices": [
            {
                "index": 0,
                "message": {"role": "assistant", "content": "polite output"},
                "finish_reason": "stop",
            }
        ],
        "usage": {"prompt_tokens": 5, "completion_tokens": 5, "total_tokens": 10},
    }
    mock_response = _make_guardrail_chat_response("rude output")
    mock_request = _make_guardrail_mock_request({
        "messages": [{"role": "user", "content": "hello"}]
    })

    failing_scorer = _SimpleScorer(passing=False)

    with (
        patch("mlflow.genai.scorers.base.Scorer.model_validate", return_value=failing_scorer),
        patch(
            "mlflow.gateway.guardrails.send_request",
            AsyncMock(
                return_value={"choices": [{"message": {"content": json.dumps(sanitized_response)}}]}
            ),
        ),
        patch(
            "mlflow.gateway.providers.openai.OpenAIProvider.chat",
            AsyncMock(return_value=mock_response),
        ),
    ):
        response = await invocations(endpoint.name, mock_request)

    assert response.choices[0].message.content == "polite output"
    assert failing_scorer.call_count == 1


@pytest.mark.asyncio
async def test_invocations_sanitize_no_action_endpoint_blocks(store: SqlAlchemyStore):
    endpoint = _setup_guardrail_endpoint(store, "ep-sanitize-no-ep")
    _setup_db_guardrail(store, "ep-sanitize-no-ep", "BEFORE", "SANITIZATION")

    mock_request = _make_guardrail_mock_request({
        "messages": [{"role": "user", "content": "bad input"}]
    })
    failing_scorer = _SimpleScorer(passing=False)

    with (
        patch("mlflow.genai.scorers.base.Scorer.model_validate", return_value=failing_scorer),
        patch("mlflow.gateway.providers.openai.OpenAIProvider.chat", AsyncMock()),
    ):
        with pytest.raises(HTTPException, match="400"):
            await invocations(endpoint.name, mock_request)


@pytest.mark.asyncio
async def test_chat_completions_before_sanitize_rewrites_request(store: SqlAlchemyStore):
    endpoint = _setup_guardrail_endpoint(store, "ep-cc-sanitize-before")
    sanitizer = _setup_guardrail_endpoint(store, "ep-cc-sanitizer")
    _setup_db_guardrail(
        store,
        "ep-cc-sanitize-before",
        "BEFORE",
        "SANITIZATION",
        action_endpoint_name=sanitizer.name,
    )

    sanitized_body = {"messages": [{"role": "user", "content": "cleaned input"}]}
    mock_response = _make_guardrail_chat_response("Response to cleaned input")
    mock_request = _make_guardrail_mock_request({
        "model": endpoint.name,
        "messages": [{"role": "user", "content": "bad input"}],
    })

    failing_scorer = _SimpleScorer(passing=False)
    captured_payloads: list[Any] = []

    async def fake_chat(payload):
        captured_payloads.append(payload)
        return mock_response

    with (
        patch("mlflow.genai.scorers.base.Scorer.model_validate", return_value=failing_scorer),
        patch(
            "mlflow.gateway.guardrails.send_request",
            AsyncMock(
                return_value={"choices": [{"message": {"content": json.dumps(sanitized_body)}}]}
            ),
        ),
        patch("mlflow.gateway.providers.openai.OpenAIProvider.chat", side_effect=fake_chat),
    ):
        from mlflow.server.gateway_api import chat_completions

        response = await chat_completions(mock_request)

    assert response.choices[0].message.content == "Response to cleaned input"
    assert failing_scorer.call_count == 1
    assert captured_payloads[0].messages[0].content == "cleaned input"


@pytest.mark.asyncio
async def test_guardrails_run_in_execution_order(store: SqlAlchemyStore):
    endpoint = _setup_guardrail_endpoint(store, "ep-order-test")

    # Register in reverse order (5→1) to ensure DB insertion order != execution order.
    for i in range(5, 0, -1):
        _setup_db_guardrail(
            store,
            "ep-order-test",
            "BEFORE",
            "VALIDATION",
            execution_order=i,
            name=f"g-order-{i}",
        )

    mock_request = _make_guardrail_mock_request({
        "messages": [{"role": "user", "content": "hello"}]
    })

    call_order: list[str] = []

    def make_scorer(label: str, passing: bool):
        def scorer(**kwargs):
            call_order.append(label)
            return "yes" if passing else "no"

        return scorer

    # Guardrails 1-4 pass; guardrail 5 blocks — so all 5 must run in order 1→2→3→4→5.
    scorers = [make_scorer(f"order-{i}", passing=(i < 5)) for i in range(1, 6)]
    call_count = {"n": 0}

    def model_validate_side_effect(serialized):
        scorer = scorers[call_count["n"]]
        call_count["n"] += 1
        return scorer

    with (
        patch(
            "mlflow.genai.scorers.base.Scorer.model_validate",
            side_effect=model_validate_side_effect,
        ),
        patch("mlflow.gateway.providers.openai.OpenAIProvider.chat", AsyncMock()),
    ):
        with pytest.raises(HTTPException, match="400"):
            await invocations(endpoint.name, mock_request)

    assert call_order == ["order-1", "order-2", "order-3", "order-4", "order-5"]


@pytest.mark.asyncio
async def test_guardrail_spans_created_when_usage_tracking_on(store: SqlAlchemyStore):
    endpoint_name = "ep-guardrail-tracing"
    experiment_id = store.create_experiment(f"gateway/{endpoint_name}")

    secret = store.create_gateway_secret(
        secret_name=f"key-{endpoint_name}",
        secret_value={"api_key": "sk-test"},
        provider="openai",
    )
    model_def = store.create_gateway_model_definition(
        name=f"model-{endpoint_name}",
        secret_id=secret.secret_id,
        provider="openai",
        model_name="gpt-4",
    )
    endpoint = store.create_gateway_endpoint(
        name=endpoint_name,
        model_configs=[
            GatewayEndpointModelConfig(
                model_definition_id=model_def.model_definition_id,
                linkage_type=GatewayModelLinkageType.PRIMARY,
                weight=1.0,
            )
        ],
        usage_tracking=True,
        experiment_id=experiment_id,
    )
    _setup_db_guardrail(store, endpoint_name, "BEFORE", "VALIDATION", name="safety-check")

    mock_response = _make_guardrail_chat_response("Safe response")
    mock_request = _make_guardrail_mock_request({
        "messages": [{"role": "user", "content": "hello"}]
    })
    passing_scorer = _SimpleScorer(passing=True)

    with (
        patch("mlflow.genai.scorers.base.Scorer.model_validate", return_value=passing_scorer),
        patch(
            "mlflow.gateway.providers.openai.OpenAIProvider.chat",
            AsyncMock(return_value=mock_response),
        ),
    ):
        response = await invocations(endpoint.name, mock_request)

    assert response.choices[0].message.content == "Safe response"

    traces = TracingClient().search_traces(locations=[experiment_id])
    assert len(traces) == 1

    span_map = {s.name: s for s in traces[0].data.spans}
    assert "guardrail/safety-check" in span_map
    assert "judge" in span_map

    gspan = span_map["guardrail/safety-check"]
    jspan = span_map["judge"]
    assert gspan.span_type == SpanType.GUARDRAIL
    assert jspan.span_type == SpanType.EVALUATOR
    assert jspan.outputs["passed"] is True
    assert jspan.parent_id == gspan.span_id