litellm/tests/unit/test_router_retry_policy_update.py

"""
Tests for the retry_policy fix on Router.update_settings (LIT-3152).

Bug: the Admin UI Model Retry Settings tab posts a global ``retry_policy``
through ``POST /config/update`` -> ``UpdateRouterConfig`` ->
``Router.update_settings``. Both the pydantic schema and
``update_settings`` were dropping the field silently:

- ``UpdateRouterConfig`` had no ``retry_policy`` attribute, so
  ``model_dump(exclude_none=True)`` returned ``{}`` for that key.
- ``Router.update_settings`` had no ``"retry_policy"`` entry in
  ``_allowed_settings``, so even when fed directly the call was a no-op
  (``Setting {} is not allowed`` debug log).

The net effect was that after saving retry counts in the UI and
reloading, every value snapped back to ``defaultRetry = num_retries``
(2 by default), exactly matching the ticket repro.

This file pins both halves of the fix.
"""

import json
from dataclasses import dataclass
from typing import Final
from unittest.mock import AsyncMock, MagicMock

import pytest
from pydantic import ValidationError


import litellm
from litellm.router_strategy.budget_limiter import RouterBudgetLimiting
from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import EncryptedContentAffinityCheck
from litellm.router_utils.pre_call_checks.model_rate_limit_check import ModelRateLimitingCheck
from litellm.router_utils.pre_call_checks.prompt_caching_deployment_check import PromptCachingDeploymentCheck
from litellm.types.router import RetryPolicy, UpdateRouterConfig


@pytest.fixture(autouse=True)
def isolate_litellm_callbacks():
    callbacks_before: Final = litellm.callbacks.copy()
    yield
    litellm.callbacks = callbacks_before  # test-quality-ok: required callback-state restoration fixture


# ---------------------------------------------------------------------------
# UpdateRouterConfig schema membership (LIT-3152 part 1)
# ---------------------------------------------------------------------------


def test_update_router_config_exposes_retry_policy_field():
    """retry_policy must be a declared field on UpdateRouterConfig.

    Without it, Pydantic silently strips the key from the /config/update
    payload before the proxy even calls llm_router.update_settings.
    """
    assert "retry_policy" in UpdateRouterConfig.model_fields


def test_update_router_config_accepts_retry_policy_payload():
    """The exact payload the Admin UI Model Retry Settings tab sends must
    round-trip through the schema's ``dict(exclude_none=True)`` form, since
    that is what /config/update writes to the LiteLLM_Config row."""
    payload = {
        "retry_policy": {
            "BadRequestErrorRetries": 5,
            "RateLimitErrorRetries": 7,
            "TimeoutErrorRetries": 3,
        }
    }
    cfg = UpdateRouterConfig(**payload)
    dumped = cfg.model_dump(exclude_none=True)
    assert "retry_policy" in dumped
    assert dumped["retry_policy"]["BadRequestErrorRetries"] == 5
    assert dumped["retry_policy"]["RateLimitErrorRetries"] == 7
    assert dumped["retry_policy"]["TimeoutErrorRetries"] == 3


def test_update_router_config_rejects_malformed_retry_policy():
    """The field is typed as RetryPolicy, so /config/update validates the
    payload at the boundary and rejects non-numeric counts with a 422 instead
    of silently persisting garbage the apply path would later have to drop."""
    with pytest.raises(ValidationError):
        UpdateRouterConfig(retry_policy={"BadRequestErrorRetries": "not-an-int"})


def test_update_router_config_rejects_malformed_model_group_retry_policy():
    """model_group_retry_policy is Dict[str, RetryPolicy], so each per-group
    policy is validated the same way."""
    with pytest.raises(ValidationError):
        UpdateRouterConfig(
            model_group_retry_policy={"gpt-4": {"RateLimitErrorRetries": "x"}}
        )


# ---------------------------------------------------------------------------
# Router.update_settings retry_policy path (LIT-3152 part 2)
# ---------------------------------------------------------------------------


def _build_router() -> litellm.Router:
    return litellm.Router(
        model_list=[
            {
                "model_name": "test-model",
                "litellm_params": {
                    "model": "openai/gpt-4",
                    "api_key": "sk-fake",
                    "api_base": "http://localhost:9999",
                },
            }
        ]
    )


def test_update_settings_adds_optional_pre_call_check_once():
    router = _build_router()

    router.update_settings(num_retries=7, optional_pre_call_checks=["prompt_caching"])
    router.update_settings(optional_pre_call_checks=["prompt_caching"])

    prompt_caching_callbacks = [
        callback for callback in router.optional_callbacks if isinstance(callback, PromptCachingDeploymentCheck)
    ]
    assert len(prompt_caching_callbacks) == 1
    assert router.num_retries == 7


def test_update_settings_clears_omitted_toggleable_pre_call_checks():
    router = _build_router()

    router.update_settings(optional_pre_call_checks=["prompt_caching"])
    router.update_settings(optional_pre_call_checks=[])

    assert not any(isinstance(callback, PromptCachingDeploymentCheck) for callback in (router.optional_callbacks or []))
    assert not any(isinstance(callback, PromptCachingDeploymentCheck) for callback in litellm.callbacks)


def test_set_optional_pre_call_checks_reconciles_callback_types():
    router = _build_router()

    router.set_optional_pre_call_checks(["prompt_caching"])
    router.set_optional_pre_call_checks([])

    assert not any(isinstance(callback, PromptCachingDeploymentCheck) for callback in (router.optional_callbacks or []))
    assert not any(isinstance(callback, PromptCachingDeploymentCheck) for callback in litellm.callbacks)


def test_remove_optional_pre_call_check_removes_local_and_global_callbacks():
    router = _build_router()

    router.set_optional_pre_call_checks(["prompt_caching"])
    router._remove_optional_callbacks_of_type(PromptCachingDeploymentCheck)

    assert not any(type(callback) is PromptCachingDeploymentCheck for callback in (router.optional_callbacks or []))
    assert not any(type(callback) is PromptCachingDeploymentCheck for callback in litellm.callbacks)


def test_remove_optional_pre_call_check_keeps_global_callback_for_another_router():
    router_a = _build_router()
    router_b = _build_router()

    router_a.update_settings(optional_pre_call_checks=["prompt_caching"])
    router_b.update_settings(optional_pre_call_checks=["prompt_caching"])

    router_a.update_settings(optional_pre_call_checks=[])

    assert not any(type(callback) is PromptCachingDeploymentCheck for callback in (router_a.optional_callbacks or []))
    assert any(type(callback) is PromptCachingDeploymentCheck for callback in (router_b.optional_callbacks or []))
    assert any(type(callback) is PromptCachingDeploymentCheck for callback in litellm.callbacks)

    router_b.update_settings(optional_pre_call_checks=[])

    assert not any(type(callback) is PromptCachingDeploymentCheck for callback in (router_b.optional_callbacks or []))
    assert not any(type(callback) is PromptCachingDeploymentCheck for callback in litellm.callbacks)


def test_remove_optional_pre_call_check_keeps_global_callback_when_second_router_clears_first():
    router_a = _build_router()
    router_b = _build_router()

    router_a.update_settings(optional_pre_call_checks=["prompt_caching"])
    router_b.update_settings(optional_pre_call_checks=["prompt_caching"])

    router_b.update_settings(optional_pre_call_checks=[])

    assert any(type(callback) is PromptCachingDeploymentCheck for callback in (router_a.optional_callbacks or []))
    assert not any(type(callback) is PromptCachingDeploymentCheck for callback in (router_b.optional_callbacks or []))
    assert any(type(callback) is PromptCachingDeploymentCheck for callback in litellm.callbacks)

    router_a.update_settings(optional_pre_call_checks=[])

    assert not any(type(callback) is PromptCachingDeploymentCheck for callback in litellm.callbacks)


def test_update_settings_replaces_toggleable_pre_call_checks():
    router = _build_router()

    router.update_settings(optional_pre_call_checks=["prompt_caching"])
    router.update_settings(optional_pre_call_checks=["enforce_model_rate_limits"])

    assert not any(isinstance(callback, PromptCachingDeploymentCheck) for callback in (router.optional_callbacks or []))
    assert not any(isinstance(callback, PromptCachingDeploymentCheck) for callback in litellm.callbacks)
    assert any(isinstance(callback, ModelRateLimitingCheck) for callback in (router.optional_callbacks or []))


def test_update_settings_clears_omitted_encrypted_content_affinity_check():
    router = _build_router()

    router.update_settings(optional_pre_call_checks=["encrypted_content_affinity"])
    router.update_settings(optional_pre_call_checks=[])

    assert not any(
        isinstance(callback, EncryptedContentAffinityCheck) for callback in (router.optional_callbacks or [])
    )
    assert not any(isinstance(callback, EncryptedContentAffinityCheck) for callback in litellm.callbacks)


@pytest.mark.asyncio
async def test_update_settings_turning_off_encrypted_content_affinity_stops_flagging_requests():
    router = _build_router()

    router.update_settings(optional_pre_call_checks=["encrypted_content_affinity"])
    enabled_request: Final = {"litellm_metadata": {}, "input": "hello"}
    await router.async_get_available_deployment(model="test-model", request_kwargs=enabled_request)
    assert enabled_request["litellm_metadata"]["encrypted_content_affinity_enabled"] is True

    router.update_settings(optional_pre_call_checks=[])
    disabled_request: Final = {"litellm_metadata": {}, "input": "hello"}
    await router.async_get_available_deployment(model="test-model", request_kwargs=disabled_request)
    assert "encrypted_content_affinity_enabled" not in disabled_request["litellm_metadata"]

    router.update_settings(optional_pre_call_checks=["encrypted_content_affinity"])
    reenabled_request: Final = {"litellm_metadata": {}, "input": "hello"}
    await router.async_get_available_deployment(model="test-model", request_kwargs=reenabled_request)
    assert reenabled_request["litellm_metadata"]["encrypted_content_affinity_enabled"] is True


def test_update_settings_keeps_per_group_encrypted_content_affinity_when_global_toggle_is_omitted():
    router = litellm.Router(
        model_list=[
            {
                "model_name": "test-model",
                "litellm_params": {
                    "model": "openai/gpt-4",
                    "api_key": "sk-fake",
                    "api_base": "http://localhost:9999",
                },
            }
        ],
        model_group_affinity_config={"test-model": ["encrypted_content_affinity"]},
    )

    router.update_settings(optional_pre_call_checks=["encrypted_content_affinity"])
    router.update_settings(optional_pre_call_checks=[])

    affinity_checks: Final = [
        callback for callback in (router.optional_callbacks or []) if isinstance(callback, EncryptedContentAffinityCheck)
    ]
    assert len(affinity_checks) == 1
    assert affinity_checks[0].enable_global_affinity is False


@pytest.mark.asyncio
async def test_update_settings_preserves_router_budget_limiting_when_omitted(monkeypatch):
    async def _disable_periodic_sync(*args, **kwargs):
        return None

    monkeypatch.setattr(
        "litellm.router_strategy.budget_limiter.RouterBudgetLimiting.periodic_sync_in_memory_spend_with_redis",
        _disable_periodic_sync,
    )
    router = _build_router()

    router.add_optional_pre_call_checks(["router_budget_limiting"])
    router.update_settings(optional_pre_call_checks=[])

    assert any(isinstance(callback, RouterBudgetLimiting) for callback in (router.optional_callbacks or []))


def test_update_settings_persists_retry_policy_dict():
    """When the proxy's ``_add_router_settings_from_db_config`` calls
    ``llm_router.update_settings(retry_policy={...})`` after reading the
    DB row, the dict must land on ``self.retry_policy`` as a typed
    ``RetryPolicy`` (mirroring ``Router.__init__`` semantics)."""
    router = _build_router()
    assert router.retry_policy is None  # baseline

    router.update_settings(
        retry_policy={
            "BadRequestErrorRetries": 5,
            "RateLimitErrorRetries": 7,
            "TimeoutErrorRetries": 3,
        }
    )

    assert isinstance(router.retry_policy, RetryPolicy)
    assert router.retry_policy.BadRequestErrorRetries == 5
    assert router.retry_policy.RateLimitErrorRetries == 7
    assert router.retry_policy.TimeoutErrorRetries == 3


def test_update_settings_accepts_retry_policy_object_unchanged():
    """A pre-built ``RetryPolicy`` instance must pass through verbatim so
    callers that already constructed one (e.g. tests or programmatic
    callers) keep working."""
    router = _build_router()

    policy = RetryPolicy(BadRequestErrorRetries=2)
    router.update_settings(retry_policy=policy)

    assert router.retry_policy is policy


def test_update_settings_ignores_malformed_retry_policy():
    """A non-dict, non-``RetryPolicy`` value (e.g. a YAML typo like
    ``retry_policy: 5`` reaching ``update_settings``) must not land on
    ``self.retry_policy``. ``Router.__init__`` already drops such inputs;
    the update path must match so a malformed config can't store garbage
    that ``get_num_retries_from_retry_policy`` would only choke on at
    request time."""
    router = _build_router()

    existing = RetryPolicy(BadRequestErrorRetries=4)
    router.update_settings(retry_policy=existing)
    assert router.retry_policy is existing

    for bad_value in (5, "RateLimitErrorRetries=7", ["BadRequestErrorRetries"]):
        router.update_settings(retry_policy=bad_value)
        assert router.retry_policy is existing


def test_update_settings_get_settings_round_trip_for_retry_policy():
    """``GET /get/config/callbacks`` serializes ``llm_router.get_settings()``
    back to the UI. After updating, the round-trip must reflect the new
    values rather than the pre-update sentinel."""
    router = _build_router()
    pre = router.get_settings().get("retry_policy")
    assert pre is None

    router.update_settings(
        retry_policy={
            "BadRequestErrorRetries": 5,
            "RateLimitErrorRetries": 7,
        }
    )
    post = router.get_settings().get("retry_policy")
    assert post is not None
    assert post.BadRequestErrorRetries == 5
    assert post.RateLimitErrorRetries == 7


def test_update_settings_unrelated_kwargs_still_skipped():
    """Regression guard: the new branch must not relax the
    ``_allowed_settings`` allowlist for unrelated keys. An unknown
    setting should still be dropped silently as before."""
    router = _build_router()
    router.update_settings(this_is_not_a_router_setting=123)
    assert not hasattr(router, "this_is_not_a_router_setting")


# ---------------------------------------------------------------------------
# End-to-end persist -> apply -> read-back (LIT-3152 part 3)
#
# The tests above pin each layer in isolation, so they would all still pass
# if a regression flipped ``ConfigYAML.router_settings`` back to a loose
# ``dict`` (silently dropping retry_policy on the DB write) or if
# ``_add_router_settings_from_db_config`` stopped pushing the stored row onto
# the live router. This drives the real handler chain an Admin UI save
# triggers — ``POST /config/update`` writes the LiteLLM_Config row,
# ``add_deployment`` applies it to ``llm_router``, and
# ``GET /get/config/callbacks`` serializes it back — so the round trip is
# pinned, not just the pieces.
# ---------------------------------------------------------------------------


@dataclass(frozen=True, slots=True)
class _FakeConfigRow:
    param_name: str
    param_value: dict


class _FakeConfigTable:
    """Stand-in for ``prisma_client.db.litellm_config``.

    Reproduces the one behavior the apply path relies on: a value written as a
    JSON string by ``/config/update`` reads back as a parsed dict, which is
    what ``_add_router_settings_from_db_config``'s
    ``isinstance(param_value, dict)`` branch requires to forward the settings.
    """

    def __init__(self):
        self.rows = {}

    async def find_first(self, where):
        return self.rows.get(where["param_name"])

    async def upsert(self, where, data):
        name = where["param_name"]
        raw = (data["update"] if name in self.rows else data["create"])["param_value"]
        self.rows[name] = _FakeConfigRow(name, json.loads(raw) if isinstance(raw, str) else raw)


@pytest.mark.asyncio
async def test_config_update_persists_and_reads_back_retry_policy(monkeypatch):
    """The exact global retry_policy save the UI performs must survive the
    real ``/config/update`` -> DB -> apply -> ``/get/config/callbacks`` path,
    not snap back to the ``num_retries`` fallback the ticket reported."""
    import litellm.proxy.proxy_server as proxy_server
    from litellm.proxy._types import ConfigYAML, LitellmUserRoles, UserAPIKeyAuth

    router = _build_router()
    assert router.retry_policy is None

    fake_table = _FakeConfigTable()
    prisma_client = MagicMock()
    prisma_client.db.litellm_config = fake_table

    async def _apply_router_settings(*args, **kwargs):
        await proxy_server.proxy_config._add_router_settings_from_db_config(
            llm_router=router, prisma_client=prisma_client
        )

    monkeypatch.setattr(proxy_server, "prisma_client", prisma_client)
    monkeypatch.setattr(proxy_server, "llm_router", router)
    monkeypatch.setattr(proxy_server.proxy_config, "add_deployment", _apply_router_settings)
    monkeypatch.setattr(proxy_server.proxy_config, "get_config", AsyncMock(return_value={}))

    posted = UpdateRouterConfig(
        retry_policy=RetryPolicy(
            BadRequestErrorRetries=5,
            TimeoutErrorRetries=3,
            RateLimitErrorRetries=7,
        )
    )
    request = MagicMock()
    request.json = AsyncMock(return_value={"router_settings": {"retry_policy": posted.model_dump()}})

    await proxy_server.update_config(
        config_info=ConfigYAML(router_settings=posted),
        request=request,
        user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1234"),
    )

    persisted = fake_table.rows["router_settings"].param_value["retry_policy"]
    assert persisted == {
        "BadRequestErrorRetries": 5,
        "TimeoutErrorRetries": 3,
        "RateLimitErrorRetries": 7,
    }

    assert isinstance(router.retry_policy, RetryPolicy)
    assert router.retry_policy.RateLimitErrorRetries == 7

    read_back = (
        await proxy_server.get_config(
            user_api_key_dict=UserAPIKeyAuth(
                user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1234"
            )
        )
    )["router_settings"]["retry_policy"]
    assert read_back.BadRequestErrorRetries == 5
    assert read_back.TimeoutErrorRetries == 3
    assert read_back.RateLimitErrorRetries == 7