scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-02 02:11:58 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 101
litellm/tests/unit/proxy/proxy_server/test_proxy_config.py
yujonglee be67fce26a
refactor(proxy): inject tracing receiver and access context (#44035) 
* refactor(proxy): inject tracing receiver and access context

* refactor(proxy): own tracing resources through FastAPI lifespan

* test(proxy): pass tracing dependency in Lens lifecycle

* refactor(proxy): stop tracing logger cooperatively

* refactor(proxy): derive tracing permissions in one place

* refactor(proxy): compose application lifespan state

* refactor(proxy): give Lens tracing storage directly

* refactor(tracing): name shared ClickHouse storage explicitly

* refactor(tracing): extract shared ClickHouse storage crate

* test(proxy): isolate db push timeout from Lens safety check

* fix(tracing): drain spend retries during shutdown
2026-10-01 13:45:32 -07:00

5286 lines
211 KiB
Python
Raw Blame History

"""Behavior pins for ProxyConfig and module-level config scrubbers.
Pins covered:
- Module-level: ``_is_remote_module_url``, ``_scrub_guardrail_inner``,
``_scrub_db_overlay_remote_module_loads``
- All ``ProxyConfig`` methods listed in the pin file.
"""
from __future__ import annotations
import asyncio
import json
import logging
import os
import re
from collections.abc import Mapping
from contextlib import nullcontext
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from types import MappingProxyType, SimpleNamespace
from typing import Any, Dict, Final
from unittest.mock import AsyncMock, MagicMock
import pytest
from pydantic import JsonValue, TypeAdapter, ValidationError
import litellm
from litellm.proxy._types import CommonProxyErrors
from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper
from litellm.proxy.proxy_server import (
ProxyConfig,
_is_remote_module_url,
_scrub_db_overlay_remote_module_loads,
_scrub_guardrail_inner,
resolve_complexity_router_plugins,
resolve_routing_plugins,
validate_auto_router_capability_limits,
validate_deployment_access_windows,
validate_deployment_complexity_router_placement,
validate_deployment_max_agentic_loops,
)
from .conftest import normalize
@pytest.mark.asyncio
@pytest.mark.parametrize("shutdown_error", [False, True])
async def test_tracing_config_automatically_logs_spend_without_callback_setting(shutdown_error: bool) -> None:
from litellm.integrations.clickhouse.clickhouse_spend_logger import ClickHouseSpendLogger
from litellm.proxy.tracing_runtime import manage_tracing
from litellm.tracing import TraceReceiver
from litellm.tracing.store import TraceStore
storage: Final = MagicMock()
storage.ensure_schema = AsyncMock()
storage.insert_rows = AsyncMock()
receiver: Final = TraceReceiver(TraceStore(storage))
outcome: Final = pytest.raises(RuntimeError, match="shutdown failure") if shutdown_error else nullcontext()
with outcome:
async with manage_tracing(enabled=True, receiver_factory=lambda: receiver):
storage.ensure_schema.assert_awaited_once()
logger: Final = next(
callback
for callback in litellm._async_success_callback
if isinstance(callback, ClickHouseSpendLogger) and callback.storage is storage
)
now: Final = datetime.now()
await logger.async_log_success_event(
{
"standard_logging_object": {
"id": "response-1",
"startTime": now.timestamp(),
"endTime": now.timestamp(),
"response_cost": 0.25,
}
},
None,
now,
now,
)
storage.insert_rows.assert_not_awaited()
if shutdown_error:
raise RuntimeError("shutdown failure")
assert storage.insert_rows.await_args.args[0] == "spend_logs"
assert storage.insert_rows.await_args.args[1][0]["spend"] == 0.25
assert logger not in litellm._async_success_callback
assert logger._flush_task is not None and logger._flush_task.done()
assert not logger._flush_task.cancelled()
# ---------------------------------------------------------------------------
# _is_remote_module_url
# ---------------------------------------------------------------------------
def test__is_remote_module_url_identifies_remote_and_local():
result = {
"s3": _is_remote_module_url("s3://bucket/key.py"),
"gcs": _is_remote_module_url("gcs://bucket/key.py"),
"local": _is_remote_module_url("my.module.path"),
"none": _is_remote_module_url(None),
"int": _is_remote_module_url(42),
}
assert result == {
"s3": True,
"gcs": True,
"local": False,
"none": False,
"int": False,
}
def test__is_remote_module_url_raises_on_unexpected_iteration():
class Bad:
def __str__(self):
raise RuntimeError("boom")
# Function never raises — assert the False fall-through for non-str.
with pytest.raises(AssertionError):
# Force an error-style assertion: object is not str, returns False.
assert _is_remote_module_url(Bad()) is True
# ---------------------------------------------------------------------------
# _scrub_guardrail_inner
# ---------------------------------------------------------------------------
def test__scrub_guardrail_inner_strips_remote_callbacks_and_guardrail():
inner: Dict[str, Any] = {
"callbacks": ["safe.mod", "s3://attacker/m.py", "gcs://x/y.py"],
"guardrail": "s3://attacker/g.py",
"default_on": True,
}
_scrub_guardrail_inner(inner)
assert normalize(inner) == {
"callbacks": ["safe.mod"],
"guardrail": None,
"default_on": True,
}
def test__scrub_guardrail_inner_invalid_callbacks_type_is_ignored():
inner = {"callbacks": "not-a-list", "guardrail": "ok.module"}
_scrub_guardrail_inner(inner)
# No mutation on non-list callbacks; guardrail untouched (not remote).
assert inner == {"callbacks": "not-a-list", "guardrail": "ok.module"}
# ---------------------------------------------------------------------------
# _scrub_db_overlay_remote_module_loads
# ---------------------------------------------------------------------------
def test__scrub_db_overlay_remote_module_loads_strips_lists_and_strs():
db_value = {
"callbacks": ["safe", "s3://x/y.py"],
"success_callback": ["gcs://a/b.py", "safe2"],
"post_call_rules": "s3://bad/m.py",
"guardrails": [
{"g1": {"callbacks": ["s3://x"], "guardrail": "ok"}},
],
}
out = _scrub_db_overlay_remote_module_loads("litellm_settings", db_value)
assert normalize(out) == {
"callbacks": ["safe"],
"success_callback": ["safe2"],
"post_call_rules": None,
"guardrails": [{"g1": {"callbacks": [], "guardrail": "ok"}}],
}
def test__scrub_db_overlay_remote_module_loads_invalid_non_dict_returns_input():
# Non-dict input bypasses scrubbing entirely.
assert _scrub_db_overlay_remote_module_loads("litellm_settings", "raw") == "raw"
# ---------------------------------------------------------------------------
# resolve_complexity_router_plugins
# ---------------------------------------------------------------------------
def test_resolve_complexity_router_plugins_no_plugins_key_is_a_noop():
config: Dict[str, Any] = {"tiers": {"SIMPLE": "gpt-4o-mini"}}
resolve_complexity_router_plugins(model_name="smart-router", complexity_router_config=config, config_file_path=None)
assert config == {"tiers": {"SIMPLE": "gpt-4o-mini"}}
def test_resolve_complexity_router_plugins_resolves_dotted_path_to_live_instance(tmp_path):
plugin_file = tmp_path / "my_plugin.py"
plugin_file.write_text(
"class _Plugin:\n async def run(self, context):\n return context\n\nmy_plugin_instance = _Plugin()\n"
)
config: Dict[str, Any] = {"plugins": ["my_plugin.my_plugin_instance"]}
resolve_complexity_router_plugins(
model_name="smart-router",
complexity_router_config=config,
config_file_path=str(tmp_path / "config.yaml"),
)
assert len(config["plugins"]) == 1
assert hasattr(config["plugins"][0], "run")
assert type(config["plugins"][0]).__name__ == "_Plugin"
def test_validate_deployment_complexity_router_placement_refuses_to_start():
"""Rejected here rather than at router build for the same reason as max_agentic_loops: the
proxy builds its router with ignore_invalid_deployments=True, so a rejection further down
turns the bad deployment into a silently missing model instead of a refusal to start."""
model = {
"model_name": "smart-router",
"litellm_params": {
"model": "auto_router/complexity_router",
"complexity_router_config": {"tiers": {"SIMPLE": "gpt-4o-mini"}},
"tier_boundaries": {"simple_medium": 0.1},
},
}
with pytest.raises(ValueError, match="tier_boundaries"):
validate_deployment_complexity_router_placement(model)
@pytest.mark.parametrize(
"litellm_params",
[
{"model": "gpt-4o"},
{"model": "openai/gpt-4o", "embedding_model": "text-embedding-3-small"},
{
"model": "auto_router/complexity_router",
"complexity_router_config": {"tiers": {"SIMPLE": "gpt-4o-mini"}, "tier_boundaries": {"simple_medium": 0.1}},
},
],
)
def test_validate_deployment_complexity_router_placement_leaves_valid_deployments_alone(litellm_params):
"""`embedding_model` is a legitimate flat param on an s3_vectors vector store, so the gate is
scoped to complexity routers rather than applied to every deployment."""
model = {"model_name": "m", "litellm_params": dict(litellm_params)}
validate_deployment_complexity_router_placement(model)
assert model["litellm_params"] == litellm_params
def _heuristic_v2_row(model_name: str, classifier_type: str = "heuristic_v2") -> dict[str, object]:
return {
"model_name": model_name,
"litellm_params": {
"model": "auto_router/complexity_router",
"complexity_router_config": {"classifier_type": classifier_type, "tiers": {"SIMPLE": "gpt-4o-mini"}},
},
}
def _custom_tier_row(model_name: str) -> dict[str, object]:
return {
"model_name": model_name,
"litellm_params": {
"model": "auto_router/complexity_router",
"complexity_router_config": {
"classifier_type": "llm",
"tier_definitions": [
{"name": "routine", "description": "routine drafting"},
{"name": "hard", "description": "hard reasoning"},
],
"tiers": {"routine": "gpt-4o-mini", "hard": "gpt-4o"},
"fallback_tier": "routine",
},
},
}
def _operator_examples_row(model_name: str) -> dict[str, object]:
return {
"model_name": model_name,
"litellm_params": {
"model": "auto_router/complexity_router",
"complexity_router_config": {
"classifier_type": "llm",
"classifier_llm_config": {"model": "gpt-4o-mini"},
"tiers": {"SIMPLE": "gpt-4o-mini"},
"classification_examples": '- "reset my password" -> SIMPLE',
},
},
}
def _custom_prompt_row(model_name: str) -> dict[str, object]:
return {
"model_name": model_name,
"litellm_params": {
"model": "auto_router/complexity_router",
"complexity_router_config": {
"classifier_type": "llm",
"classifier_llm_config": {"model": "gpt-4o-mini", "system_prompt": "judge it my way"},
"tiers": {"SIMPLE": "gpt-4o-mini"},
},
},
}
@pytest.mark.parametrize(
"over_limit_rows,subject",
[
([_heuristic_v2_row("a"), _heuristic_v2_row("b"), _heuristic_v2_row("c", "heuristic")], "heuristic_v2"),
([_custom_tier_row("a"), _custom_tier_row("b"), _heuristic_v2_row("c", "heuristic")], "tier_definitions"),
(
[_custom_prompt_row("a"), _custom_prompt_row("b"), _heuristic_v2_row("c", "heuristic")],
"operator-written classifier prompt",
),
(
[_custom_tier_row("a"), _custom_prompt_row("b"), _heuristic_v2_row("c", "heuristic")],
"operator-written classifier prompt",
),
(
[_operator_examples_row("a"), _custom_tier_row("b"), _heuristic_v2_row("c", "heuristic")],
"operator-written classifier prompt",
),
],
)
def test_validate_auto_router_capability_limits_refuses_to_start_over_the_limit(
over_limit_rows: list[dict[str, object]], subject: str
) -> None:
"""Same reason as the two validators above: the proxy router swallows registration errors, so
an over-limit config.yaml must fail here instead of booting with a silently missing router."""
with pytest.raises(ValueError, match=re.escape("At most 1 auto-router")) as exc_info:
validate_auto_router_capability_limits(over_limit_rows, limit=1)
assert subject in str(exc_info.value)
assert "'auto_router' feature lifts the limit" in str(exc_info.value)
@pytest.mark.parametrize(
"model_list,limit",
[
([_heuristic_v2_row("a"), _heuristic_v2_row("b")], None),
([_heuristic_v2_row("a"), _heuristic_v2_row("c", "heuristic")], 1),
([{"model_name": "gpt-4o", "litellm_params": {"model": "gpt-4o"}}], 1),
([_custom_tier_row("a"), _custom_tier_row("b")], None),
([_custom_tier_row("a"), _heuristic_v2_row("b")], 1),
],
)
def test_validate_auto_router_capability_limits_leaves_configs_within_the_limit_alone(
model_list: list[dict[str, object]], limit: int | None
) -> None:
"""The last case is the separate-ceiling invariant: one router of each capability fits under a limit of one."""
assert validate_auto_router_capability_limits(model_list, limit=limit) is None
_TWO_HEURISTIC_V2_ROUTERS_YAML = (
"model_list:\n"
" - model_name: gpt-4o-mini\n"
" litellm_params:\n"
" model: openai/gpt-4o-mini\n"
" api_key: k\n"
" - model_name: v2-a\n"
" litellm_params:\n"
" model: auto_router/complexity_router\n"
" complexity_router_config:\n"
" classifier_type: heuristic_v2\n"
" tiers: {SIMPLE: gpt-4o-mini}\n"
" - model_name: v2-b\n"
" litellm_params:\n"
" model: auto_router/complexity_router\n"
" complexity_router_config:\n"
" classifier_type: heuristic_v2\n"
" tiers: {SIMPLE: gpt-4o-mini}\n"
"router_settings:\n"
" auto_router_capability_limit: 99\n"
)
@pytest.mark.asyncio
@pytest.mark.parametrize("license_limit", [1, None])
@pytest.mark.parametrize("classifier_type", ["heuristic_v2", "capability", "llm_v2"])
async def test_ProxyConfig_load_config_takes_the_classifier_limit_from_the_license_only(
tmp_path, monkeypatch, license_limit: int | None, classifier_type: str
) -> None:
"""`router_settings.auto_router_capability_limit` is managed outside config.yaml: an operator
cannot grant the entitlement by editing the config, and a licensed proxy boots both routers."""
f = tmp_path / "c.yaml"
forecast_settings = {
"capability": (
" classifier_llm_config: {model: gpt-4o-mini}\n"
" capability_classifier_config: {efficient_tier: SIMPLE, capable_tier: REASONING, base_threshold: 0.7}\n"
),
"llm_v2": (
" classifier_llm_config: {model: gpt-4o-mini}\n"
" adaptive: false\n"
" llm_v2_config: {efficient_profile: Small solver, capable_profile: Large solver, harness: One attempt, max_quality_gap: 0.05}\n"
),
}
config_yaml = _TWO_HEURISTIC_V2_ROUTERS_YAML.replace(
"classifier_type: heuristic_v2\n",
f"classifier_type: {classifier_type}\n{forecast_settings.get(classifier_type, '')}",
).replace("tiers: {SIMPLE: gpt-4o-mini}", "tiers: {SIMPLE: gpt-4o-mini, REASONING: gpt-4o}")
f.write_text(config_yaml)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False)
monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False)
monkeypatch.setattr("litellm.proxy.proxy_server._license_check.auto_router_capability_limit", lambda: license_limit)
if license_limit is None:
router, _model_list, _general_settings = await ProxyConfig().load_config(router=None, config_file_path=str(f))
assert router.auto_router_capability_limit is not None
assert router.auto_router_capability_limit() is None
assert sorted(router.complexity_routers) == ["v2-a", "v2-b"]
return
with pytest.raises(ValueError, match=re.escape("config.yaml model_list: At most 1 auto-router")):
await ProxyConfig().load_config(router=None, config_file_path=str(f))
@pytest.mark.asyncio
async def test_ProxyConfig_load_config_router_refuses_a_db_heuristic_v2_router_beyond_the_license(
tmp_path, monkeypatch
) -> None:
"""config.yaml holds the one allowed heuristic_v2 router; a second one arriving later from the DB
is refused at registration because the router was built with the license's ceiling."""
from litellm.types.router import Deployment
f = tmp_path / "c.yaml"
f.write_text(
_TWO_HEURISTIC_V2_ROUTERS_YAML.replace(" - model_name: v2-b\n", " - model_name: v1-b\n", 1).replace(
"classifier_type: heuristic_v2\n tiers: {SIMPLE: gpt-4o-mini}\nrouter_settings",
"classifier_type: heuristic\n tiers: {SIMPLE: gpt-4o-mini}\nrouter_settings",
)
)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False)
monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False)
monkeypatch.setattr("litellm.proxy.proxy_server._license_check.auto_router_capability_limit", lambda: 1)
router, _model_list, _general_settings = await ProxyConfig().load_config(router=None, config_file_path=str(f))
assert router.auto_router_capability_limit is not None
assert router.auto_router_capability_limit() == 1
assert sorted(router.complexity_routers) == ["v1-b", "v2-a"]
db_row = Deployment(**_heuristic_v2_row("v2-from-db"), model_info={"id": "db-id"})
assert router.upsert_deployment(db_row) is None
assert sorted(router.complexity_routers) == ["v1-b", "v2-a"]
def test_validate_deployment_max_agentic_loops_allows_a_deployment_without_the_key():
model = {"model_name": "gpt-4o", "litellm_params": {"model": "gpt-4o"}}
validate_deployment_max_agentic_loops(model)
assert "max_agentic_loops" not in model["litellm_params"]
def test_validate_deployment_max_agentic_loops_leaves_a_valid_ceiling_alone():
model = {"model_name": "gpt-4o", "litellm_params": {"model": "gpt-4o", "max_agentic_loops": 5}}
validate_deployment_max_agentic_loops(model)
assert model["litellm_params"]["max_agentic_loops"] == 5
def test_validate_deployment_max_agentic_loops_rejects_zero():
"""
A per-deployment 0 used to be swallowed by an `or 3` and read as the default
ceiling of 3, handing the loosest setting to whoever asked for the tightest.
"""
with pytest.raises(ValueError, match="must be at least 1, got 0"):
validate_deployment_max_agentic_loops(
{"model_name": "gpt-4o", "litellm_params": {"model": "gpt-4o", "max_agentic_loops": 0}}
)
def test_validate_deployment_max_agentic_loops_rejects_a_non_integer():
"""
A per-deployment non-integer used to let the proxy boot and then fail every
request to that model with `invalid literal for int() with base 10`.
"""
with pytest.raises(TypeError, match="must be an integer"):
validate_deployment_max_agentic_loops(
{"model_name": "gpt-4o", "litellm_params": {"model": "gpt-4o", "max_agentic_loops": "three"}}
)
def test_validate_deployment_max_agentic_loops_rejects_a_bool():
with pytest.raises(TypeError, match="must be an integer"):
validate_deployment_max_agentic_loops(
{"model_name": "gpt-4o", "litellm_params": {"model": "gpt-4o", "max_agentic_loops": True}}
)
def test_validate_deployment_max_agentic_loops_accepts_a_ceiling_from_an_env_var():
"""
`max_agentic_loops: os.environ/MAX_AGENTIC_LOOPS` is resolved to a string
before this check runs, and the old `int(... or 3)` accepted that, so
refusing it here would stop an already working proxy from booting.
"""
model = {"model_name": "gpt-4o", "litellm_params": {"model": "gpt-4o", "max_agentic_loops": "5"}}
validate_deployment_max_agentic_loops(model)
assert model["litellm_params"]["max_agentic_loops"] == "5"
def test_validate_deployment_max_agentic_loops_names_the_offending_model():
with pytest.raises(ValueError, match="on model 'claude-sonnet-4-5'"):
validate_deployment_max_agentic_loops(
{"model_name": "claude-sonnet-4-5", "litellm_params": {"max_agentic_loops": -1}}
)
def test_resolve_complexity_router_plugins_rejects_non_routing_plugin_object(tmp_path):
plugin_file = tmp_path / "bad_plugin.py"
plugin_file.write_text("not_a_plugin = object()\n")
config: Dict[str, Any] = {"plugins": ["bad_plugin.not_a_plugin"]}
with pytest.raises(ValueError, match="does not implement the RoutingPlugin interface"):
resolve_complexity_router_plugins(
model_name="smart-router",
complexity_router_config=config,
config_file_path=str(tmp_path / "config.yaml"),
)
def test_resolve_complexity_router_plugins_rejects_synchronous_run_method(tmp_path):
"""Regression: @runtime_checkable only checks that `run` exists as an attribute,
not that it's a coroutine function. A plugin with a synchronous `run` passes a bare
isinstance() check and would only fail at request time with a confusing
`TypeError: object RoutingContext can't be used in 'await' expression`. Reported
by Greptile on PR #33251."""
plugin_file = tmp_path / "sync_plugin.py"
plugin_file.write_text(
"class _SyncPlugin:\n"
" def run(self, context):\n"
" return context\n"
"\n"
"sync_plugin_instance = _SyncPlugin()\n"
)
config: Dict[str, Any] = {"plugins": ["sync_plugin.sync_plugin_instance"]}
with pytest.raises(ValueError, match="does not implement the RoutingPlugin interface"):
resolve_complexity_router_plugins(
model_name="smart-router",
complexity_router_config=config,
config_file_path=str(tmp_path / "config.yaml"),
)
def test_resolve_complexity_router_plugins_resolves_classifier_plugin_dotted_path(tmp_path):
plugin_file = tmp_path / "my_classifier.py"
plugin_file.write_text(
"class _Classifier:\n"
" async def classify(self, context):\n"
" return 'SIMPLE'\n"
"\n"
"my_classifier_instance = _Classifier()\n"
)
config: dict[str, Any] = {
"classifier_type": "custom",
"classifier_plugin": "my_classifier.my_classifier_instance",
}
resolve_complexity_router_plugins(
model_name="smart-router",
complexity_router_config=config,
config_file_path=str(tmp_path / "config.yaml"),
)
assert hasattr(config["classifier_plugin"], "classify")
assert type(config["classifier_plugin"]).__name__ == "_Classifier"
def test_resolve_complexity_router_plugins_rejects_non_classifier_object(tmp_path):
plugin_file = tmp_path / "bad_classifier.py"
plugin_file.write_text("not_a_classifier = object()\n")
config: dict[str, Any] = {"classifier_plugin": "bad_classifier.not_a_classifier"}
with pytest.raises(ValueError, match="does not implement the ClassifierPlugin interface"):
resolve_complexity_router_plugins(
model_name="smart-router",
complexity_router_config=config,
config_file_path=str(tmp_path / "config.yaml"),
)
def test_resolve_complexity_router_plugins_rejects_synchronous_classify_method(tmp_path):
"""A synchronous `classify` passes the runtime_checkable isinstance and would only fail on
the first classified request, so reject it at config load like the sync-run case above."""
plugin_file = tmp_path / "sync_classifier.py"
plugin_file.write_text(
"class _SyncClassifier:\n"
" def classify(self, context):\n"
" return 'SIMPLE'\n"
"\n"
"sync_classifier_instance = _SyncClassifier()\n"
)
config: dict[str, Any] = {"classifier_plugin": "sync_classifier.sync_classifier_instance"}
with pytest.raises(ValueError, match="does not implement the ClassifierPlugin interface"):
resolve_complexity_router_plugins(
model_name="smart-router",
complexity_router_config=config,
config_file_path=str(tmp_path / "config.yaml"),
)
def test_resolve_complexity_router_plugins_leaves_live_classifier_instance_alone():
class _Classifier:
async def classify(self, context):
return "SIMPLE"
instance = _Classifier()
config: dict[str, Any] = {"classifier_plugin": instance}
resolve_complexity_router_plugins(model_name="smart-router", complexity_router_config=config, config_file_path=None)
assert config["classifier_plugin"] is instance
# ---------------------------------------------------------------------------
# resolve_routing_plugins
# ---------------------------------------------------------------------------
def test_resolve_routing_plugins_resolves_dotted_paths(tmp_path):
plugin_file = tmp_path / "rs_plugin.py"
plugin_file.write_text(
"class _Plugin:\n async def run(self, context):\n return context\n\nrs_plugin_instance = _Plugin()\n"
)
resolved = resolve_routing_plugins(
plugin_paths=["rs_plugin.rs_plugin_instance"],
config_file_path=str(tmp_path / "config.yaml"),
source_label="router_settings.plugins",
)
assert len(resolved) == 1
assert type(resolved[0]).__name__ == "_Plugin"
def test_resolve_routing_plugins_passes_through_instances(tmp_path):
class _Plugin:
async def run(self, context):
return context
instance = _Plugin()
resolved = resolve_routing_plugins(
plugin_paths=[instance],
config_file_path=None,
source_label="router_settings.plugins",
)
assert resolved == [instance]
def test_resolve_routing_plugins_rejects_non_routing_plugin(tmp_path):
plugin_file = tmp_path / "bad_rs_plugin.py"
plugin_file.write_text("not_a_plugin = object()\n")
with pytest.raises(ValueError, match=re.escape("router_settings.plugins")):
resolve_routing_plugins(
plugin_paths=["bad_rs_plugin.not_a_plugin"],
config_file_path=str(tmp_path / "config.yaml"),
source_label="router_settings.plugins",
)
def test_resolve_routing_plugins_rejects_synchronous_run(tmp_path):
plugin_file = tmp_path / "sync_rs_plugin.py"
plugin_file.write_text(
"class _SyncPlugin:\n"
" def run(self, context):\n"
" return context\n"
"\n"
"sync_plugin_instance = _SyncPlugin()\n"
)
with pytest.raises(ValueError, match="does not implement the RoutingPlugin interface"):
resolve_routing_plugins(
plugin_paths=["sync_rs_plugin.sync_plugin_instance"],
config_file_path=str(tmp_path / "config.yaml"),
source_label="router_settings.plugins",
)
# ---------------------------------------------------------------------------
# ProxyConfig.__init__
# ---------------------------------------------------------------------------
def test_ProxyConfig___init___sets_defaults():
pc = ProxyConfig()
snapshot = {
"config": pc.config,
"last_semantic_filter_config": pc._last_semantic_filter_config,
"worker_registry": pc.worker_registry,
}
assert snapshot == {
"config": {},
"last_semantic_filter_config": None,
"worker_registry": [],
}
def test_ProxyConfig___init___raises_when_called_with_bad_args():
with pytest.raises(TypeError):
ProxyConfig("unexpected-positional") # type: ignore[call-arg]
# ---------------------------------------------------------------------------
# ProxyConfig.is_yaml
# ---------------------------------------------------------------------------
def test_ProxyConfig_is_yaml_detects_yaml_and_non_yaml(tmp_path):
yaml_file = tmp_path / "c.yaml"
yaml_file.write_text("model_list: []\n")
yml_file = tmp_path / "c.yml"
yml_file.write_text("model_list: []\n")
json_file = tmp_path / "c.json"
json_file.write_text("{}")
pc = ProxyConfig()
result = {
"yaml": pc.is_yaml(str(yaml_file)),
"yml": pc.is_yaml(str(yml_file)),
"json": pc.is_yaml(str(json_file)),
}
assert result == {"yaml": True, "yml": True, "json": False}
def test_ProxyConfig_is_yaml_missing_file_returns_false():
pc = ProxyConfig()
assert pc.is_yaml("/no/such/path/here.yaml") is False
# ---------------------------------------------------------------------------
# ProxyConfig._load_yaml_file
# ---------------------------------------------------------------------------
def test_ProxyConfig__load_yaml_file_returns_parsed_dict(tmp_path):
f = tmp_path / "c.yaml"
f.write_text("a: 1\nb: two\nc:\n - x\n - y\n")
pc = ProxyConfig()
result = pc._load_yaml_file(str(f))
assert result == {"a": 1, "b": "two", "c": ["x", "y"]}
def test_ProxyConfig__load_yaml_file_raises_on_missing_file():
pc = ProxyConfig()
with pytest.raises(Exception, match="Error loading yaml file"):
pc._load_yaml_file("/no/such/file.yaml")
# ---------------------------------------------------------------------------
# ProxyConfig._get_config_from_file
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_ProxyConfig__get_config_from_file_loads_yaml(tmp_path):
f = tmp_path / "c.yaml"
f.write_text("model_list: []\ngeneral_settings: {}\nlitellm_settings:\n drop_params: true\n")
pc = ProxyConfig()
result = await pc._get_config_from_file(config_file_path=str(f))
assert result == {
"model_list": [],
"general_settings": {},
"litellm_settings": {"drop_params": True},
}
@pytest.mark.asyncio
async def test_ProxyConfig__get_config_from_file_missing_path_raises():
pc = ProxyConfig()
with pytest.raises(Exception, match="Config file not found"):
await pc._get_config_from_file(config_file_path="/no/such/file.yaml")
# ---------------------------------------------------------------------------
# ProxyConfig._process_includes
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_ProxyConfig__process_includes_merges_files(tmp_path):
inc = tmp_path / "models.yaml"
inc.write_text("model_list:\n - model_name: gpt-4\n")
pc = ProxyConfig()
cfg = {"include": ["models.yaml"], "model_list": [], "litellm_settings": {}}
result = await pc._process_includes(cfg, config_file_path=str(tmp_path / "config.yaml"))
assert result == {
"model_list": [{"model_name": "gpt-4"}],
"litellm_settings": {},
}
@pytest.mark.asyncio
async def test_ProxyConfig__process_includes_missing_file_raises(tmp_path):
pc = ProxyConfig()
with pytest.raises(FileNotFoundError):
await pc._process_includes({"include": ["nope.yaml"]}, config_file_path=str(tmp_path / "config.yaml"))
@pytest.mark.asyncio
async def test_ProxyConfig__process_includes_follows_nested_includes(tmp_path):
(tmp_path / "models.yaml").write_text("include:\n - more_models.yaml\nmodel_list:\n - model_name: first\n")
(tmp_path / "more_models.yaml").write_text("model_list:\n - model_name: second\n")
result = await ProxyConfig()._process_includes(
{"include": ["models.yaml"]}, config_file_path=str(tmp_path / "config.yaml")
)
assert result == {"model_list": [{"model_name": "first"}, {"model_name": "second"}]}
@pytest.mark.asyncio
async def test_ProxyConfig__process_includes_resolves_a_nested_include_next_to_its_own_file(tmp_path):
(tmp_path / "shared").mkdir()
(tmp_path / "shared" / "models.yaml").write_text(
"include:\n - more_models.yaml\nmodel_list:\n - model_name: first\n"
)
(tmp_path / "shared" / "more_models.yaml").write_text("model_list:\n - model_name: second\n")
(tmp_path / "more_models.yaml").write_text("model_list:\n - model_name: wrong-directory\n")
result = await ProxyConfig()._process_includes(
{"include": ["shared/models.yaml"]}, config_file_path=str(tmp_path / "config.yaml")
)
assert result == {"model_list": [{"model_name": "first"}, {"model_name": "second"}]}
@pytest.mark.asyncio
async def test_ProxyConfig__process_includes_still_reads_a_nested_include_left_beside_the_root_config(tmp_path):
(tmp_path / "shared").mkdir()
(tmp_path / "shared" / "models.yaml").write_text(
"include:\n - more_models.yaml\nmodel_list:\n - model_name: first\n"
)
(tmp_path / "more_models.yaml").write_text("model_list:\n - model_name: second\n")
result = await ProxyConfig()._process_includes(
{"include": ["shared/models.yaml"]}, config_file_path=str(tmp_path / "config.yaml")
)
assert result == {"model_list": [{"model_name": "first"}, {"model_name": "second"}]}
@pytest.mark.asyncio
async def test_ProxyConfig__process_includes_names_both_files_when_a_nested_include_matches_two(tmp_path, caplog):
(tmp_path / "shared").mkdir()
(tmp_path / "shared" / "models.yaml").write_text(
"include:\n - more_models.yaml\nmodel_list:\n - model_name: first\n"
)
(tmp_path / "shared" / "more_models.yaml").write_text("model_list:\n - model_name: next-to-the-declaring-file\n")
(tmp_path / "more_models.yaml").write_text("model_list:\n - model_name: next-to-the-root-config\n")
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
result = await ProxyConfig()._process_includes(
{"include": ["shared/models.yaml"]}, config_file_path=str(tmp_path / "config.yaml")
)
assert result == {"model_list": [{"model_name": "first"}, {"model_name": "next-to-the-declaring-file"}]}
assert [
record
for record in caplog.records
if str(tmp_path / "shared" / "more_models.yaml") in record.getMessage()
and str(tmp_path / "more_models.yaml") in record.getMessage()
]
@pytest.mark.asyncio
async def test_ProxyConfig__process_includes_merges_a_shared_file_once(tmp_path):
(tmp_path / "shared.yaml").write_text("model_list:\n - model_name: shared\n")
(tmp_path / "a.yaml").write_text("include:\n - shared.yaml\n")
(tmp_path / "b.yaml").write_text("include:\n - ./shared.yaml\n")
result = await ProxyConfig()._process_includes(
{"include": ["a.yaml", "b.yaml"]}, config_file_path=str(tmp_path / "config.yaml")
)
assert result == {"model_list": [{"model_name": "shared"}]}
@pytest.mark.asyncio
async def test_ProxyConfig__process_includes_names_the_file_when_it_is_not_a_mapping(tmp_path):
(tmp_path / "models.yaml").write_text("- model_name: gpt-4\n")
with pytest.raises(ValueError, match=re.escape(str(tmp_path / "models.yaml"))):
await ProxyConfig()._process_includes(
{"include": ["models.yaml"]}, config_file_path=str(tmp_path / "config.yaml")
)
@pytest.mark.asyncio
async def test_ProxyConfig__process_includes_terminates_on_a_cycle(tmp_path):
(tmp_path / "a.yaml").write_text("include:\n - b.yaml\nmodel_list:\n - model_name: from-a\n")
(tmp_path / "b.yaml").write_text("include:\n - a.yaml\nmodel_list:\n - model_name: from-b\n")
result = await asyncio.wait_for(
ProxyConfig()._process_includes({"include": ["a.yaml"]}, config_file_path=str(tmp_path / "config.yaml")),
timeout=10,
)
assert result == {"model_list": [{"model_name": "from-a"}, {"model_name": "from-b"}]}
# ---------------------------------------------------------------------------
# ProxyConfig.save_config
# ---------------------------------------------------------------------------
_CONFIG_VALUE: Final = TypeAdapter(dict[str, JsonValue])
@dataclass(frozen=True, slots=True)
class _ConfigRow:
param_value: dict[str, JsonValue] | str
class _ConfigTable:
def __init__(self, rows: Mapping[str, Mapping[str, JsonValue] | str]) -> None:
self.rows = {
param_name: value if isinstance(value, str) else _CONFIG_VALUE.validate_python(value)
for param_name, value in rows.items()
}
self.upserted_param_names: list[str] = []
self._section_lock = asyncio.Lock()
async def find_first(self, *, where: Mapping[str, str]) -> _ConfigRow | None:
value: Final = self.rows.get(where["param_name"])
await asyncio.sleep(0)
return _ConfigRow(param_value=value) if value is not None else None
async def upsert(self, *, where: Mapping[str, str], data: Mapping[str, Mapping[str, str]]) -> _ConfigRow:
param_name: Final = where["param_name"]
value: Final = _CONFIG_VALUE.validate_json(data["update"]["param_value"])
self.rows[param_name] = value
self.upserted_param_names.append(param_name)
return _ConfigRow(param_value=value)
class _ConfigTransaction:
def __init__(self, table: _ConfigTable) -> None:
self.litellm_config: Final = table
self._section_lock: Final = table._section_lock
self._locked = False
async def __aenter__(self) -> _ConfigTransaction:
return self
async def __aexit__(self, *_: object) -> None:
if self._locked:
self._section_lock.release()
async def query_raw(self, _: str, __: str) -> None:
await self._section_lock.acquire()
self._locked = True
@dataclass(frozen=True, slots=True)
class _ConfigDb:
litellm_config: _ConfigTable
def tx(self) -> _ConfigTransaction:
return _ConfigTransaction(self.litellm_config)
@dataclass(frozen=True, slots=True)
class _ConfigPrisma:
db: _ConfigDb
def tx(self) -> _ConfigTransaction:
return self.db.tx()
async def insert_data(self, *, data: Mapping[str, object], table_name: str) -> None:
if table_name != "config":
raise AssertionError(f"Expected config write, got {table_name}")
for param_name, value in data.items():
self.db.litellm_config.rows[param_name] = _CONFIG_VALUE.validate_python(value)
self.db.litellm_config.upserted_param_names.append(param_name)
def _db_backed_proxy_config(
monkeypatch, rows: Mapping[str, Mapping[str, JsonValue]]
) -> tuple[ProxyConfig, _ConfigTable]:
table: Final = _ConfigTable(rows)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", _ConfigPrisma(db=_ConfigDb(litellm_config=table)))
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True)
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"store_model_in_db": True})
monkeypatch.setattr("litellm.proxy.proxy_server.invalidate_config_param", AsyncMock())
return ProxyConfig(), table
@pytest.mark.asyncio
async def test_ProxyConfig_save_config_merges_changed_keys_without_copying_file_settings(monkeypatch):
proxy_config, table = _db_backed_proxy_config(monkeypatch, {"general_settings": {"db_only": "stored"}})
baseline: Final = {
"model_list": [],
"general_settings": {"max_parallel_requests": 5, "file_only": "yaml", "allowed_ips": []},
"router_settings": {"num_retries": 1},
"litellm_settings": {"drop_params": True},
}
proxy_config.update_config_state(config=baseline)
changed: Final = {
**baseline,
"general_settings": {**baseline["general_settings"], "allowed_ips": ["127.0.0.1"]},
}
await proxy_config.save_config(changed)
assert table.rows == {"general_settings": {"db_only": "stored", "allowed_ips": ["127.0.0.1"]}}
assert table.upserted_param_names == ["general_settings"]
@pytest.mark.asyncio
async def test_ProxyConfig_save_config_skips_unchanged_config(monkeypatch):
proxy_config, table = _db_backed_proxy_config(monkeypatch, {"general_settings": {"db_only": "stored"}})
baseline: Final = {
"model_list": [],
"general_settings": {"max_parallel_requests": 5},
"router_settings": {"num_retries": 1},
"litellm_settings": {"drop_params": True},
}
proxy_config.update_config_state(config=baseline)
await proxy_config.save_config(baseline)
assert table.rows == {"general_settings": {"db_only": "stored"}}
assert table.upserted_param_names == []
@pytest.mark.asyncio
async def test_ProxyConfig_save_config_skips_unchanged_unmanaged_values(monkeypatch):
proxy_config, table = _db_backed_proxy_config(monkeypatch, {})
baseline: Final = {"general_settings": {}, "guardrails": {"enabled": True}}
proxy_config.update_config_state(config=baseline)
await proxy_config.save_config(baseline)
assert table.rows == {}
assert table.upserted_param_names == []
@pytest.mark.asyncio
async def test_ProxyConfig_save_config_leaves_omitted_sections_unchanged(monkeypatch):
proxy_config, table = _db_backed_proxy_config(
monkeypatch,
{"general_settings": {"allowed_ips": ["10.0.0.1"], "db_only": "stored"}},
)
proxy_config.update_config_state(
config={"general_settings": {"allowed_ips": ["10.0.0.1"]}, "router_settings": {"num_retries": 1}}
)
await proxy_config.save_config({"router_settings": {"num_retries": 2}})
assert table.rows == {
"general_settings": {"allowed_ips": ["10.0.0.1"], "db_only": "stored"},
"router_settings": {"num_retries": 2},
}
assert table.upserted_param_names == ["router_settings"]
@pytest.mark.asyncio
async def test_ProxyConfig_save_config_decodes_a_serialized_config_row(monkeypatch):
proxy_config, table = _db_backed_proxy_config(monkeypatch, {"general_settings": '{"db_only":"stored"}'})
proxy_config.update_config_state(config={"general_settings": {"allowed_ips": []}})
await proxy_config.save_config({"general_settings": {"allowed_ips": ["127.0.0.1"]}})
assert table.rows == {"general_settings": {"db_only": "stored", "allowed_ips": ["127.0.0.1"]}}
assert table.upserted_param_names == ["general_settings"]
@pytest.mark.asyncio
async def test_ProxyConfig_save_config_serializes_concurrent_changes_to_one_section(monkeypatch):
first, table = _db_backed_proxy_config(monkeypatch, {"general_settings": {"a": 0, "b": 0}})
second: Final = ProxyConfig()
baseline: Final = {"general_settings": {"a": 0, "b": 0}}
first.update_config_state(config=baseline)
second.update_config_state(config=baseline)
await asyncio.gather(
first.save_config({"general_settings": {"a": 1, "b": 0}}),
second.save_config({"general_settings": {"a": 0, "b": 1}}),
)
assert table.rows == {"general_settings": {"a": 1, "b": 1}}
@pytest.mark.asyncio
async def test_ProxyConfig_save_config_updates_the_baseline_after_a_save(monkeypatch):
proxy_config, table = _db_backed_proxy_config(monkeypatch, {})
proxy_config.update_config_state(config={"general_settings": {}})
await proxy_config.save_config({"general_settings": {"removed_key": True}})
await proxy_config.save_config({"general_settings": {}})
assert table.rows == {"general_settings": {}}
@pytest.mark.asyncio
async def test_ProxyConfig_save_config_keeps_omitted_sections_in_its_next_baseline(monkeypatch):
proxy_config, table = _db_backed_proxy_config(monkeypatch, {"general_settings": {"allowed_ips": ["10.0.0.1"]}})
proxy_config.update_config_state(
config={"general_settings": {"allowed_ips": ["10.0.0.1"]}, "router_settings": {"num_retries": 1}}
)
await proxy_config.save_config({"router_settings": {"num_retries": 2}})
await proxy_config.save_config({"general_settings": {}})
assert table.rows == {"general_settings": {}, "router_settings": {"num_retries": 2}}
@pytest.mark.asyncio
async def test_ProxyConfig_save_config_uses_the_baseline_from_the_loaded_config(tmp_path, monkeypatch):
config_file: Final = tmp_path / "config.yaml"
config_file.write_text("general_settings:\n yaml_only: true\n")
proxy_config: Final = ProxyConfig()
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False)
first: Final = await proxy_config.get_config(config_file_path=str(config_file))
second: Final = await proxy_config.get_config(config_file_path=str(config_file))
table: Final = _ConfigTable({})
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", _ConfigPrisma(db=_ConfigDb(litellm_config=table)))
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True)
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"store_model_in_db": True})
monkeypatch.setattr("litellm.proxy.proxy_server.invalidate_config_param", AsyncMock())
first["general_settings"]["first"] = True
second["general_settings"]["second"] = True
await proxy_config.save_config(second)
await proxy_config.save_config(first)
assert table.rows == {"general_settings": {"second": True, "first": True}}
@pytest.mark.asyncio
async def test_ProxyConfig_save_config_accepts_non_json_model_metadata(monkeypatch):
proxy_config, table = _db_backed_proxy_config(monkeypatch, {})
proxy_config.update_config_state(config={"general_settings": {"allowed_ips": []}})
config: Final = {
"model_list": [{"model_name": "date-model", "model_info": {"created_at": datetime(2026, 1, 1)}}],
"general_settings": {"allowed_ips": ["127.0.0.1"]},
}
await proxy_config.save_config(config)
assert table.rows == {"general_settings": {"allowed_ips": ["127.0.0.1"]}}
@pytest.mark.asyncio
async def test_ProxyConfig_save_config_writes_only_changed_router_settings(monkeypatch):
proxy_config, table = _db_backed_proxy_config(monkeypatch, {"router_settings": {"db_only": "stored"}})
baseline: Final = {
"model_list": [],
"general_settings": {"max_parallel_requests": 5},
"router_settings": {"num_retries": 1},
"litellm_settings": {"drop_params": True},
}
proxy_config.update_config_state(config=baseline)
changed: Final = {**baseline, "router_settings": {"num_retries": 2}}
await proxy_config.save_config(changed)
assert table.rows == {"router_settings": {"db_only": "stored", "num_retries": 2}}
assert table.upserted_param_names == ["router_settings"]
@pytest.mark.asyncio
async def test_ProxyConfig_save_config_removes_a_key_only_when_the_db_has_it(monkeypatch):
proxy_config, table = _db_backed_proxy_config(
monkeypatch, {"general_settings": {"removed_key": "db", "db_only": "stored"}}
)
baseline: Final = {"general_settings": {"removed_key": "yaml", "file_only": "yaml"}}
proxy_config.update_config_state(config=baseline)
changed: Final = {"general_settings": {"file_only": "yaml"}}
await proxy_config.save_config(changed)
assert table.rows == {"general_settings": {"db_only": "stored"}}
assert table.upserted_param_names == ["general_settings"]
@pytest.mark.asyncio
async def test_ProxyConfig_save_config_keeps_an_unstored_removed_key_as_a_noop(monkeypatch):
proxy_config, table = _db_backed_proxy_config(monkeypatch, {"general_settings": {"db_only": "stored"}})
baseline: Final = {"general_settings": {"file_only": "yaml"}}
proxy_config.update_config_state(config=baseline)
await proxy_config.save_config({"general_settings": {}})
assert table.rows == {"general_settings": {"db_only": "stored"}}
assert table.upserted_param_names == []
@pytest.mark.asyncio
async def test_ProxyConfig_get_config_keeps_state_separate_from_returned_config(tmp_path, monkeypatch):
config_file: Final = tmp_path / "config.yaml"
config_file.write_text("general_settings:\n max_parallel_requests: 5\n")
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False)
monkeypatch.setattr("litellm.proxy.proxy_server.user_config_file_path", str(config_file))
proxy_config: Final = ProxyConfig()
loaded: Final = await proxy_config.get_config(config_file_path=str(config_file))
loaded["general_settings"]["max_parallel_requests"] = 6
assert proxy_config.get_config_state()["general_settings"]["max_parallel_requests"] == 5
def test_ProxyConfig_update_config_state_keeps_a_copy_of_its_input():
source: Final = {"general_settings": {"max_parallel_requests": 5}}
proxy_config: Final = ProxyConfig()
proxy_config.update_config_state(config=source)
source["general_settings"]["max_parallel_requests"] = 6
assert proxy_config.get_config_state()["general_settings"]["max_parallel_requests"] == 5
@pytest.mark.asyncio
async def test_ProxyConfig_save_config_writes_yaml_when_no_db(tmp_path, monkeypatch):
target = tmp_path / "out.yaml"
monkeypatch.setattr("litellm.proxy.proxy_server.user_config_file_path", str(target))
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False)
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
pc = ProxyConfig()
cfg = {"model_list": [], "general_settings": {"a": 1}, "litellm_settings": {}}
await pc.save_config(cfg)
import yaml as _yaml
loaded = _yaml.safe_load(target.read_text())
assert loaded == cfg
@pytest.mark.asyncio
async def test_ProxyConfig_save_config_writes_a_loadable_yaml_for_a_loaded_config(tmp_path, monkeypatch):
config_file: Final = tmp_path / "config.yaml"
config_file.write_text("general_settings:\n max_parallel_requests: 5\n")
monkeypatch.setattr("litellm.proxy.proxy_server.user_config_file_path", str(config_file))
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False)
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
proxy_config: Final = ProxyConfig()
loaded_config: Final = await proxy_config.get_config(config_file_path=str(config_file))
loaded_config["general_settings"]["max_parallel_requests"] = 6
await proxy_config.save_config(loaded_config)
import yaml as _yaml
assert _yaml.safe_load(config_file.read_text()) == {"general_settings": {"max_parallel_requests": 6}}
@pytest.mark.asyncio
async def test_ProxyConfig_save_config_invalid_path_raises(monkeypatch):
monkeypatch.setattr(
"litellm.proxy.proxy_server.user_config_file_path",
"/no/such/dir/out.yaml",
)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False)
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
pc = ProxyConfig()
with pytest.raises(FileNotFoundError):
await pc.save_config({"x": 1})
@pytest.mark.asyncio
async def test_ProxyConfig_save_config_db_omits_environment_variables_by_default(monkeypatch):
proxy_config, table = _db_backed_proxy_config(monkeypatch, {})
baseline: Final = {"model_list": [], "litellm_settings": {}}
proxy_config.update_config_state(config=baseline)
config: Final = {
"model_list": [{"model_name": "gpt-4o"}],
"litellm_settings": {"success_callback": ["langfuse"]},
"environment_variables": {"OPENAI_API_KEY": "sk-from-yaml"},
}
await proxy_config.save_config(config)
assert table.rows == {"litellm_settings": {"success_callback": ["langfuse"]}}
assert table.upserted_param_names == ["litellm_settings"]
assert config["environment_variables"] == {"OPENAI_API_KEY": "sk-from-yaml"}
@pytest.mark.asyncio
async def test_ProxyConfig_save_config_db_persists_environment_variables_when_opted_in(monkeypatch):
proxy_config, table = _db_backed_proxy_config(monkeypatch, {})
proxy_config.update_config_state(config={"litellm_settings": {}})
monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "sk-test-salt-key")
config: Final = {"litellm_settings": {}, "environment_variables": {"OPENAI_API_KEY": "sk-explicit"}}
await proxy_config.save_config(config, include_env_vars=True)
assert set(table.rows["environment_variables"]) == {"OPENAI_API_KEY"}
assert table.rows["environment_variables"]["OPENAI_API_KEY"] != "sk-explicit"
assert table.upserted_param_names == ["environment_variables"]
@pytest.mark.asyncio
async def test_ProxyConfig_save_config_persists_unchanged_environment_variables_when_opted_in(monkeypatch):
proxy_config, table = _db_backed_proxy_config(monkeypatch, {})
config: Final = {
"litellm_settings": {},
"environment_variables": {"OPENAI_API_KEY": "sk-explicit"},
}
monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "sk-test-salt-key")
await proxy_config.save_config(config)
assert table.rows == {}
assert table.upserted_param_names == []
await proxy_config.save_config(config, include_env_vars=True)
assert set(table.rows["environment_variables"]) == {"OPENAI_API_KEY"}
assert table.upserted_param_names == ["environment_variables"]
def _install_fake_config_repo(monkeypatch, existing_row):
"""Route ProxyConfig's ConfigRepository through an in-memory fake that
records the value written to the environment_variables row."""
captured: dict = {}
class _FakeTable:
async def find_first(self, where):
return SimpleNamespace(param_value=existing_row) if existing_row is not None else None
async def upsert(self, where, data):
captured["value"] = json.loads(data["update"]["param_value"])
class _FakeRepo:
def __init__(self, client):
self.table = _FakeTable()
monkeypatch.setattr("litellm.proxy.proxy_server.ConfigRepository", _FakeRepo)
monkeypatch.setattr("litellm.proxy.proxy_server.invalidate_config_param", AsyncMock())
return captured
@pytest.mark.asyncio
async def test_ProxyConfig_save_environment_variables_merges_sets_and_deletes(monkeypatch):
"""The per-key env-var write updates/deletes only the named keys and leaves
every other stored key untouched, so an unrelated env var is never lost or
snapshotted."""
captured = _install_fake_config_repo(
monkeypatch,
existing_row={"EXISTING_KEY": "ciphertext-existing", "UI_LOGO_PATH": "old-logo", "LITELLM_FAVICON_URL": "old"},
)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MagicMock())
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True)
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "sk-test-salt-key")
pc = ProxyConfig()
await pc.save_environment_variables({"UI_LOGO_PATH": "new-logo", "LITELLM_FAVICON_URL": None})
written = captured["value"]
# unrelated key preserved byte-for-byte
assert written["EXISTING_KEY"] == "ciphertext-existing"
# set key updated and encrypted (not the plaintext)
assert "UI_LOGO_PATH" in written and written["UI_LOGO_PATH"] != "new-logo"
# None-valued key deleted
assert "LITELLM_FAVICON_URL" not in written
@pytest.mark.asyncio
async def test_ProxyConfig_save_environment_variables_noop_without_db(monkeypatch):
"""With no DB configured the per-key write must do nothing (never touch the
config repository)."""
captured = _install_fake_config_repo(monkeypatch, existing_row={})
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False)
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
pc = ProxyConfig()
await pc.save_environment_variables({"UI_LOGO_PATH": "x"})
assert "value" not in captured
# ---------------------------------------------------------------------------
# ProxyConfig._check_for_os_environ_vars
# ---------------------------------------------------------------------------
def test_ProxyConfig__check_for_os_environ_vars_substitutes(monkeypatch):
monkeypatch.setenv("MY_TEST_VAR", "secret-value")
pc = ProxyConfig()
cfg = {
"a": "os.environ/MY_TEST_VAR",
"b": 2,
"nested": {"c": "os.environ/MY_TEST_VAR"},
}
out = pc._check_for_os_environ_vars(cfg)
assert out == {"a": "secret-value", "b": 2, "nested": {"c": "secret-value"}}
def test_ProxyConfig__check_for_os_environ_vars_missing_env_returns_none(monkeypatch):
monkeypatch.delenv("NONEXISTENT_TEST_VAR_X", raising=False)
pc = ProxyConfig()
cfg = {"a": "os.environ/NONEXISTENT_TEST_VAR_X"}
out = pc._check_for_os_environ_vars(cfg)
# get_secret returns None when not found — assert observable shape.
assert out["a"] is None
# ---------------------------------------------------------------------------
# ProxyConfig._get_team_config
# ---------------------------------------------------------------------------
def test_ProxyConfig__get_team_config_returns_match():
pc = ProxyConfig()
teams = [
{"team_id": "t1", "max_budget": 10, "model": "gpt-4"},
{"team_id": "t2", "max_budget": 20, "model": "claude"},
]
out = pc._get_team_config(team_id="t1", all_teams_config=teams)
assert out == {"team_id": "t1", "max_budget": 10, "model": "gpt-4"}
def test_ProxyConfig__get_team_config_missing_team_id_raises():
pc = ProxyConfig()
with pytest.raises(Exception, match="team_id missing from team"):
pc._get_team_config(team_id="t1", all_teams_config=[{"no_id_field": True}])
# ---------------------------------------------------------------------------
# ProxyConfig.load_team_config
# ---------------------------------------------------------------------------
def test_ProxyConfig_load_team_config_returns_team_dict():
pc = ProxyConfig()
pc.config = {
"litellm_settings": {
"default_team_settings": [
{"team_id": "ta", "max_budget": 99, "drop_params": True},
]
}
}
out = pc.load_team_config(team_id="ta")
assert out == {"team_id": "ta", "max_budget": 99, "drop_params": True}
def test_ProxyConfig_load_team_config_no_settings_returns_empty():
pc = ProxyConfig()
pc.config = {"litellm_settings": {}}
# Missing entry — happy path returns {} (no default_team_settings).
out = pc.load_team_config(team_id="missing")
assert out == {}
# Error-style: a misconfigured team list without team_id raises.
pc.config = {"litellm_settings": {"default_team_settings": [{"no_id": True}]}}
with pytest.raises(Exception, match="team_id missing from team"):
pc.load_team_config(team_id="anything")
# ---------------------------------------------------------------------------
# ProxyConfig._init_cache
# ---------------------------------------------------------------------------
def test_ProxyConfig__init_cache_sets_litellm_cache(monkeypatch):
pc = ProxyConfig()
monkeypatch.setattr(litellm, "cache", None, raising=False)
pc._init_cache(cache_params={"type": "local"})
snapshot = {
"cache_is_set": litellm.cache is not None,
"cache_type_name": type(litellm.cache).__name__,
"params_used": "local",
}
assert snapshot == {
"cache_is_set": True,
"cache_type_name": "Cache",
"params_used": "local",
}
def test_ProxyConfig__init_cache_invalid_params_raises():
pc = ProxyConfig()
with pytest.raises(AttributeError):
pc._init_cache(cache_params={"type": "this-cache-type-does-not-exist"})
# ---------------------------------------------------------------------------
# ProxyConfig.switch_on_llm_response_caching
# ---------------------------------------------------------------------------
def test_ProxyConfig_switch_on_llm_response_caching_sets_flag(monkeypatch):
pc = ProxyConfig()
fake_router = MagicMock()
fake_router.cache_responses = False
fake_cache = MagicMock()
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", fake_router)
monkeypatch.setattr(litellm, "cache", fake_cache, raising=False)
pc.switch_on_llm_response_caching()
snapshot = {
"cache_responses": fake_router.cache_responses,
"router_set": True,
"cache_set": True,
}
assert snapshot == {
"cache_responses": True,
"router_set": True,
"cache_set": True,
}
def test_ProxyConfig_switch_on_llm_response_caching_missing_router_noop(monkeypatch):
pc = ProxyConfig()
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None)
monkeypatch.setattr(litellm, "cache", None, raising=False)
# No router and no cache — should silently no-op (no raise).
pc.switch_on_llm_response_caching()
# Error-style: prove no router was created.
with pytest.raises(AttributeError):
_ = pc.does_not_exist # type: ignore[attr-defined]
# ---------------------------------------------------------------------------
# ProxyConfig.get_config
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_ProxyConfig_get_config_loads_from_file(tmp_path, monkeypatch):
f = tmp_path / "c.yaml"
f.write_text("model_list: []\ngeneral_settings: {}\nlitellm_settings: {}\n")
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False)
monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False)
pc = ProxyConfig()
cfg = await pc.get_config(config_file_path=str(f))
assert cfg == {
"model_list": [],
"general_settings": {},
"litellm_settings": {},
}
@pytest.mark.asyncio
async def test_ProxyConfig_get_config_from_a_bucket_merges_includes(monkeypatch):
objects = {
"lit6982/config.yaml": {
"include": ["model_config.yaml"],
"general_settings": {"master_key": "sk-1234"},
},
"lit6982/model_config.yaml": {"model_list": [{"model_name": "included-model"}]},
}
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False)
monkeypatch.setattr(
"litellm.proxy.common_utils.load_config_utils.s3_object_reader",
lambda bucket_name: objects.get,
)
monkeypatch.setenv("LITELLM_CONFIG_BUCKET_NAME", "litellm-configs")
monkeypatch.setenv("LITELLM_CONFIG_BUCKET_OBJECT_KEY", "lit6982/config.yaml")
monkeypatch.setenv("LITELLM_CONFIG_BUCKET_TYPE", "s3")
cfg = await ProxyConfig().get_config()
assert cfg["model_list"] == [{"model_name": "included-model"}]
assert "include" not in cfg
@pytest.mark.asyncio
async def test_ProxyConfig_get_config_missing_file_raises(monkeypatch):
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False)
monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False)
pc = ProxyConfig()
with pytest.raises(Exception, match="Config file not found"):
await pc.get_config(config_file_path="/no/such/path.yaml")
# ---------------------------------------------------------------------------
# ProxyConfig._initialize_secret_manager_from_raw_config
# ---------------------------------------------------------------------------
VAULT_SECRET_MANAGER_MODULE = """
import os
from litellm.integrations.custom_secret_manager import CustomSecretManager
VAULT = {"LITELLM_MASTER_KEY": "master-from-vault", "MY_PROVIDER_KEY": "provider-from-vault"}
class VaultSecretManager(CustomSecretManager):
def __init__(self):
super().__init__()
# The loader re-executes this module on every construction, so an in-module counter
# would reset. Append to a file instead, to count constructions across the whole load.
with open(os.environ["VAULT_CONSTRUCTION_LOG"], "a") as f:
f.write("constructed\\n")
def sync_read_secret(self, secret_name, optional_params=None, timeout=None, **kwargs):
return VAULT.get(secret_name)
async def async_read_secret(self, secret_name, optional_params=None, timeout=None, **kwargs):
return VAULT.get(secret_name)
"""
VAULT_BACKED_CONFIG = """
model_list:
- model_name: my-model
litellm_params:
model: openai/gpt-4o-mini
api_key: os.environ/MY_PROVIDER_KEY
general_settings:
master_key: os.environ/LITELLM_MASTER_KEY
key_management_system: custom
key_management_settings:
custom_secret_manager: vault_secret_manager.VaultSecretManager
hosted_keys:
- LITELLM_MASTER_KEY
- MY_PROVIDER_KEY
"""
def _write_vault_backed_config(tmp_path, monkeypatch, config_yaml: str) -> str:
"""Write a config whose secrets live only in a custom secret manager, never in the env."""
(tmp_path / "vault_secret_manager.py").write_text(VAULT_SECRET_MANAGER_MODULE)
config_file = tmp_path / "c.yaml"
config_file.write_text(config_yaml)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False)
monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False)
monkeypatch.delenv("LITELLM_MASTER_KEY", raising=False)
monkeypatch.delenv("MY_PROVIDER_KEY", raising=False)
monkeypatch.setenv("VAULT_CONSTRUCTION_LOG", str(tmp_path / "constructions.log"))
monkeypatch.setattr(litellm, "secret_manager_client", None)
return str(config_file)
def _construction_count(tmp_path) -> int:
log = tmp_path / "constructions.log"
return len(log.read_text().splitlines()) if log.exists() else 0
@pytest.mark.asyncio
async def test_ProxyConfig_get_config_resolves_keys_held_only_by_the_secret_manager(tmp_path, monkeypatch):
"""Regression for GH #35239.
get_config() used to resolve every ``os.environ/<KEY>`` reference and write the result
back into the config before the secret manager was initialized, so any key that lived
only in the manager became a permanent ``None``.
"""
config_file_path = _write_vault_backed_config(tmp_path, monkeypatch, VAULT_BACKED_CONFIG)
cfg = await ProxyConfig().get_config(config_file_path=config_file_path)
assert {
"master_key": cfg["general_settings"]["master_key"],
"api_key": cfg["model_list"][0]["litellm_params"]["api_key"],
"hosted_keys": litellm._key_management_settings.hosted_keys,
} == {
"master_key": "master-from-vault",
"api_key": "provider-from-vault",
"hosted_keys": ["LITELLM_MASTER_KEY", "MY_PROVIDER_KEY"],
}
@pytest.mark.asyncio
async def test_ProxyConfig_load_config_builds_the_secret_manager_exactly_once(tmp_path, monkeypatch):
"""The full startup path must not build the manager, then throw it away and build another.
A discarded client costs a Vault/CyberArk re-auth and leaks a gRPC channel on Google KMS.
"""
config_file_path = _write_vault_backed_config(tmp_path, monkeypatch, VAULT_BACKED_CONFIG)
_router, _model_list, general_settings = await ProxyConfig().load_config(
router=None, config_file_path=config_file_path
)
assert {
"constructions": _construction_count(tmp_path),
"master_key": general_settings["master_key"],
} == {"constructions": 1, "master_key": "master-from-vault"}
@pytest.mark.asyncio
async def test_ProxyConfig_get_config_reuses_an_already_initialized_secret_manager(tmp_path, monkeypatch):
"""get_config() also runs on management-endpoint request paths.
Rebuilding the client on every call would re-execute the custom manager module, drop the
Vault/CyberArk token caches, and leak a gRPC channel per request on Google KMS.
"""
config_file_path = _write_vault_backed_config(tmp_path, monkeypatch, VAULT_BACKED_CONFIG)
await ProxyConfig().get_config(config_file_path=config_file_path)
first_client = litellm.secret_manager_client
second = await ProxyConfig().get_config(config_file_path=config_file_path)
assert {
"client_reused": litellm.secret_manager_client is first_client,
"master_key": second["general_settings"]["master_key"],
} == {"client_reused": True, "master_key": "master-from-vault"}
@pytest.mark.asyncio
async def test_ProxyConfig_get_config_without_key_management_system_leaves_secret_manager_unset(tmp_path, monkeypatch):
"""No ``key_management_system`` means no manager, an unresolvable reference stays None, and
nothing is warned about: with no manager there is nothing to have been absent from."""
config_yaml = VAULT_BACKED_CONFIG.replace(" key_management_system: custom\n", "")
config_file_path = _write_vault_backed_config(tmp_path, monkeypatch, config_yaml)
warn = MagicMock()
monkeypatch.setattr("litellm.proxy.proxy_server.verbose_proxy_logger.warning", warn)
cfg = await ProxyConfig().get_config(config_file_path=config_file_path)
assert {
"master_key": cfg["general_settings"]["master_key"],
"api_key": cfg["model_list"][0]["litellm_params"]["api_key"],
"client": litellm.secret_manager_client,
"warned_about": [call.args[1] for call in warn.call_args_list],
} == {"master_key": None, "api_key": None, "client": None, "warned_about": []}
@pytest.mark.asyncio
async def test_ProxyConfig_get_config_warns_when_a_reference_is_missing_from_the_secret_manager(tmp_path, monkeypatch):
"""A reference the manager cannot resolve is logged, instead of silently becoming None."""
config_yaml = VAULT_BACKED_CONFIG.replace("MY_PROVIDER_KEY", "NOT_IN_VAULT")
config_file_path = _write_vault_backed_config(tmp_path, monkeypatch, config_yaml)
warn = MagicMock()
monkeypatch.setattr("litellm.proxy.proxy_server.verbose_proxy_logger.warning", warn)
cfg = await ProxyConfig().get_config(config_file_path=config_file_path)
assert {
"api_key": cfg["model_list"][0]["litellm_params"]["api_key"],
"warned_about": [call.args[1] for call in warn.call_args_list],
} == {"api_key": None, "warned_about": ["os.environ/NOT_IN_VAULT"]}
@pytest.mark.asyncio
async def test_ProxyConfig_get_config_does_not_warn_for_a_name_outside_hosted_keys(tmp_path, monkeypatch):
"""``hosted_keys`` is an allowlist, so a name outside it is never looked up in the manager.
Warning about it would claim a lookup that never happened, on every optional env-only
reference, on every config reload.
"""
config_yaml = VAULT_BACKED_CONFIG.replace("api_key: os.environ/MY_PROVIDER_KEY", "api_key: os.environ/ENV_ONLY")
config_file_path = _write_vault_backed_config(tmp_path, monkeypatch, config_yaml)
warn = MagicMock()
monkeypatch.setattr("litellm.proxy.proxy_server.verbose_proxy_logger.warning", warn)
cfg = await ProxyConfig().get_config(config_file_path=config_file_path)
assert {
"api_key": cfg["model_list"][0]["litellm_params"]["api_key"],
"client_is_up": litellm.secret_manager_client is not None,
"warned_about": [call.args[1] for call in warn.call_args_list],
} == {"api_key": None, "client_is_up": True, "warned_about": []}
@pytest.mark.asyncio
async def test_ProxyConfig_get_config_does_not_warn_under_write_only_access_mode(tmp_path, monkeypatch):
"""``write_only`` means reads never reach the manager, so an absent name is not its fault.
That mode exists so the manager can store virtual keys while config secrets stay in the
environment, which makes env-only references the expected state rather than an error.
"""
config_yaml = VAULT_BACKED_CONFIG.replace(
" key_management_settings:\n", " key_management_settings:\n access_mode: write_only\n"
)
config_file_path = _write_vault_backed_config(tmp_path, monkeypatch, config_yaml)
warn = MagicMock()
monkeypatch.setattr("litellm.proxy.proxy_server.verbose_proxy_logger.warning", warn)
cfg = await ProxyConfig().get_config(config_file_path=config_file_path)
assert {
"master_key": cfg["general_settings"]["master_key"],
"client_is_up": litellm.secret_manager_client is not None,
"warned_about": [call.args[1] for call in warn.call_args_list],
} == {"master_key": None, "client_is_up": True, "warned_about": []}
# ---------------------------------------------------------------------------
# ProxyConfig.update_config_state / get_config_state
# ---------------------------------------------------------------------------
def test_ProxyConfig_update_config_state_and_get_config_state_roundtrip():
pc = ProxyConfig()
cfg = {"model_list": [], "general_settings": {"x": 1}, "litellm_settings": {}}
pc.update_config_state(config=cfg)
out = pc.get_config_state()
assert out == cfg
# Mutating the returned dict must not affect internal state.
out["model_list"].append({"new": True})
assert pc.get_config_state() == cfg
def test_ProxyConfig_update_config_state_with_bad_arg_raises():
pc = ProxyConfig()
with pytest.raises(TypeError):
pc.update_config_state() # type: ignore[call-arg]
def test_ProxyConfig_get_config_state_handles_undeepcopyable(monkeypatch):
# Pins ProxyConfig.get_config_state — see source for behavior.
pc = ProxyConfig()
class NoCopy:
def __deepcopy__(self, memo):
raise RuntimeError("nope")
pc.config = {"x": NoCopy()} # type: ignore[assignment]
# Exception is caught internally and an empty dict returned.
assert pc.get_config_state() == {}
# ---------------------------------------------------------------------------
# ProxyConfig.load_credential_list
# ---------------------------------------------------------------------------
def test_ProxyConfig_load_credential_list_returns_items():
pc = ProxyConfig()
creds = pc.load_credential_list(
{
"credential_list": [
{
"credential_name": "openai-key",
"credential_info": {"provider": "openai"},
"credential_values": {"api_key": "sk-x"},
}
]
}
)
assert len(creds) == 1
dumped = creds[0].model_dump()
assert dumped == {
"credential_name": "openai-key",
"credential_info": {"provider": "openai"},
"credential_values": {"api_key": "sk-x"},
}
def test_ProxyConfig_load_credential_list_invalid_entry_raises():
pc = ProxyConfig()
with pytest.raises(ValidationError):
pc.load_credential_list({"credential_list": [{"missing_required": True}]})
# ---------------------------------------------------------------------------
# ProxyConfig.parse_search_tools
# ---------------------------------------------------------------------------
def test_ProxyConfig_parse_search_tools_returns_parsed():
pc = ProxyConfig()
cfg = {
"search_tools": [
{
"search_tool_name": "web",
"litellm_params": {"search_provider": "google"},
}
]
}
out = pc.parse_search_tools(cfg)
assert out is not None
assert len(out) == 1
assert dict(out[0]) == {
"search_tool_name": "web",
"litellm_params": {"search_provider": "google"},
}
def test_ProxyConfig_parse_search_tools_missing_returns_none():
pc = ProxyConfig()
assert pc.parse_search_tools({}) is None
def test_ProxyConfig_merge_config_and_db_search_tools_returns_superset():
config_tools = [
{
"search_tool_name": "config-search",
"litellm_params": {"search_provider": "tavily"},
}
]
db_tools = [
{
"search_tool_name": "db-search",
"litellm_params": {
"search_provider": "exa_ai",
"api_key": "fake-db-key",
},
}
]
merged = ProxyConfig._merge_config_and_db_search_tools(
config_search_tools=config_tools,
db_search_tools=db_tools,
)
assert [tool["search_tool_name"] for tool in merged] == ["config-search", "db-search"]
assert merged[1]["litellm_params"]["api_key"] == "fake-db-key"
def test_ProxyConfig_merge_config_and_db_search_tools_prefers_db_duplicate():
config_tools = [
{
"search_tool_name": "shared-search",
"litellm_params": {"search_provider": "tavily"},
},
{
"search_tool_name": "config-only",
"litellm_params": {"search_provider": "perplexity"},
},
]
db_tools = [
{
"search_tool_name": "shared-search",
"litellm_params": {
"search_provider": "exa_ai",
"api_key": "fake-db-key",
},
}
]
merged = ProxyConfig._merge_config_and_db_search_tools(
config_search_tools=config_tools,
db_search_tools=db_tools,
)
assert [tool["search_tool_name"] for tool in merged] == ["config-only", "shared-search"]
assert merged[1]["litellm_params"]["search_provider"] == "exa_ai"
assert merged[1]["litellm_params"]["api_key"] == "fake-db-key"
@pytest.mark.asyncio
async def test_ProxyConfig__init_search_tools_in_db_loads_merged_tools(monkeypatch):
from litellm.proxy import proxy_server
from litellm.router_utils.search_api_router import SearchAPIRouter
pc = ProxyConfig()
pc.update_config_state(
{
"search_tools": [
{
"search_tool_name": "shared-search",
"litellm_params": {"search_provider": "tavily"},
},
{
"search_tool_name": "config-only",
"litellm_params": {"search_provider": "perplexity"},
},
]
}
)
db_tools = [
{
"search_tool_name": "shared-search",
"litellm_params": {
"search_provider": "exa_ai",
"api_key": "fake-db-key",
},
}
]
fake_router = MagicMock()
mock_get_db_tools = AsyncMock(return_value=db_tools)
mock_update_router = AsyncMock()
monkeypatch.setattr(proxy_server, "llm_router", fake_router)
monkeypatch.setattr(
"litellm.proxy.search_endpoints.search_tool_registry.SearchToolRegistry.get_all_search_tools_from_db",
mock_get_db_tools,
)
monkeypatch.setattr(SearchAPIRouter, "update_router_search_tools", mock_update_router)
await pc._init_search_tools_in_db(prisma_client=MagicMock())
mock_get_db_tools.assert_awaited_once()
mock_update_router.assert_awaited_once()
update_kwargs = mock_update_router.await_args.kwargs
assert update_kwargs["router_instance"] is fake_router
assert [tool["search_tool_name"] for tool in update_kwargs["search_tools"]] == [
"config-only",
"shared-search",
]
assert update_kwargs["search_tools"][1]["litellm_params"]["api_key"] == "fake-db-key"
@pytest.mark.asyncio
async def test_ProxyConfig__init_search_tools_in_db_clears_router_when_last_tool_is_deleted(monkeypatch):
"""Deleting the last search tool must clear the router, not leave the tool live in memory."""
from litellm.proxy import proxy_server
pc = ProxyConfig()
pc.update_config_state({})
fake_router = MagicMock()
fake_router.search_tools = [{"search_tool_name": "deleted-search", "litellm_params": {}}]
mock_get_db_tools = AsyncMock(return_value=[])
monkeypatch.setattr(proxy_server, "llm_router", fake_router)
monkeypatch.setattr(
"litellm.proxy.search_endpoints.search_tool_registry.SearchToolRegistry.get_all_search_tools_from_db",
mock_get_db_tools,
)
await pc._init_search_tools_in_db(prisma_client=MagicMock())
mock_get_db_tools.assert_awaited_once()
assert fake_router.search_tools == []
@pytest.mark.asyncio
async def test_ProxyConfig_reload_search_tools_from_db_refreshes_router(monkeypatch):
from litellm.proxy import proxy_server
pc = ProxyConfig()
mock_init = AsyncMock()
monkeypatch.setattr(pc, "_init_search_tools_in_db", mock_init)
monkeypatch.setattr(proxy_server, "prisma_client", MagicMock())
await pc.reload_search_tools_from_db()
mock_init.assert_awaited_once()
@pytest.mark.asyncio
async def test_ProxyConfig_reload_search_tools_from_db_honors_supported_db_objects(monkeypatch):
from litellm.proxy import proxy_server
pc = ProxyConfig()
mock_init = AsyncMock()
monkeypatch.setattr(pc, "_init_search_tools_in_db", mock_init)
monkeypatch.setattr(proxy_server, "prisma_client", MagicMock())
monkeypatch.setattr(proxy_server, "general_settings", {"supported_db_objects": ["models"]})
await pc.reload_search_tools_from_db()
mock_init.assert_not_awaited()
@pytest.mark.asyncio
async def test_ProxyConfig_reload_search_tools_from_db_serializes_overlapping_refreshes(monkeypatch):
"""An older snapshot must not land last and restore a tool a newer refresh deleted."""
import asyncio
from litellm.proxy import proxy_server
pc = ProxyConfig()
pc.update_config_state({})
fake_router = MagicMock()
fake_router.search_tools = []
stale_read_started = asyncio.Event()
fresh_write_committed = asyncio.Event()
snapshots = iter(
(
[{"search_tool_name": "doomed-search", "litellm_params": {}}],
[],
)
)
async def _read_db(**_):
snapshot = next(snapshots)
if not stale_read_started.is_set():
stale_read_started.set()
await fresh_write_committed.wait()
return snapshot
monkeypatch.setattr(proxy_server, "llm_router", fake_router)
monkeypatch.setattr(proxy_server, "prisma_client", MagicMock())
monkeypatch.setattr(
"litellm.proxy.search_endpoints.search_tool_registry.SearchToolRegistry.get_all_search_tools_from_db",
_read_db,
)
stale = asyncio.create_task(pc.reload_search_tools_from_db())
await stale_read_started.wait()
deleter = asyncio.create_task(pc.reload_search_tools_from_db())
await asyncio.sleep(0)
fresh_write_committed.set()
await asyncio.gather(stale, deleter)
assert fake_router.search_tools == []
@pytest.mark.asyncio
async def test_ProxyConfig_reload_search_tools_from_db_noops_without_prisma(monkeypatch):
from litellm.proxy import proxy_server
pc = ProxyConfig()
mock_init = AsyncMock()
monkeypatch.setattr(pc, "_init_search_tools_in_db", mock_init)
monkeypatch.setattr(proxy_server, "prisma_client", None)
await pc.reload_search_tools_from_db()
mock_init.assert_not_awaited()
# ---------------------------------------------------------------------------
# ProxyConfig._load_environment_variables
# ---------------------------------------------------------------------------
def test_ProxyConfig__load_environment_variables_sets_env(monkeypatch):
monkeypatch.delenv("TEST_LOAD_ENV_X", raising=False)
pc = ProxyConfig()
pc._load_environment_variables({"environment_variables": {"TEST_LOAD_ENV_X": "hello"}})
result = {
"TEST_LOAD_ENV_X": os.environ.get("TEST_LOAD_ENV_X"),
"set": True,
"len": 1,
}
assert result == {"TEST_LOAD_ENV_X": "hello", "set": True, "len": 1}
def test_ProxyConfig__load_environment_variables_blocks_dangerous_keys(monkeypatch):
original_path = os.environ.get("PATH", "")
pc = ProxyConfig()
pc._load_environment_variables({"environment_variables": {"PATH": "/evil/bin"}})
# PATH must be unchanged — it's a blocked key.
assert os.environ.get("PATH", "") == original_path
# ---------------------------------------------------------------------------
# ProxyConfig.load_config
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
@pytest.mark.parametrize(
("flag", "system"), (("use_google_kms", "google_kms"), ("use_azure_key_vault", "azure_key_vault"))
)
async def test_load_config_legacy_secret_manager_flags_capture_the_initialized_client(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch, flag: str, system: str
) -> None:
if system == "azure_key_vault":
client_type: Final = pytest.importorskip("azure.keyvault.secrets").SecretClient
else:
client_type: Final = pytest.importorskip("google.cloud.kms_v1").KeyManagementServiceClient
from litellm.rust_bridge.secret_manager import native_secret_manager_config
credentials_file: Final = tmp_path / "credentials.json"
credentials_file.write_text(
json.dumps(
{
"type": "authorized_user",
"client_id": "test-client",
"client_secret": "test-secret",
"refresh_token": "test",
}
)
)
config_file: Final = tmp_path / "legacy-secret-manager.yaml"
config_file.write_text(
f"model_list: []\ngeneral_settings:\n {flag}: true\n key_management_settings:\n access_mode: write_only\n"
)
monkeypatch.setenv("GOOGLE_APPLICATION_CREDENTIALS", str(credentials_file))
monkeypatch.setenv("GOOGLE_KMS_RESOURCE_NAME", "projects/test/locations/global/keyRings/test/cryptoKeys/test")
monkeypatch.setenv("AZURE_KEY_VAULT_URI", "https://test.vault.azure.net")
monkeypatch.setattr(litellm, "secret_manager_client", None)
monkeypatch.setattr(litellm, "_key_management_system", None)
monkeypatch.setattr(litellm, "_google_kms_resource_name", None)
monkeypatch.setattr(litellm, "_key_management_settings", litellm._key_management_settings)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False)
monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False)
await ProxyConfig().load_config(router=None, config_file_path=str(config_file))
client: Final = litellm.secret_manager_client
assert isinstance(client, client_type)
try:
captured: Final = native_secret_manager_config(client)
assert captured is not None
assert captured.system == system
assert dict(captured.environment)["GOOGLE_APPLICATION_CREDENTIALS"] == str(credentials_file)
assert litellm._key_management_system is not None
assert litellm._key_management_system.value == system
finally:
if system == "azure_key_vault":
client.close()
else:
client.transport.close()
@pytest.mark.asyncio
@pytest.mark.parametrize("flag", ("null", "false"))
async def test_load_config_disabled_google_kms_does_not_initialize_a_manager(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch, flag: str
) -> None:
config_file: Final = tmp_path / "disabled-kms.yaml"
config_file.write_text(f"model_list: []\ngeneral_settings:\n use_google_kms: {flag}\n")
monkeypatch.setattr(litellm, "secret_manager_client", None)
monkeypatch.setattr(litellm, "_key_management_system", None)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False)
monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False)
monkeypatch.delenv("GOOGLE_APPLICATION_CREDENTIALS", raising=False)
_router, model_list, general_settings = await ProxyConfig().load_config(
router=None, config_file_path=str(config_file)
)
assert model_list == []
assert general_settings["use_google_kms"] is (None if flag == "null" else False)
assert litellm.secret_manager_client is None
assert litellm._key_management_system is None
@pytest.mark.asyncio
async def test_ProxyConfig_load_config_minimal_yaml(tmp_path, monkeypatch):
f = tmp_path / "c.yaml"
f.write_text("model_list: []\ngeneral_settings: {}\nlitellm_settings: {}\n")
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False)
monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False)
pc = ProxyConfig()
try:
await pc.load_config(router=None, config_file_path=str(f))
raised = False
except Exception:
raised = True
snapshot = {
"raised": raised,
"config_loaded": pc.config is not None,
"model_list_key_present": "model_list" in pc.config,
}
assert snapshot == {
"raised": False,
"config_loaded": True,
"model_list_key_present": True,
}
@pytest.mark.asyncio
@pytest.mark.parametrize("setting", ["true", "false", "null", "'true'", None])
async def test_load_config_logs_disabled_budget_reservation_once(tmp_path, monkeypatch, caplog, setting):
config_file = tmp_path / "budget.yaml"
flag = f" disable_budget_reservation: {setting}\n" if setting is not None else ""
config_file.write_text("model_list: []\nlitellm_settings: {}\ngeneral_settings:\n master_key: null\n" + flag)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False)
monkeypatch.setattr("litellm.constants.budget_reservation_disabled_info_emitted", False)
monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False)
config = ProxyConfig()
with caplog.at_level(logging.INFO, logger="LiteLLM Proxy"):
for _ in range(3):
await config.load_config(router=None, config_file_path=str(config_file))
records = [record for record in caplog.records if "disable_budget_reservation is enabled" in record.message]
assert [record.levelno for record in records] == ([logging.INFO] if setting == "true" else [])
@pytest.mark.asyncio
async def test_ProxyConfig_load_config_resolves_router_settings_plugins(tmp_path, monkeypatch):
"""Regression: router_settings.plugins dotted-path strings must be resolved to
live RoutingPlugin instances on the created Router. Previously they were passed
through as raw strings and only blew up at request time when the pipeline tried
to `await "some.string".run(context)`."""
plugin_file = tmp_path / "rs_plugin.py"
plugin_file.write_text(
"class _Plugin:\n async def run(self, context):\n return context\n\nrs_plugin_instance = _Plugin()\n"
)
f = tmp_path / "c.yaml"
f.write_text(
"model_list: []\n"
"general_settings: {}\n"
"litellm_settings: {}\n"
"router_settings:\n"
" plugins:\n"
" - rs_plugin.rs_plugin_instance\n"
)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False)
monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False)
router, _model_list, _general_settings = await ProxyConfig().load_config(router=None, config_file_path=str(f))
assert len(router.routing_plugins) == 1
assert type(router.routing_plugins[0]).__name__ == "_Plugin"
@pytest.mark.asyncio
async def test_ProxyConfig_load_config_rejects_bad_router_settings_plugin(tmp_path, monkeypatch):
plugin_file = tmp_path / "bad_rs_plugin.py"
plugin_file.write_text("not_a_plugin = object()\n")
f = tmp_path / "c.yaml"
f.write_text(
"model_list: []\n"
"general_settings: {}\n"
"litellm_settings: {}\n"
"router_settings:\n"
" plugins:\n"
" - bad_rs_plugin.not_a_plugin\n"
)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False)
monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False)
with pytest.raises(ValueError, match="does not implement the RoutingPlugin interface"):
await ProxyConfig().load_config(router=None, config_file_path=str(f))
@pytest.mark.asyncio
async def test_ProxyConfig_load_config_wires_general_settings_url_validation(tmp_path, monkeypatch):
"""Regression for #26599: SSRF settings in general_settings must reach litellm globals."""
f = tmp_path / "c.yaml"
f.write_text(
"model_list: []\n"
"general_settings:\n"
" user_url_validation: false\n"
" user_url_allowed_hosts:\n"
" - internal.corp\n"
" provider_url_destination_allowed_hosts:\n"
" - api.example.com\n"
)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False)
monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False)
original_validation = litellm.user_url_validation
original_hosts = list(litellm.user_url_allowed_hosts)
original_provider_hosts = list(litellm.provider_url_destination_allowed_hosts)
try:
await ProxyConfig().load_config(router=None, config_file_path=str(f))
assert litellm.user_url_validation is False
assert litellm.user_url_allowed_hosts == ["internal.corp"]
assert litellm.provider_url_destination_allowed_hosts == ["api.example.com"]
finally:
litellm.user_url_validation = original_validation
litellm.user_url_allowed_hosts = original_hosts
litellm.provider_url_destination_allowed_hosts = original_provider_hosts
@pytest.mark.asyncio
async def test_ssrf_block_message_names_a_config_section_load_config_honors(tmp_path, monkeypatch):
"""Regression for LIT-8349: the remediation in the SSRF block message must point at a section that works."""
from litellm.litellm_core_utils.url_utils import SSRFError, validate_url
monkeypatch.setattr(litellm, "user_url_allowed_hosts", [])
monkeypatch.setattr(litellm, "user_url_validation", True)
with pytest.raises(SSRFError) as blocked:
validate_url("http://10.96.3.245:10002/agent.json")
section_match = re.search(r"add the host to `user_url_allowed_hosts` in (\w+)\.", str(blocked.value))
assert section_match is not None, str(blocked.value)
section: Final = section_match.group(1)
assert section == "litellm_settings", f"block message points admins at {section}, which the docs contradict"
f = tmp_path / "c.yaml"
f.write_text(f"model_list: []\n{section}:\n user_url_allowed_hosts:\n - '10.96.3.245:10002'\n")
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False)
monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False)
await ProxyConfig().load_config(router=None, config_file_path=str(f))
assert litellm.user_url_allowed_hosts == ["10.96.3.245:10002"], f"{section} did not apply the allowlist"
assert validate_url("http://10.96.3.245:10002/agent.json") == (
"http://10.96.3.245:10002/agent.json",
"10.96.3.245:10002",
)
@pytest.mark.asyncio
async def test_ProxyConfig_load_config_wires_config_reload_interval(tmp_path, monkeypatch):
"""general_settings.proxy_config_reload_interval_seconds must reach the proxy_server
module global that schedules the DB config-reload jobs, so operators can tune multi-pod
convergence from config.yaml."""
import litellm.proxy.proxy_server as proxy_server
f = tmp_path / "c.yaml"
f.write_text(
"model_list: []\ngeneral_settings:\n proxy_config_reload_interval_seconds: 47\nlitellm_settings: {}\n"
)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False)
monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False)
original = proxy_server.proxy_config_reload_interval_seconds
try:
await ProxyConfig().load_config(router=None, config_file_path=str(f))
assert proxy_server.proxy_config_reload_interval_seconds == 47
finally:
proxy_server.proxy_config_reload_interval_seconds = original
@pytest.mark.asyncio
async def test_ProxyConfig_load_config_missing_file_raises(monkeypatch):
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False)
monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False)
pc = ProxyConfig()
with pytest.raises(Exception, match="Config file not found"):
await pc.load_config(router=None, config_file_path="/no/file.yaml")
@pytest.mark.asyncio
async def test_ProxyConfig_load_config_forwards_callback_specific_params(tmp_path, monkeypatch):
"""Regression: callback_settings from config must be forwarded to
initialize_callbacks_on_proxy as callback_specific_params.
Callbacks like DatadogCostManagementLogger read their init params (e.g.
cost_tag_keys) from callback_specific_params[<callback_name>]. If the
argument is dropped at the call site, they silently initialize with empty
params and the configured allowlist never takes effect.
"""
f = tmp_path / "c.yaml"
f.write_text(
"model_list: []\n"
"general_settings: {}\n"
"callback_settings:\n"
" datadog_cost_management:\n"
" cost_tag_keys:\n"
" - capability\n"
" - platform\n"
" - ai_product\n"
"litellm_settings:\n"
' callbacks: ["datadog_cost_management"]\n'
)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False)
monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False)
captured = {}
def _fake_initialize_callbacks_on_proxy(**kwargs):
captured.update(kwargs)
monkeypatch.setattr(
"litellm.proxy.proxy_server.initialize_callbacks_on_proxy",
_fake_initialize_callbacks_on_proxy,
)
pc = ProxyConfig()
await pc.load_config(router=None, config_file_path=str(f))
# The callbacks branch must forward the loaded callback_settings.
assert captured.get("callback_specific_params") == {
"datadog_cost_management": {"cost_tag_keys": ["capability", "platform", "ai_product"]}
}
@pytest.mark.asyncio
async def test_ProxyConfig_load_config_blank_callback_settings_does_not_crash(tmp_path, monkeypatch):
"""Regression: `callback_settings:` with no body loads as None because
dict.get() only falls back to the default when the key is absent. The None
was forwarded verbatim to initialize_callbacks_on_proxy, where the first
`"<name>" in callback_specific_params` membership test raised
TypeError: argument of type 'NoneType' is not iterable, aborting startup.
Startup must succeed and the callback must initialize with its defaults.
"""
f = tmp_path / "c.yaml"
f.write_text(
"model_list: []\n"
"general_settings: {}\n"
"callback_settings:\n"
"litellm_settings:\n"
' callbacks: ["compression_interception"]\n'
)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False)
monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False)
from litellm.integrations.compression_interception.handler import (
CompressionInterceptionLogger,
)
original_callbacks = list(litellm.callbacks) if isinstance(litellm.callbacks, list) else []
litellm.callbacks = []
try:
pc = ProxyConfig()
await pc.load_config(router=None, config_file_path=str(f))
assert any(isinstance(c, CompressionInterceptionLogger) for c in litellm.callbacks)
finally:
litellm.callbacks = original_callbacks
# ---------------------------------------------------------------------------
# ProxyConfig._init_non_llm_configs
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_ProxyConfig__init_non_llm_configs_empty_config():
pc = ProxyConfig()
try:
await pc._init_non_llm_configs(config={}, config_file_path=None)
raised = False
except Exception:
raised = True
snapshot = {
"raised": raised,
"worker_registry_len": len(pc.worker_registry),
"is_list": isinstance(pc.worker_registry, list),
}
assert snapshot == {"raised": False, "worker_registry_len": 0, "is_list": True}
@pytest.mark.asyncio
async def test_ProxyConfig__init_non_llm_configs_premium_invalid_worker_registry_raises(monkeypatch):
monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True)
pc = ProxyConfig()
with pytest.raises(ValidationError):
await pc._init_non_llm_configs(
config={"worker_registry": [{"totally": "invalid"}]},
config_file_path=None,
)
@pytest.mark.asyncio
async def test_ProxyConfig__init_non_llm_configs_worker_registry_requires_premium(monkeypatch):
monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", False)
pc = ProxyConfig()
with pytest.raises(ValueError, match="Trying to use `worker_registry`You must be a LiteLLM") as exc_info:
await pc._init_non_llm_configs(
config={"worker_registry": [{"worker_id": "worker-a", "name": "Worker A", "url": "http://localhost:4001"}]},
config_file_path=None,
)
message = str(exc_info.value)
assert "worker_registry" in message
assert CommonProxyErrors.not_premium_user.value in message
assert pc.worker_registry == []
@pytest.mark.asyncio
async def test_ProxyConfig__init_non_llm_configs_worker_registry_loads_for_premium(monkeypatch):
monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True)
pc = ProxyConfig()
await pc._init_non_llm_configs(
config={
"worker_registry": [
{"worker_id": "worker-a", "name": "Worker A", "url": "http://localhost:4001"},
{"worker_id": "worker-b", "name": "Worker B", "url": "https://worker-b.example.com"},
]
},
config_file_path=None,
)
assert [(w.worker_id, w.name, w.url) for w in pc.worker_registry] == [
("worker-a", "Worker A", "http://localhost:4001"),
("worker-b", "Worker B", "https://worker-b.example.com"),
]
@pytest.mark.parametrize("premium", [True, False])
@pytest.mark.asyncio
async def test_ProxyConfig__init_non_llm_configs_no_worker_registry_is_never_gated(monkeypatch, premium):
monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", premium)
pc = ProxyConfig()
await pc._init_non_llm_configs(config={}, config_file_path=None)
assert pc.worker_registry == []
# ---------------------------------------------------------------------------
# ProxyConfig._init_policy_engine
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_ProxyConfig__init_policy_engine_no_policies_noop():
pc = ProxyConfig()
try:
await pc._init_policy_engine(config={}, prisma_client=None, llm_router=None)
raised = False
except Exception:
raised = True
assert {"raised": raised, "called": True, "skipped": True} == {
"raised": False,
"called": True,
"skipped": True,
}
@pytest.mark.asyncio
async def test_ProxyConfig__init_policy_engine_none_config_noop():
pc = ProxyConfig()
# None config returns early without raising.
await pc._init_policy_engine(config=None, prisma_client=None, llm_router=None)
# Error-style: invalid policies value should raise.
with pytest.raises(AttributeError):
await pc._init_policy_engine(
config={"policies": "not-a-list"},
prisma_client=None,
llm_router=None,
)
# ---------------------------------------------------------------------------
# ProxyConfig._load_alerting_settings
# ---------------------------------------------------------------------------
def test_ProxyConfig__load_alerting_settings_noop_when_no_alerting():
pc = ProxyConfig()
try:
pc._load_alerting_settings({})
raised = False
except Exception:
raised = True
assert {"raised": raised, "called": True, "no_alerting": True} == {
"raised": False,
"called": True,
"no_alerting": True,
}
def test_ProxyConfig__load_alerting_settings_invalid_alerting_raises():
pc = ProxyConfig()
with pytest.raises(RuntimeError):
# alerting must be iterable — int triggers an error.
pc._load_alerting_settings({"alerting": 12345})
def test_ProxyConfig__load_alerting_settings_does_not_log_general_settings_dict(monkeypatch):
"""Regression for LIT-4152.
``_load_alerting_settings`` used to log ``general_settings`` verbatim in a
line labelled ``_alerting_callbacks:``, leaking ``master_key``,
``database_url``, and any other secret sitting in ``general_settings`` in
cleartext at DEBUG. The fix logs only the alerting callback list.
The regression check runs with the last-line-of-defense regex scrubber
(``SecretRedactionFilter``) DISABLED, since defense in depth is the point.
The caller must not construct the leaky string, so consumers of the log
stream that bypass the module filter (versions before it existed,
``LITELLM_DISABLE_REDACT_SECRETS=true`` operators, downstream handlers
that snapshot the record pre-filter) still do not see the secret. Uses a
dedicated handler rather than caplog because caplog is unreliable under
pytest-xdist.
"""
import logging
import litellm._logging as _logging_module
from litellm._logging import verbose_proxy_logger
monkeypatch.setattr(_logging_module, "_ENABLE_SECRET_REDACTION", False)
class LogRecordHandler(logging.Handler):
def __init__(self) -> None:
super().__init__()
self.records: list[logging.LogRecord] = []
def emit(self, record: logging.LogRecord) -> None:
self.records.append(record)
master_key_secret = "sk-lit4152-regression-master-key-abcdef1234567890"
db_url_secret = "postgresql://leak_user:leak_password_9090@leak-host.internal:5432/leak_db"
settings = {
"alerting": ["slack"],
"alerting_threshold": 300,
"master_key": master_key_secret,
"database_url": db_url_secret,
}
handler = LogRecordHandler()
handler.setLevel(logging.DEBUG)
original_level = verbose_proxy_logger.level
verbose_proxy_logger.setLevel(logging.DEBUG)
verbose_proxy_logger.addHandler(handler)
try:
try:
ProxyConfig()._load_alerting_settings(settings)
except Exception:
pass # downstream init may fail without full env; the debug log fires first
rendered = " ".join(record.getMessage() for record in handler.records)
finally:
verbose_proxy_logger.removeHandler(handler)
verbose_proxy_logger.setLevel(original_level)
assert master_key_secret not in rendered, f"master_key leaked in logs: {rendered!r}"
assert db_url_secret not in rendered, f"database_url leaked in logs: {rendered!r}"
assert "leak_password_9090" not in rendered
assert any("['slack']" in r.getMessage() for r in handler.records), (
f"expected the alerting callback list to appear in a debug record; got {[r.getMessage() for r in handler.records]!r}"
)
# ---------------------------------------------------------------------------
# ProxyConfig._warn_on_misplaced_jwt_keys
# ---------------------------------------------------------------------------
def _capture_proxy_warnings(config: dict) -> tuple[tuple[str, ...], list[str]]:
"""Run ``_warn_on_misplaced_jwt_keys`` and return (result, warning messages).
Uses a dedicated handler rather than caplog because caplog is unreliable
under pytest-xdist (see the LIT-4152 alerting test above).
"""
import logging
from litellm._logging import verbose_proxy_logger
class LogRecordHandler(logging.Handler):
def __init__(self) -> None:
super().__init__()
self.records: list[logging.LogRecord] = []
def emit(self, record: logging.LogRecord) -> None:
self.records.append(record)
handler = LogRecordHandler()
handler.setLevel(logging.WARNING)
original_level = verbose_proxy_logger.level
verbose_proxy_logger.setLevel(logging.WARNING)
verbose_proxy_logger.addHandler(handler)
try:
result = ProxyConfig()._warn_on_misplaced_jwt_keys(config=config)
finally:
verbose_proxy_logger.removeHandler(handler)
verbose_proxy_logger.setLevel(original_level)
warnings = [r.getMessage() for r in handler.records if r.levelno == logging.WARNING]
return result, warnings
def test_ProxyConfig__warn_on_misplaced_jwt_keys_warns_on_top_level_keys():
"""LIT-4584 Issue 3: JWT keys at the YAML top level are silently dropped, so
load_config must warn. Both recognized keys are reported."""
result, warnings = _capture_proxy_warnings(
{"enable_jwt_auth": True, "litellm_jwtauth": {"team_id_jwt_field": "client_id"}}
)
assert result == ("enable_jwt_auth", "litellm_jwtauth")
assert len(warnings) == 1
assert "enable_jwt_auth" in warnings[0]
assert "litellm_jwtauth" in warnings[0]
assert "general_settings" in warnings[0]
def test_ProxyConfig__warn_on_misplaced_jwt_keys_warns_even_when_also_under_general_settings():
"""A stale top-level copy is dead config even when the correct copy lives
under general_settings, so the warning must still fire on dual placement."""
result, warnings = _capture_proxy_warnings(
{
"enable_jwt_auth": True,
"general_settings": {"enable_jwt_auth": True},
}
)
assert result == ("enable_jwt_auth",)
assert len(warnings) == 1
assert "enable_jwt_auth" in warnings[0]
def test_ProxyConfig__warn_on_misplaced_jwt_keys_silent_when_correctly_placed():
"""Keys living only under general_settings are valid, so no warning fires."""
result, warnings = _capture_proxy_warnings({"general_settings": {"enable_jwt_auth": True, "litellm_jwtauth": {}}})
assert result == ()
assert warnings == []
# ---------------------------------------------------------------------------
# ProxyConfig.initialize_secret_manager
# ---------------------------------------------------------------------------
def test_ProxyConfig_initialize_secret_manager_none_noop():
pc = ProxyConfig()
try:
pc.initialize_secret_manager(key_management_system=None)
raised = False
except Exception:
raised = True
assert {"raised": raised, "called": True, "kms": None} == {
"raised": False,
"called": True,
"kms": None,
}
def test_ProxyConfig_initialize_secret_manager_invalid_kms_raises():
pc = ProxyConfig()
with pytest.raises(ValueError, match="Invalid Key Management System selected"):
pc.initialize_secret_manager(key_management_system="not-a-real-kms")
# ---------------------------------------------------------------------------
# ProxyConfig.get_model_info_with_id
# ---------------------------------------------------------------------------
def test_ProxyConfig_get_model_info_with_id_returns_router_model_info():
pc = ProxyConfig()
model = SimpleNamespace(
model_id="m-1",
model_info={"id": "m-1"},
blocked=False,
)
out = pc.get_model_info_with_id(model=model, db_model=True)
dumped = out.model_dump()
snapshot = {
"id": dumped.get("id"),
"db_model": dumped.get("db_model"),
"blocked": dumped.get("blocked"),
}
assert snapshot == {"id": "m-1", "db_model": True, "blocked": False}
PINNED_MODEL_INFO: Final = MappingProxyType(
{
"id": "pinned-row",
"key": "gpt-5.6",
"mode": "chat",
"access_groups": ["prod"],
"input_cost_per_token": 4e-06,
"output_cost_per_token": 2e-05,
"cache_read_input_token_cost_above_272k_tokens": 8e-07,
}
)
def test_ProxyConfig_get_model_info_with_id_ignores_cost_map_pricing_echoed_into_model_info():
"""LIT-8064. A pre-1.102 Admin UI save wrote the whole ``/model/info`` response back into
the row's ``model_info``, cost-map pricing included. Only that response carries ``key``, so
a stored blob with it holds a copy of the map, not a price anyone typed, and the deployment
must keep following the live cost map."""
pc = ProxyConfig()
model = SimpleNamespace(model_id="pinned-row", model_info=dict(PINNED_MODEL_INFO), blocked=False)
out = pc.get_model_info_with_id(model=model, db_model=True).model_dump(exclude_none=True)
assert out["access_groups"] == ["prod"]
assert out["mode"] == "chat"
for field in ("input_cost_per_token", "output_cost_per_token", "cache_read_input_token_cost_above_272k_tokens"):
assert field not in out, f"{field} still pins the deployment to the cost map of the day it was saved"
def test_ProxyConfig_get_model_info_with_id_keeps_pricing_typed_into_model_info():
"""A custom-priced deployment the cost map does not know never got ``key``, so its
``model_info`` pricing is the operator's own and stays."""
pc = ProxyConfig()
model = SimpleNamespace(
model_id="custom-row",
model_info={"id": "custom-row", "input_cost_per_token": 7e-06, "output_cost_per_token": 9e-06},
blocked=False,
)
out = pc.get_model_info_with_id(model=model, db_model=True).model_dump(exclude_none=True)
assert (out["input_cost_per_token"], out["output_cost_per_token"]) == (7e-06, 9e-06)
def test_ProxyConfig__add_deployment_pinned_row_follows_the_cost_map_across_reloads(monkeypatch, local_model_cost_map):
"""The customer's symptom end to end: a row pinned before 1.102 must bill at the live cost
map price on boot and again after Reload Price Data, while a price typed on
``litellm_params`` keeps overriding it."""
monkeypatch.setattr(
"litellm.proxy.proxy_server.decrypt_value_helper",
lambda value, key, return_original_value: value,
)
router = litellm.Router(model_list=[])
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router)
pinned = SimpleNamespace(
model_id="pinned-row",
model_name="gpt-5.6",
model_info=dict(PINNED_MODEL_INFO),
litellm_params={"model": "openai/gpt-5.6", "api_key": "sk-test"},
blocked=False,
)
typed = SimpleNamespace(
model_id="typed-row",
model_name="gpt-5.6-typed",
model_info={"id": "typed-row", "key": "gpt-5.6", "input_cost_per_token": 4e-06},
litellm_params={"model": "openai/gpt-5.6", "api_key": "sk-test", "input_cost_per_token": 3e-06},
blocked=False,
)
assert ProxyConfig()._add_deployment(db_models=[pinned, typed]) == 2
monkeypatch.setitem(litellm.model_cost["gpt-5.6"], "input_cost_per_token", 1e-06)
router._replay_model_cost_registrations()
assert litellm.model_cost.get("pinned-row", {}).get("input_cost_per_token") is None
assert router.get_deployment(model_id="pinned-row").model_info.input_cost_per_token is None
assert litellm.get_model_info("openai/gpt-5.6")["input_cost_per_token"] == 1e-06
assert litellm.model_cost["typed-row"]["input_cost_per_token"] == 3e-06
def test_ProxyConfig__add_deployment_ptu_row_with_a_cost_map_copy_still_bills_zero(monkeypatch, local_model_cost_map):
"""A PTU deployment bills nothing per token: the proxy writes zeros to both blobs. When such
a row also carries the echoed cost map, dropping the ``model_info`` copy must not send it
back to the per-token price, because the ``litellm_params`` zeros are the operator's."""
monkeypatch.setattr(
"litellm.proxy.proxy_server.decrypt_value_helper",
lambda value, key, return_original_value: value,
)
router = litellm.Router(model_list=[])
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router)
ptu = SimpleNamespace(
model_id="ptu-row",
model_name="gpt-5.6-ptu",
model_info={**PINNED_MODEL_INFO, "id": "ptu-row", "input_cost_per_token": 0.0, "output_cost_per_token": 0.0},
litellm_params={
"model": "openai/gpt-5.6",
"api_key": "sk-test",
"input_cost_per_token": 0.0,
"output_cost_per_token": 0.0,
},
blocked=False,
)
assert ProxyConfig()._add_deployment(db_models=[ptu]) == 1
router._replay_model_cost_registrations()
assert litellm.model_cost["ptu-row"]["input_cost_per_token"] == 0.0
assert litellm.model_cost["ptu-row"]["output_cost_per_token"] == 0.0
assert router.get_deployment(model_id="ptu-row").model_info.input_cost_per_token == 0.0
def test_ProxyConfig_get_model_info_with_id_missing_model_id_raises(monkeypatch):
monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", False)
pc = ProxyConfig()
# model with no model_id, no model_info — accessing .model_id will fail.
bad = SimpleNamespace(model_info=None)
with pytest.raises(AttributeError):
pc.get_model_info_with_id(model=bad)
# ---------------------------------------------------------------------------
# ProxyConfig._delete_deployment
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_ProxyConfig__delete_deployment_no_router_returns_none(monkeypatch):
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None)
pc = ProxyConfig()
result = await pc._delete_deployment(db_models=[])
snapshot = {"still_desired": result, "router_was": "none", "empty_db_models": True}
assert snapshot == {"still_desired": None, "router_was": "none", "empty_db_models": True}
@pytest.mark.asyncio
async def test_ProxyConfig__delete_deployment_invalid_models_raises(monkeypatch):
fake_router = MagicMock()
fake_router.get_model_ids = MagicMock(return_value=[])
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", fake_router)
pc = ProxyConfig()
with pytest.raises(AttributeError):
# Non-model objects without expected attrs trigger an error.
await pc._delete_deployment(db_models=[{"not_a_model": True}])
# ---------------------------------------------------------------------------
# ProxyConfig._add_deployment
# ---------------------------------------------------------------------------
def test_ProxyConfig__add_deployment_no_router_returns_zero(monkeypatch):
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None)
pc = ProxyConfig()
result = pc._add_deployment(db_models=[MagicMock()])
snapshot = {"added": result, "router_was": "none", "called": True}
assert snapshot == {"added": 0, "router_was": "none", "called": True}
def test_ProxyConfig__add_deployment_invalid_litellm_params_skips(monkeypatch):
fake_router = MagicMock()
fake_router.upsert_deployment = MagicMock(return_value=None)
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", fake_router)
pc = ProxyConfig()
bad = SimpleNamespace(litellm_params="not-a-dict", model_name="x", model_id="x")
# invalid params logs and continues — assert zero added (error-style branch).
assert pc._add_deployment(db_models=[bad]) == 0
def test_ProxyConfig__add_deployment_resolves_env_refs_after_db_decrypt(monkeypatch):
"""Every ``os.environ/`` value on an admin-scoped DB row resolves at
load time, regardless of the field name. Replaces the earlier
behavior where only fields in ``_DB_LITELLM_PARAM_ENV_REF_KEYS``
resolved: the whitelist has been removed so the resolver applies to
every string field."""
monkeypatch.setenv("LITELLM_DB_MODEL_API_KEY", "resolved-secret")
monkeypatch.setenv("LITELLM_MASTER_KEY", "master-secret")
monkeypatch.setattr(
"litellm.proxy.proxy_server.decrypt_value_helper",
lambda value, key, return_original_value: value,
)
fake_router = MagicMock()
fake_router.upsert_deployment = MagicMock(return_value=True)
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", fake_router)
pc = ProxyConfig()
db_model = SimpleNamespace(
model_id="model-1",
model_name="env-model",
model_info={"id": "model-1"},
litellm_params={
"model": "openai/gpt-4o-mini",
"api_key": "os.environ/LITELLM_DB_MODEL_API_KEY",
"api_base": "os.environ/LITELLM_MASTER_KEY",
},
blocked=False,
)
added = pc._add_deployment(db_models=[db_model])
deployment = fake_router.upsert_deployment.call_args.kwargs["deployment"]
assert added == 1
assert deployment.litellm_params.api_key == "resolved-secret"
assert deployment.litellm_params.api_base == "master-secret"
def test_ProxyConfig__add_deployment_resolves_team_env_refs(monkeypatch):
"""Team-scoped DB rows now resolve ``os.environ/`` refs the same way
admin rows do. The prior team-scoped short-circuit and the
field-by-field whitelist have both been removed; the write-side team
auth check in ``ModelManagementAuthChecks.can_user_make_model_call``
remains the single trust boundary. A literal (non-``os.environ/``)
value still passes through unchanged."""
monkeypatch.setenv("LITELLM_MASTER_KEY", "master-secret")
monkeypatch.setattr(
"litellm.proxy.proxy_server.decrypt_value_helper",
lambda value, key, return_original_value: value,
)
fake_router = MagicMock()
fake_router.upsert_deployment = MagicMock(return_value=True)
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", fake_router)
pc = ProxyConfig()
db_model = SimpleNamespace(
model_id="model-1",
model_name="model_name_team-1_abc",
model_info={"id": "model-1", "team_id": "team-1"},
litellm_params={
"model": "openai/gpt-4o-mini",
"api_key": "os.environ/LITELLM_MASTER_KEY",
"api_base": "https://team.example",
},
blocked=False,
)
added = pc._add_deployment(db_models=[db_model])
deployment = fake_router.upsert_deployment.call_args.kwargs["deployment"]
assert added == 1
assert deployment.litellm_params.api_key == "master-secret"
assert deployment.litellm_params.api_base == "https://team.example"
def test_ProxyConfig__resolve_db_litellm_param_skips_non_string_values(monkeypatch):
def fail_on_call(value, key, return_original_value):
raise AssertionError("decrypt_value_helper should only receive strings")
monkeypatch.setattr(
"litellm.proxy.proxy_server.decrypt_value_helper",
fail_on_call,
)
pc = ProxyConfig()
assert pc._resolve_db_litellm_param(key="tpm", value=100) == 100
def test_ProxyConfig__add_deployment_resolves_env_refs_for_aws_bedrock_auth_params(
monkeypatch,
):
"""Regression: DB-stored Bedrock/SageMaker auth params like
``aws_role_name: os.environ/BEDROCK_ASSUME_ROLE_ARN`` must resolve at
DB-load time. PR #30867 removed request-time expansion in
``BaseAWSLLM.get_credentials``; without DB-load resolution the literal
string reaches STS and fails with ``ValidationError: ... is invalid``."""
aws_env = {
"aws_session_token": ("BEDROCK_SESSION_TOKEN", "resolved-session-token"),
"aws_region_name": ("BEDROCK_REGION", "us-east-1"),
"aws_session_name": ("BEDROCK_SESSION_NAME", "resolved-session"),
"aws_profile_name": ("BEDROCK_PROFILE", "resolved-profile"),
"aws_role_name": (
"BEDROCK_ASSUME_ROLE_ARN",
"arn:aws:iam::123456789012:role/resolved",
),
"aws_web_identity_token": ("BEDROCK_WEB_IDENTITY_TOKEN", "resolved-token"),
"aws_sts_endpoint": (
"BEDROCK_STS_ENDPOINT",
"https://sts.us-east-1.amazonaws.com",
),
"aws_external_id": ("BEDROCK_EXTERNAL_ID", "resolved-external-id"),
"aws_bedrock_runtime_endpoint": (
"BEDROCK_RUNTIME_ENDPOINT",
"https://bedrock-runtime.us-east-1.amazonaws.com",
),
"aws_bedrock_project_id": ("BEDROCK_PROJECT_ID", "resolved-project-id"),
"aws_batch_role_arn": (
"BEDROCK_BATCH_ROLE_ARN",
"arn:aws:iam::123456789012:role/batch",
),
"aws_workspace_id": ("BEDROCK_WORKSPACE_ID", "resolved-workspace-id"),
}
for _, (env_name, env_value) in aws_env.items():
monkeypatch.setenv(env_name, env_value)
monkeypatch.setattr(
"litellm.proxy.proxy_server.decrypt_value_helper",
lambda value, key, return_original_value: value,
)
fake_router = MagicMock()
fake_router.upsert_deployment = MagicMock(return_value=True)
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", fake_router)
pc = ProxyConfig()
litellm_params: Dict[str, Any] = {"model": "bedrock/anthropic.claude-v2"}
for key, (env_name, _) in aws_env.items():
litellm_params[key] = f"os.environ/{env_name}"
db_model = SimpleNamespace(
model_id="model-1",
model_name="bedrock-model",
model_info={"id": "model-1"},
litellm_params=litellm_params,
blocked=False,
)
added = pc._add_deployment(db_models=[db_model])
deployment = fake_router.upsert_deployment.call_args.kwargs["deployment"]
assert added == 1
for key, (_, expected) in aws_env.items():
assert getattr(deployment.litellm_params, key) == expected, key
def test_ProxyConfig__add_deployment_resolves_env_refs_on_arbitrary_field(monkeypatch):
"""A made-up field name that was never on the removed whitelist still
resolves ``os.environ/`` refs. Pins the "no whitelist" invariant:
the resolver applies to every string field, not a curated list."""
monkeypatch.setenv("SOME_CUSTOM_ENV", "resolved-custom-value")
monkeypatch.setattr(
"litellm.proxy.proxy_server.decrypt_value_helper",
lambda value, key, return_original_value: value,
)
fake_router = MagicMock()
fake_router.upsert_deployment = MagicMock(return_value=True)
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", fake_router)
pc = ProxyConfig()
db_model = SimpleNamespace(
model_id="model-1",
model_name="custom-field-model",
model_info={"id": "model-1"},
litellm_params={
"model": "openai/gpt-4o-mini",
"some_future_field": "os.environ/SOME_CUSTOM_ENV",
},
blocked=False,
)
added = pc._add_deployment(db_models=[db_model])
deployment = fake_router.upsert_deployment.call_args.kwargs["deployment"]
assert added == 1
assert deployment.litellm_params.some_future_field == "resolved-custom-value"
@pytest.mark.parametrize(
"stored_drop_params",
["true", "os.environ/DROP_PARAMS_FLAG"],
)
def test_ProxyConfig__add_deployment_turns_stored_drop_params_string_into_bool(monkeypatch, stored_drop_params):
monkeypatch.setenv("LITELLM_SALT_KEY", "sk-1234")
monkeypatch.setenv("DROP_PARAMS_FLAG", "true")
fake_router = MagicMock()
fake_router.upsert_deployment = MagicMock(return_value=True)
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", fake_router)
pc = ProxyConfig()
db_model = SimpleNamespace(
model_id="model-1",
model_name="gpt-5-nano",
model_info={"id": "model-1"},
litellm_params={
"model": encrypt_value_helper(value="openai/gpt-5-nano"),
"drop_params": encrypt_value_helper(value=stored_drop_params),
},
blocked=False,
)
added = pc._add_deployment(db_models=[db_model])
deployment = fake_router.upsert_deployment.call_args.kwargs["deployment"]
assert added == 1
assert deployment.litellm_params.drop_params is True
def test_ProxyConfig__add_deployment_keeps_loading_rows_after_a_non_flag_drop_params(monkeypatch):
monkeypatch.setenv("LITELLM_SALT_KEY", "sk-1234")
fake_router = MagicMock()
fake_router.upsert_deployment = MagicMock(return_value=True)
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", fake_router)
pc = ProxyConfig()
def db_model(model_id, drop_params):
return SimpleNamespace(
model_id=model_id,
model_name="gpt-5-nano",
model_info={"id": model_id},
litellm_params={
"model": encrypt_value_helper(value="openai/gpt-5-nano"),
"drop_params": encrypt_value_helper(value=drop_params),
},
blocked=False,
)
added = pc._add_deployment(db_models=[db_model("bad-row", 2), db_model("good-after", "true")])
deployments = [call.kwargs["deployment"] for call in fake_router.upsert_deployment.call_args_list]
assert added == 2
assert [d.litellm_params.drop_params for d in deployments] == [None, True]
@pytest.mark.asyncio
@pytest.mark.parametrize("configured, expected", [("true", True), ("false", False)])
async def test_ProxyConfig_load_config_turns_litellm_settings_drop_params_string_into_bool(
tmp_path, monkeypatch, configured, expected
):
f = tmp_path / "c.yaml"
f.write_text(f'model_list: []\nlitellm_settings:\n drop_params: "{configured}"\n')
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False)
monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False)
monkeypatch.setattr(litellm, "drop_params", not expected)
await ProxyConfig().load_config(router=None, config_file_path=str(f))
assert litellm.drop_params is expected
@pytest.mark.asyncio
async def test_ProxyConfig_load_config_resolves_a_litellm_settings_drop_params_env_ref(tmp_path, monkeypatch):
f = tmp_path / "c.yaml"
f.write_text("model_list: []\nlitellm_settings:\n drop_params: os.environ/DROP_PARAMS_FROM_ENV\n")
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False)
monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False)
monkeypatch.setenv("DROP_PARAMS_FROM_ENV", "true")
monkeypatch.setattr(litellm, "drop_params", False)
await ProxyConfig().load_config(router=None, config_file_path=str(f))
assert litellm.drop_params is True
@pytest.mark.asyncio
async def test_ProxyConfig_load_config_warns_and_turns_off_a_non_flag_litellm_settings_drop_params(
tmp_path, monkeypatch, caplog
):
f = tmp_path / "c.yaml"
f.write_text("model_list: []\nlitellm_settings:\n drop_params: ture\n")
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False)
monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False)
monkeypatch.setattr(litellm, "drop_params", True)
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
await ProxyConfig().load_config(router=None, config_file_path=str(f))
assert litellm.drop_params is False
assert "litellm_settings.drop_params='ture' is not a flag value, treating it as off" in caplog.text
# ---------------------------------------------------------------------------
# ProxyConfig.decrypt_model_list_from_db
# ---------------------------------------------------------------------------
def test_ProxyConfig_decrypt_model_list_from_db_returns_decrypted(monkeypatch):
monkeypatch.setattr(
"litellm.proxy.proxy_server.decrypt_value_helper",
lambda value, key, return_original_value: value,
)
pc = ProxyConfig()
m = SimpleNamespace(
model_id="m-1",
model_name="gpt-4",
model_info={"id": "m-1"},
litellm_params={"api_key": "sk-x", "model": "gpt-4"},
blocked=False,
)
out = pc.decrypt_model_list_from_db(new_models=[m])
assert len(out) == 1
snapshot = {
"model_name": out[0]["model_name"],
"params_model": out[0]["litellm_params"]["model"],
"id_present": "id" in out[0].get("model_info", {}),
}
assert snapshot == {
"model_name": "gpt-4",
"params_model": "gpt-4",
"id_present": True,
}
def test_ProxyConfig_decrypt_model_list_from_db_resolves_env_refs_after_db_decrypt(
monkeypatch,
):
"""Path B (feeding /v2/model/info fallback and /model/info fallback)
resolves every ``os.environ/`` field on admin-scoped rows, mirroring
path A. Both paths now share the same universal-resolution shape."""
monkeypatch.setenv("LITELLM_DB_MODEL_API_KEY", "resolved-secret")
monkeypatch.setenv("LITELLM_MASTER_KEY", "master-secret")
monkeypatch.setattr(
"litellm.proxy.proxy_server.decrypt_value_helper",
lambda value, key, return_original_value: (
"os.environ/LITELLM_DB_MODEL_API_KEY"
if key == "api_key"
else "os.environ/LITELLM_MASTER_KEY"
if key == "api_base"
else value
),
)
pc = ProxyConfig()
m = SimpleNamespace(
model_id="model-1",
model_name="env-model",
model_info={"id": "model-1"},
litellm_params={
"api_key": "encrypted-env-ref",
"api_base": "encrypted-api-base-env-ref",
"model": "openai/gpt-4o-mini",
},
blocked=False,
)
out = pc.decrypt_model_list_from_db(new_models=[m])
assert out[0]["litellm_params"]["api_key"] == "resolved-secret"
assert out[0]["litellm_params"]["api_base"] == "master-secret"
def test_ProxyConfig_decrypt_model_list_from_db_resolves_team_env_refs_after_db_decrypt(
monkeypatch,
):
"""Team-scoped rows on path B resolve ``os.environ/`` refs just like
admin rows do. Pairs with
``test_ProxyConfig__add_deployment_resolves_team_env_refs`` on path
A — both paths now agree on the trust model."""
monkeypatch.setenv("LITELLM_MASTER_KEY", "master-secret")
monkeypatch.setattr(
"litellm.proxy.proxy_server.decrypt_value_helper",
lambda value, key, return_original_value: "os.environ/LITELLM_MASTER_KEY" if key == "api_key" else value,
)
pc = ProxyConfig()
m = SimpleNamespace(
model_id="model-1",
model_name="model_name_team-1_abc",
model_info={"id": "model-1", "team_id": "team-1"},
litellm_params={
"api_key": "encrypted-env-ref",
"api_base": "https://team.example",
"model": "openai/gpt-4o-mini",
},
blocked=False,
)
out = pc.decrypt_model_list_from_db(new_models=[m])
assert out[0]["litellm_params"]["api_key"] == "master-secret"
assert out[0]["litellm_params"]["api_base"] == "https://team.example"
def test_ProxyConfig_decrypt_model_list_from_db_invalid_params_skips():
pc = ProxyConfig()
bad = SimpleNamespace(model_id="m-1", model_name="x", model_info={}, litellm_params="not-a-dict")
out = pc.decrypt_model_list_from_db(new_models=[bad])
# Invalid entries skipped — empty list returned.
assert out == []
# ---------------------------------------------------------------------------
# ProxyConfig._update_llm_router
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_ProxyConfig__update_llm_router_no_models_smoke(monkeypatch):
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None)
monkeypatch.setattr("litellm.proxy.proxy_server.master_key", "sk-master")
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
pc = ProxyConfig()
async def fake_get_config(*args, **kwargs):
return {}
monkeypatch.setattr(pc, "get_config", fake_get_config)
monkeypatch.setattr(
"litellm.proxy.proxy_server.proxy_config",
pc,
)
try:
await pc._update_llm_router(new_models=[], proxy_logging_obj=MagicMock())
raised = False
except Exception:
raised = True
snapshot = {"raised": raised, "called": True, "models": "empty"}
assert snapshot == {"raised": False, "called": True, "models": "empty"}
# ---------------------------------------------------------------------------
# ProxyConfig._add_callback_from_db_to_in_memory_litellm_callbacks
# ---------------------------------------------------------------------------
def test_ProxyConfig__add_callback_from_db_to_in_memory_litellm_callbacks_adds(
monkeypatch,
):
monkeypatch.setattr(litellm, "callbacks", [], raising=False)
pc = ProxyConfig()
pc._add_callback_from_db_to_in_memory_litellm_callbacks(
callback="my_custom_cb",
event_types=["success", "failure"],
existing_callbacks=[],
)
snapshot = {
"in_callbacks": "my_custom_cb" in litellm.callbacks,
"count": len(litellm.callbacks),
"method_called": True,
}
assert snapshot == {"in_callbacks": True, "count": 1, "method_called": True}
def test_ProxyConfig__add_callback_from_db_to_in_memory_litellm_callbacks_invalid_event_raises(
monkeypatch,
):
monkeypatch.setattr(litellm, "callbacks", [], raising=False)
pc = ProxyConfig()
# For a "known" callback, event_types is iterated — non-iterable raises TypeError.
with pytest.raises(TypeError):
pc._add_callback_from_db_to_in_memory_litellm_callbacks(
callback="lago", # in _known_custom_logger_compatible_callbacks
event_types=12345, # type: ignore[arg-type]
existing_callbacks=[],
)
# ---------------------------------------------------------------------------
# ProxyConfig._add_callbacks_from_db_config
# ---------------------------------------------------------------------------
def test_ProxyConfig__add_callbacks_from_db_config_processes_lists(monkeypatch):
monkeypatch.setattr(litellm, "callbacks", [], raising=False)
monkeypatch.setattr(litellm, "success_callback", [], raising=False)
monkeypatch.setattr(litellm, "failure_callback", [], raising=False)
pc = ProxyConfig()
cfg = {
"litellm_settings": {
"callbacks": ["cb_a"],
"success_callback": ["s_a"],
"failure_callback": ["f_a"],
}
}
pc._add_callbacks_from_db_config(cfg)
snapshot = {
"cb_added": "cb_a" in litellm.callbacks,
"success_added": "s_a" in litellm.success_callback,
"failure_added": "f_a" in litellm.failure_callback,
}
assert snapshot == {
"cb_added": True,
"success_added": True,
"failure_added": True,
}
def test_ProxyConfig__add_callbacks_from_db_config_bad_config_raises():
pc = ProxyConfig()
with pytest.raises(AttributeError):
# Non-dict input — .get will fail.
pc._add_callbacks_from_db_config(None) # type: ignore[arg-type]
# ---------------------------------------------------------------------------
# ProxyConfig._encrypt_env_variables
# ---------------------------------------------------------------------------
def test_ProxyConfig__encrypt_env_variables_returns_dict(monkeypatch):
monkeypatch.setattr(
"litellm.proxy.proxy_server.encrypt_value_helper",
lambda value, new_encryption_key=None: f"ENC[{value}]",
)
pc = ProxyConfig()
out = pc._encrypt_env_variables({"A": "1", "B": "2", "C": "3"})
assert out == {"A": "ENC[1]", "B": "ENC[2]", "C": "ENC[3]"}
def test_ProxyConfig__encrypt_env_variables_invalid_raises():
pc = ProxyConfig()
with pytest.raises(AttributeError):
# Non-dict input — .items() fails.
pc._encrypt_env_variables(None) # type: ignore[arg-type]
# ---------------------------------------------------------------------------
# ProxyConfig._decrypt_and_set_db_env_variables
# ---------------------------------------------------------------------------
def test_ProxyConfig__decrypt_and_set_db_env_variables_sets_env(monkeypatch):
monkeypatch.setattr(
"litellm.proxy.proxy_server.decrypt_value_helper",
lambda value, key, return_original_value=False: value + "-dec",
)
monkeypatch.delenv("KEY_X", raising=False)
monkeypatch.delenv("KEY_Y", raising=False)
pc = ProxyConfig()
out = pc._decrypt_and_set_db_env_variables({"KEY_X": "x", "KEY_Y": "y"})
snapshot = {
"KEY_X_env": os.environ.get("KEY_X"),
"KEY_Y_env": os.environ.get("KEY_Y"),
"returned_keys": sorted(out.keys()),
}
assert snapshot == {
"KEY_X_env": "x-dec",
"KEY_Y_env": "y-dec",
"returned_keys": ["KEY_X", "KEY_Y"],
}
def test_ProxyConfig__decrypt_and_set_db_env_variables_invalid_dict_raises():
pc = ProxyConfig()
with pytest.raises(AttributeError):
pc._decrypt_and_set_db_env_variables("not-a-dict") # type: ignore[arg-type]
# ---------------------------------------------------------------------------
# ProxyConfig._decrypt_db_variables
# ---------------------------------------------------------------------------
def test_ProxyConfig__decrypt_db_variables_returns_decrypted(monkeypatch):
monkeypatch.setattr(
"litellm.proxy.proxy_server.decrypt_value_helper",
lambda value, key, return_original_value: f"D({value})",
)
pc = ProxyConfig()
out = pc._decrypt_db_variables({"a": "1", "b": "2", "c": "3"})
assert out == {"a": "D(1)", "b": "D(2)", "c": "D(3)"}
def test_ProxyConfig__decrypt_db_variables_invalid_raises():
pc = ProxyConfig()
with pytest.raises(AttributeError):
pc._decrypt_db_variables(None) # type: ignore[arg-type]
# ---------------------------------------------------------------------------
# ProxyConfig._encrypt_env_variables_for_db
# ---------------------------------------------------------------------------
def test_ProxyConfig__encrypt_env_variables_for_db_idempotent(monkeypatch):
monkeypatch.setattr(
"litellm.proxy.proxy_server.decrypt_value_helper",
lambda value, key, return_original_value: value,
)
monkeypatch.setattr(
"litellm.proxy.proxy_server.encrypt_value_helper",
lambda value, new_encryption_key=None: f"ENC[{value}]",
)
pc = ProxyConfig()
out = pc._encrypt_env_variables_for_db({"A": "1", "B": "2", "C": "3"})
assert out == {"A": "ENC[1]", "B": "ENC[2]", "C": "ENC[3]"}
def test_ProxyConfig__encrypt_env_variables_for_db_invalid_raises():
pc = ProxyConfig()
with pytest.raises(AttributeError):
pc._encrypt_env_variables_for_db(None) # type: ignore[arg-type]
# ---------------------------------------------------------------------------
# ProxyConfig._parse_router_settings_value
# ---------------------------------------------------------------------------
def test_ProxyConfig__parse_router_settings_value_handles_inputs():
result = {
"dict": ProxyConfig._parse_router_settings_value({"a": 1}),
"yaml_string": ProxyConfig._parse_router_settings_value("a: 1\nb: 2"),
"none": ProxyConfig._parse_router_settings_value(None),
}
assert result == {
"dict": {"a": 1},
"yaml_string": {"a": 1, "b": 2},
"none": None,
}
def test_ProxyConfig__parse_router_settings_value_invalid_returns_none():
# Non-dict, non-parseable scalar -> None.
assert ProxyConfig._parse_router_settings_value(12345) is None
# Empty dict -> None (not truthy).
assert ProxyConfig._parse_router_settings_value({}) is None
# ---------------------------------------------------------------------------
# ProxyConfig._get_hierarchical_router_settings
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_ProxyConfig__get_hierarchical_router_settings_key_wins():
pc = ProxyConfig()
fake_key = SimpleNamespace(
router_settings={"timeout": 30, "retries": 2, "model": "gpt-4"},
team_id=None,
)
out = await pc._get_hierarchical_router_settings(
user_api_key_dict=fake_key,
prisma_client=None,
proxy_logging_obj=None,
)
assert out == {"timeout": 30, "retries": 2, "model": "gpt-4"}
@pytest.mark.asyncio
async def test_ProxyConfig__get_hierarchical_router_settings_missing_returns_none():
pc = ProxyConfig()
fake_key = SimpleNamespace(router_settings=None, team_id=None)
out = await pc._get_hierarchical_router_settings(
user_api_key_dict=fake_key,
prisma_client=None,
proxy_logging_obj=None,
)
assert out is None
@pytest.mark.asyncio
async def test_ProxyConfig__get_hierarchical_router_settings_falls_back_to_team(monkeypatch):
"""A key with no router_settings inherits the team's, so a team-level
model_group_alias reaches the request path at all."""
pc = ProxyConfig()
fake_key = SimpleNamespace(router_settings=None, team_id="team-1")
team_settings = {"model_group_alias": {"group-a": "group-b"}}
monkeypatch.setattr(
"litellm.proxy.proxy_server.get_team_object",
AsyncMock(return_value=SimpleNamespace(router_settings=team_settings)),
)
out = await pc._get_hierarchical_router_settings(
user_api_key_dict=fake_key,
prisma_client=None,
proxy_logging_obj=None,
)
assert out == team_settings
@pytest.mark.asyncio
async def test_ProxyConfig__get_hierarchical_router_settings_key_shadows_team_entirely(monkeypatch):
"""Resolution returns whichever object it finds first, it does not merge
per field, so a key that sets any router setting hides every team setting
including an alias the key itself never set."""
pc = ProxyConfig()
fake_key = SimpleNamespace(router_settings={"num_retries": 3}, team_id="team-1")
team_lookup = AsyncMock(return_value=SimpleNamespace(router_settings={"model_group_alias": {"group-a": "group-b"}}))
monkeypatch.setattr("litellm.proxy.proxy_server.get_team_object", team_lookup)
out = await pc._get_hierarchical_router_settings(
user_api_key_dict=fake_key,
prisma_client=None,
proxy_logging_obj=None,
)
assert out == {"num_retries": 3}
assert "model_group_alias" not in out
team_lookup.assert_not_called()
# ---------------------------------------------------------------------------
# ProxyConfig._add_router_settings_from_db_config
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_ProxyConfig__add_router_settings_from_db_config_updates_router():
pc = ProxyConfig()
fake_router = MagicMock()
fake_router.update_settings = MagicMock()
fake_prisma = MagicMock()
fake_prisma.db.litellm_config.find_first = AsyncMock(
return_value=SimpleNamespace(param_value={"timeout": 30, "retries": 2, "fallbacks": []})
)
pc.router_settings.load_yaml({"timeout": 10})
await pc._add_router_settings_from_db_config(
llm_router=fake_router,
prisma_client=fake_prisma,
)
snapshot = {
"called": fake_router.update_settings.called,
"call_count": fake_router.update_settings.call_count,
"kwargs_keys": sorted(list(fake_router.update_settings.call_args.kwargs.keys())),
}
assert snapshot == {
"called": True,
"call_count": 1,
"kwargs_keys": ["fallbacks", "retries", "timeout"],
}
@pytest.mark.asyncio
async def test_ProxyConfig__add_router_settings_from_db_config_none_router_noop():
pc = ProxyConfig()
# No router and no prisma — should silently return.
await pc._add_router_settings_from_db_config(llm_router=None, prisma_client=None)
# Error-style: bad call signature raises.
with pytest.raises(TypeError):
await pc._add_router_settings_from_db_config() # type: ignore[call-arg]
# ---------------------------------------------------------------------------
# ProxyConfig.add_deployment
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_ProxyConfig_add_deployment_applies_db_router_settings(monkeypatch):
from litellm.proxy import proxy_server
pc = ProxyConfig()
fake_router = MagicMock()
fake_router.get_model_list = MagicMock(return_value=[])
fake_prisma = MagicMock()
fake_prisma.db.litellm_config.find_first = AsyncMock(
return_value=SimpleNamespace(param_value={"routing_strategy": "latency-based-routing"})
)
async def fake_get_config(*args, **kwargs):
return {}
monkeypatch.setattr(pc, "get_config", fake_get_config)
monkeypatch.setattr(pc, "_get_models_from_db", AsyncMock(return_value=[]))
monkeypatch.setattr(pc, "_init_non_llm_objects_in_db", AsyncMock())
monkeypatch.setattr(proxy_server, "prefetch_config_params", AsyncMock())
monkeypatch.setattr(proxy_server, "get_config_param", AsyncMock(return_value=None))
monkeypatch.setattr(proxy_server, "llm_router", fake_router)
monkeypatch.setattr(proxy_server, "master_key", "sk-master")
monkeypatch.setattr(proxy_server, "prisma_client", fake_prisma)
monkeypatch.setattr(proxy_server, "general_settings", {})
monkeypatch.setattr(proxy_server, "proxy_config", pc)
await pc.add_deployment(prisma_client=fake_prisma, proxy_logging_obj=MagicMock())
fake_router.update_settings.assert_called_once_with(routing_strategy="latency-based-routing")
def _stub_add_deployment_collaborators(
monkeypatch: pytest.MonkeyPatch, pc: ProxyConfig, fake_prisma: MagicMock
) -> None:
from litellm.proxy import proxy_server
fake_router = MagicMock()
fake_router.get_model_list = MagicMock(return_value=[])
async def fake_get_config(*args: object, **kwargs: object) -> dict[str, object]:
return {}
monkeypatch.setattr(litellm, "credential_list", [])
monkeypatch.setattr(pc, "get_config", fake_get_config)
monkeypatch.setattr(pc, "_init_non_llm_objects_in_db", AsyncMock())
monkeypatch.setattr(proxy_server, "prefetch_config_params", AsyncMock())
monkeypatch.setattr(proxy_server, "get_config_param", AsyncMock(return_value=None))
monkeypatch.setattr(proxy_server, "llm_router", fake_router)
monkeypatch.setattr(proxy_server, "master_key", "sk-master")
monkeypatch.setattr(proxy_server, "prisma_client", fake_prisma)
monkeypatch.setattr(proxy_server, "proxy_config", pc)
monkeypatch.delenv("LITELLM_SALT_KEY", raising=False)
def _encrypted_credential_row(credential_name: str, api_key: str) -> dict[str, object]:
from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper
return {
"credential_name": credential_name,
"credential_values": {"api_key": encrypt_value_helper(api_key, new_encryption_key="sk-master")},
"credential_info": {"custom_llm_provider": "openai"},
}
def _fake_prisma_with_encrypted_credential(credential_name: str, api_key: str) -> MagicMock:
fake_prisma = MagicMock()
fake_prisma.db.litellm_credentialstable.find_many = AsyncMock(
return_value=[_encrypted_credential_row(credential_name, api_key)]
)
return fake_prisma
@pytest.mark.asyncio
async def test_ProxyConfig_add_deployment_loads_db_credentials_before_reconciling_models(
monkeypatch: pytest.MonkeyPatch,
) -> None:
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
from litellm.proxy import proxy_server
from litellm.utils import load_credentials_from_list
pc = ProxyConfig()
fake_prisma = MagicMock()
fake_prisma.db.litellm_credentialstable.find_many = AsyncMock(return_value=[])
_stub_add_deployment_collaborators(monkeypatch, pc, fake_prisma)
monkeypatch.setattr(proxy_server, "general_settings", {})
installed = MagicMock()
async def read_models_while_a_credential_lands(prisma_client: object) -> list[MagicMock]:
fake_prisma.db.litellm_credentialstable.find_many.return_value = [
_encrypted_credential_row("openai-cred", "sk-from-db")
]
return [MagicMock()]
async def install_models(new_models: object, proxy_logging_obj: object) -> None:
installed(credential=CredentialAccessor.get_credential_values("openai-cred"))
monkeypatch.setattr(pc, "_get_models_from_db", read_models_while_a_credential_lands)
monkeypatch.setattr(pc, "_update_llm_router", install_models)
await pc.add_deployment(prisma_client=fake_prisma, proxy_logging_obj=MagicMock())
installed.assert_called_once_with(credential={"api_key": "sk-from-db"})
assert CredentialAccessor.get_credential_values("openai-cred") == {"api_key": "sk-from-db"}
request_kwargs = {"litellm_credential_name": "openai-cred"}
load_credentials_from_list(request_kwargs)
assert request_kwargs == {"litellm_credential_name": "openai-cred", "api_key": "sk-from-db"}
@pytest.mark.asyncio
async def test_ProxyConfig_add_deployment_loads_db_credentials_even_when_models_are_not_db_objects(
monkeypatch: pytest.MonkeyPatch,
) -> None:
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
from litellm.proxy import proxy_server
pc = ProxyConfig()
fake_prisma = _fake_prisma_with_encrypted_credential("openai-cred", "sk-from-db")
_stub_add_deployment_collaborators(monkeypatch, pc, fake_prisma)
monkeypatch.setattr(proxy_server, "general_settings", {"supported_db_objects": ["mcp"]})
models_fetch = AsyncMock(return_value=[])
monkeypatch.setattr(pc, "_get_models_from_db", models_fetch)
await pc.add_deployment(prisma_client=fake_prisma, proxy_logging_obj=MagicMock())
models_fetch.assert_not_awaited()
assert CredentialAccessor.get_credential_values("openai-cred") == {"api_key": "sk-from-db"}
@pytest.mark.asyncio
async def test_ProxyConfig_get_credentials_reads_from_writer_not_replica(monkeypatch: pytest.MonkeyPatch) -> None:
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
from litellm.proxy.db.prisma_client import PrismaWrapper
from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper
pc = ProxyConfig()
writer_inner = MagicMock(name="writer_prisma")
reader_inner = MagicMock(name="reader_prisma")
writer_inner.litellm_credentialstable.find_many = AsyncMock(
return_value=[_encrypted_credential_row("openai-cred", "sk-from-writer")]
)
reader_inner.litellm_credentialstable.find_many = AsyncMock(return_value=[])
fake_prisma = MagicMock()
fake_prisma.db = RoutingPrismaWrapper(
writer=PrismaWrapper(original_prisma=writer_inner, iam_token_db_auth=False),
reader=PrismaWrapper(original_prisma=reader_inner, iam_token_db_auth=False),
)
_stub_add_deployment_collaborators(monkeypatch, pc, fake_prisma)
await pc.get_credentials(prisma_client=fake_prisma)
assert CredentialAccessor.get_credential_values("openai-cred") == {"api_key": "sk-from-writer"}
reader_inner.litellm_credentialstable.find_many.assert_not_awaited()
# ---------------------------------------------------------------------------
# ProxyConfig._reschedule_spend_log_cleanup_job
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_ProxyConfig__reschedule_spend_log_cleanup_job_no_scheduler(monkeypatch):
monkeypatch.setattr("litellm.proxy.proxy_server.scheduler", None)
pc = ProxyConfig()
try:
await pc._reschedule_spend_log_cleanup_job()
raised = False
except Exception:
raised = True
snapshot = {"raised": raised, "called": True, "scheduler_was": "none"}
assert snapshot == {"raised": False, "called": True, "scheduler_was": "none"}
@pytest.mark.asyncio
async def test_ProxyConfig__reschedule_spend_log_cleanup_job_invalid_cron(monkeypatch):
fake_scheduler = MagicMock()
fake_scheduler.remove_job = MagicMock()
fake_scheduler.add_job = MagicMock()
monkeypatch.setattr("litellm.proxy.proxy_server.scheduler", fake_scheduler)
monkeypatch.setattr(
"litellm.proxy.proxy_server.general_settings",
{
"maximum_spend_logs_retention_period": "1d",
"maximum_spend_logs_cleanup_cron": "INVALID CRON STRING",
},
)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
pc = ProxyConfig()
# Invalid cron is caught and logged — does not raise outward.
await pc._reschedule_spend_log_cleanup_job()
# But add_job should not have been called for the invalid cron path.
assert fake_scheduler.add_job.call_count == 0
@pytest.mark.asyncio
async def test_ProxyConfig__reschedule_spend_log_cleanup_job_health_check_retention(monkeypatch):
fake_scheduler = MagicMock()
monkeypatch.setattr("litellm.proxy.proxy_server.scheduler", fake_scheduler)
monkeypatch.setattr(
"litellm.proxy.proxy_server.general_settings",
{"maximum_health_check_retention_period": "30d"},
)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
pc = ProxyConfig()
await pc._reschedule_spend_log_cleanup_job()
assert fake_scheduler.add_job.call_count == 1
assert fake_scheduler.add_job.call_args.kwargs["id"] == "spend_log_cleanup_job"
@pytest.mark.asyncio
async def test_ProxyConfig__update_general_settings_updates_health_check_retention(monkeypatch):
settings = {}
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", settings)
monkeypatch.setattr("litellm.proxy.proxy_server.scheduler", MagicMock(**{"get_job.return_value": None}))
pc = ProxyConfig()
reschedule = AsyncMock()
monkeypatch.setattr(pc, "_reschedule_spend_log_cleanup_job", reschedule)
await pc._update_general_settings({"maximum_health_check_retention_period": "30d"})
from litellm.proxy import proxy_server
assert proxy_server.general_settings["maximum_health_check_retention_period"] == "30d"
reschedule.assert_awaited_once()
def _paused_scheduler(monkeypatch):
from apscheduler.schedulers.asyncio import AsyncIOScheduler
real_scheduler = AsyncIOScheduler()
real_scheduler.start(paused=True)
monkeypatch.setattr("litellm.proxy.proxy_server.scheduler", real_scheduler)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
return real_scheduler
def _scheduler_whose_first_add_job_raises(monkeypatch):
from apscheduler.schedulers.asyncio import AsyncIOScheduler
class FirstAddJobRaises(AsyncIOScheduler):
raised = False
def add_job(self, *args, **kwargs):
if not self.raised:
self.raised = True
raise RuntimeError("scheduler busy")
return super().add_job(*args, **kwargs)
real_scheduler = FirstAddJobRaises()
real_scheduler.start(paused=True)
monkeypatch.setattr("litellm.proxy.proxy_server.scheduler", real_scheduler)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
return real_scheduler
@pytest.mark.asyncio
async def test_ProxyConfig__reschedule_spend_log_cleanup_job_daily_tag_spend_retention(monkeypatch):
real_scheduler = _paused_scheduler(monkeypatch)
monkeypatch.setattr(
"litellm.proxy.proxy_server.general_settings",
{"maximum_daily_tag_spend_retention_period": "90d"},
)
pc = ProxyConfig()
try:
await pc._reschedule_spend_log_cleanup_job()
job = real_scheduler.get_job("spend_log_cleanup_job")
assert job is not None, "daily tag spend retention alone did not schedule the cleanup job"
assert job.func.__name__ == "cleanup_old_spend_logs"
finally:
real_scheduler.shutdown(wait=False)
@pytest.mark.asyncio
async def test_ProxyConfig__update_general_settings_updates_daily_tag_spend_retention(monkeypatch):
real_scheduler = _paused_scheduler(monkeypatch)
pc = ProxyConfig()
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", pc.settings)
try:
await pc._update_general_settings({"maximum_daily_tag_spend_retention_period": "90d"})
from litellm.proxy import proxy_server
assert proxy_server.general_settings["maximum_daily_tag_spend_retention_period"] == "90d"
assert real_scheduler.get_job("spend_log_cleanup_job") is not None, "runtime retention did not schedule cleanup"
finally:
real_scheduler.shutdown(wait=False)
@pytest.mark.asyncio
async def test_ProxyConfig__update_general_settings_schedules_cleanup_when_db_row_was_already_applied(monkeypatch):
"""A config reload applies the db row to the store before the side effects run, so the
before/after snapshot is equal; the job must still be scheduled when none is running."""
real_scheduler = _paused_scheduler(monkeypatch)
pc = ProxyConfig()
pc.settings.apply_db_row("general_settings", {"maximum_daily_tag_spend_retention_period": "90d"})
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", pc.settings)
try:
await pc._update_general_settings({"maximum_daily_tag_spend_retention_period": "90d"})
assert real_scheduler.get_job("spend_log_cleanup_job") is not None, "DB-only retention never scheduled cleanup"
finally:
real_scheduler.shutdown(wait=False)
@pytest.mark.asyncio
async def test_ProxyConfig__update_general_settings_retries_a_failed_schedule_once_per_settings_value(
monkeypatch, caplog
):
"""An unparseable cron leaves no job behind; reloads must not retry it every tick, only when the
cron or a retention value changes."""
real_scheduler = _paused_scheduler(monkeypatch)
pc = ProxyConfig()
bad_cron = {"maximum_daily_tag_spend_retention_period": "90d", "maximum_spend_logs_cleanup_cron": "not a cron"}
pc.settings.apply_db_row("general_settings", bad_cron)
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", pc.settings)
try:
with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"):
for _ in range(3):
await pc._update_general_settings(bad_cron)
assert real_scheduler.get_job("spend_log_cleanup_job") is None
cron_errors = [r for r in caplog.records if "maximum_spend_logs_cleanup_cron" in r.getMessage()]
assert len(cron_errors) == 1, f"invalid cron was retried on every reload: {len(cron_errors)} error lines"
await pc._update_general_settings({**bad_cron, "maximum_spend_logs_cleanup_cron": "* * * * *"})
job = real_scheduler.get_job("spend_log_cleanup_job")
assert job is not None, "a corrected cron did not schedule cleanup"
assert "minute='*'" in str(job.trigger)
finally:
real_scheduler.shutdown(wait=False)
@pytest.mark.asyncio
async def test_ProxyConfig__update_general_settings_retries_a_schedule_that_raised(monkeypatch):
"""A transient add_job failure must not be remembered as a completed attempt; the next
reload with the same settings tries again."""
real_scheduler = _scheduler_whose_first_add_job_raises(monkeypatch)
pc = ProxyConfig()
retention = {"maximum_daily_tag_spend_retention_period": "90d"}
pc.settings.apply_db_row("general_settings", retention)
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", pc.settings)
try:
await pc._update_general_settings(retention)
assert real_scheduler.get_job("spend_log_cleanup_job") is None
await pc._update_general_settings(retention)
assert real_scheduler.get_job("spend_log_cleanup_job") is not None, "raised add_job was not retried"
finally:
real_scheduler.shutdown(wait=False)
@pytest.mark.asyncio
async def test_ProxyConfig__update_general_settings_retries_a_failed_replacement_of_the_live_job(monkeypatch):
"""A cron change whose add_job raised keeps the old job running, so the next reload with the
same settings must try the replacement again instead of leaving the new cron unapplied."""
real_scheduler = _scheduler_whose_first_add_job_raises(monkeypatch)
pc = ProxyConfig()
pc.settings.load_yaml({"maximum_daily_tag_spend_retention_period": "90d"})
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", pc.settings)
real_scheduler.raised = True
await pc._reschedule_spend_log_cleanup_job()
real_scheduler.raised = False
try:
new_cron = {"maximum_spend_logs_cleanup_cron": "0 3 * * *"}
await pc._update_general_settings(new_cron)
assert "hour='3'" not in str(real_scheduler.get_job("spend_log_cleanup_job").trigger), "old job was lost"
await pc._update_general_settings(new_cron)
assert "hour='3'" in str(real_scheduler.get_job("spend_log_cleanup_job").trigger), (
"failed replacement was not retried on the next sync"
)
finally:
real_scheduler.shutdown(wait=False)
@pytest.mark.asyncio
async def test_ProxyConfig__update_general_settings_leaves_a_changed_db_schedule_to_startup_while_scheduler_is_stopped(
monkeypatch,
):
"""The first DB sync runs before the scheduler starts and usually differs from the yaml; it
must still leave registration to the startup block instead of adding a job it will replace."""
from apscheduler.schedulers.asyncio import AsyncIOScheduler
real_scheduler = AsyncIOScheduler()
monkeypatch.setattr("litellm.proxy.proxy_server.scheduler", real_scheduler)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
pc = ProxyConfig()
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", pc.settings)
await pc._update_general_settings({"maximum_daily_tag_spend_retention_period": "90d"})
assert real_scheduler.get_jobs() == [], "DB sync registered the cleanup job before the scheduler started"
@pytest.mark.asyncio
async def test_ProxyConfig__update_general_settings_leaves_first_registration_to_startup_while_scheduler_is_stopped(
monkeypatch,
):
"""The DB sync that runs before the scheduler starts must not register the cleanup job; the
startup block does, once, so the cross-replica stagger it applies to pending jobs survives."""
from apscheduler.schedulers.asyncio import AsyncIOScheduler
real_scheduler = AsyncIOScheduler()
monkeypatch.setattr("litellm.proxy.proxy_server.scheduler", real_scheduler)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
pc = ProxyConfig()
pc.settings.load_yaml({"maximum_daily_tag_spend_retention_period": "90d"})
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", pc.settings)
await pc._update_general_settings({"unrelated_key": "value"})
assert real_scheduler.get_jobs() == [], "DB sync registered the cleanup job before the scheduler started"
@pytest.mark.asyncio
async def test_ProxyConfig__update_general_settings_runtime_interval_job_carries_the_stagger_offset(monkeypatch):
"""Once the scheduler is running the sync owns registration and the job it adds is staggered."""
from apscheduler.schedulers.asyncio import AsyncIOScheduler
from litellm.proxy.common_utils.scheduled_job_stagger import _OffsetTrigger
real_scheduler = AsyncIOScheduler()
real_scheduler.start(paused=True)
monkeypatch.setattr("litellm.proxy.proxy_server.scheduler", real_scheduler)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
pc = ProxyConfig()
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", pc.settings)
try:
await pc._update_general_settings({"maximum_daily_tag_spend_retention_period": "90d"})
jobs = real_scheduler.get_jobs()
assert [job.id for job in jobs] == ["spend_log_cleanup_job"]
assert isinstance(jobs[0].trigger, _OffsetTrigger), repr(jobs[0].trigger)
finally:
real_scheduler.shutdown(wait=False)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"bad_schedule",
[
{"maximum_spend_logs_cleanup_cron": "not a cron"},
{"maximum_spend_logs_cleanup_cron": "0 0 * * * *"},
{"maximum_spend_logs_retention_interval": "soon"},
{"maximum_spend_logs_retention_interval": 86400},
],
)
async def test_ProxyConfig__update_general_settings_keeps_the_live_cleanup_job_when_the_new_schedule_is_invalid(
monkeypatch, bad_schedule
):
"""A schedule edit that does not parse must leave the old cleanup job running and must not
stop the rest of the general settings sync."""
from apscheduler.schedulers.asyncio import AsyncIOScheduler
real_scheduler = AsyncIOScheduler()
real_scheduler.start(paused=True)
monkeypatch.setattr("litellm.proxy.proxy_server.scheduler", real_scheduler)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
ssrf_sync = MagicMock()
monkeypatch.setattr("litellm.proxy.proxy_server._apply_ssrf_general_settings", ssrf_sync)
pc = ProxyConfig()
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", pc.settings)
try:
await pc._update_general_settings({"maximum_daily_tag_spend_retention_period": "90d"})
old_trigger = real_scheduler.get_job("spend_log_cleanup_job").trigger
ssrf_sync.reset_mock()
for _ in range(2):
await pc._update_general_settings({"maximum_daily_tag_spend_retention_period": "90d", **bad_schedule})
live_job = real_scheduler.get_job("spend_log_cleanup_job")
assert live_job is not None, "invalid schedule removed the cleanup job"
assert live_job.trigger is old_trigger
assert ssrf_sync.call_count == 2, "schedule error blocked the rest of the settings sync"
finally:
real_scheduler.shutdown(wait=False)
@pytest.mark.asyncio
async def test_ProxyConfig__update_general_settings_logs_an_overflowing_interval_once(monkeypatch, caplog):
"""An interval that parses but overflows the trigger must keep the live job and log one
error, not a traceback on every sync."""
from apscheduler.schedulers.asyncio import AsyncIOScheduler
real_scheduler = AsyncIOScheduler()
real_scheduler.start(paused=True)
monkeypatch.setattr("litellm.proxy.proxy_server.scheduler", real_scheduler)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
pc = ProxyConfig()
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", pc.settings)
try:
await pc._update_general_settings({"maximum_daily_tag_spend_retention_period": "90d"})
old_trigger = real_scheduler.get_job("spend_log_cleanup_job").trigger
overflowing = {
"maximum_daily_tag_spend_retention_period": "90d",
"maximum_spend_logs_retention_interval": "99999999999d",
}
with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"):
for _ in range(5):
await pc._update_general_settings(overflowing)
errors = [record for record in caplog.records if record.levelno >= logging.ERROR]
assert len(errors) == 1, [record.getMessage() for record in errors]
assert real_scheduler.get_job("spend_log_cleanup_job").trigger is old_trigger
finally:
real_scheduler.shutdown(wait=False)
@pytest.mark.asyncio
async def test_ProxyConfig__update_general_settings_reschedules_when_only_the_cron_changes(monkeypatch):
real_scheduler = _paused_scheduler(monkeypatch)
pc = ProxyConfig()
pc.settings.load_yaml({"maximum_daily_tag_spend_retention_period": "90d"})
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", pc.settings)
await pc._reschedule_spend_log_cleanup_job()
try:
interval_job = real_scheduler.get_job("spend_log_cleanup_job")
assert interval_job is not None and "hour='3'" not in str(interval_job.trigger)
await pc._update_general_settings({"maximum_spend_logs_cleanup_cron": "0 3 * * *"})
cron_job = real_scheduler.get_job("spend_log_cleanup_job")
assert "hour='3'" in str(cron_job.trigger), "cron-only change did not reschedule"
await pc._update_general_settings({"maximum_spend_logs_cleanup_cron": "0 3 * * *"})
assert real_scheduler.get_job("spend_log_cleanup_job") is cron_job, "unchanged cron replaced the job"
finally:
real_scheduler.shutdown(wait=False)
@pytest.mark.asyncio
async def test_ProxyConfig__update_general_settings_reschedules_a_cron_edit_the_reload_path_already_applied(
monkeypatch,
):
"""The periodic reload applies the DB row through _update_config_from_db before
_update_general_settings snapshots the previous schedule, so a cron edited in the DB must
still replace the live job's trigger."""
from apscheduler.schedulers.asyncio import AsyncIOScheduler
real_scheduler = AsyncIOScheduler()
real_scheduler.start(paused=True)
monkeypatch.setattr("litellm.proxy.proxy_server.scheduler", real_scheduler)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
pc = ProxyConfig()
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", pc.settings)
try:
first_row = {"maximum_daily_tag_spend_retention_period": "90d", "maximum_spend_logs_cleanup_cron": "0 3 * * *"}
pc.settings.apply_db_row("general_settings", first_row)
await pc._update_general_settings(first_row)
assert "hour='3'" in str(real_scheduler.get_job("spend_log_cleanup_job").trigger)
edited_row = {**first_row, "maximum_spend_logs_cleanup_cron": "0 5 * * *"}
pc.settings.apply_db_row("general_settings", edited_row)
await pc._update_general_settings(edited_row)
assert "hour='5'" in str(real_scheduler.get_job("spend_log_cleanup_job").trigger), "DB cron edit was ignored"
pc.settings.apply_db_row("general_settings", edited_row)
await pc._update_general_settings(edited_row)
assert "hour='5'" in str(real_scheduler.get_job("spend_log_cleanup_job").trigger)
finally:
real_scheduler.shutdown(wait=False)
# ---------------------------------------------------------------------------
# ProxyConfig._update_general_settings
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_ProxyConfig__update_general_settings_updates_max_parallel(monkeypatch):
monkeypatch.setattr(
"litellm.proxy.proxy_server.general_settings",
{},
)
pc = ProxyConfig()
await pc._update_general_settings(
{
"max_parallel_requests": 7,
"global_max_parallel_requests": 99,
"ui_access_mode": "admin_only",
}
)
from litellm.proxy import proxy_server as ps
snapshot = {
"max_parallel_requests": ps.general_settings.get("max_parallel_requests"),
"global_max_parallel_requests": ps.general_settings.get("global_max_parallel_requests"),
"ui_access_mode": ps.general_settings.get("ui_access_mode"),
}
assert snapshot == {
"max_parallel_requests": 7,
"global_max_parallel_requests": 99,
"ui_access_mode": "admin_only",
}
@pytest.mark.asyncio
async def test_ProxyConfig__update_general_settings_applies_db_max_batch_file_size_mb(monkeypatch):
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
pc = ProxyConfig()
await pc._update_general_settings({"max_batch_file_size_mb": 5})
from litellm.proxy import proxy_server as ps
assert ps.general_settings.get("max_batch_file_size_mb") == 5
@pytest.mark.asyncio
async def test_ProxyConfig__update_general_settings_yaml_max_batch_file_size_mb_wins_over_db(monkeypatch):
monkeypatch.setattr(
"litellm.proxy.proxy_server.general_settings",
{"max_batch_file_size_mb": 3},
)
pc = ProxyConfig()
await pc._update_general_settings({"max_batch_file_size_mb": 5})
from litellm.proxy import proxy_server as ps
assert ps.general_settings.get("max_batch_file_size_mb") == 3
@pytest.mark.asyncio
async def test_ProxyConfig__update_general_settings_cleared_db_max_batch_file_size_mb_lifts_cap(monkeypatch):
monkeypatch.setattr(
"litellm.proxy.proxy_server.general_settings",
{"max_batch_file_size_mb": 8},
)
pc = ProxyConfig()
await pc._update_general_settings({"max_parallel_requests": 1})
from litellm.proxy import proxy_server as ps
assert ps.general_settings.get("max_batch_file_size_mb") == 8
@pytest.mark.asyncio
async def test_ProxyConfig__update_general_settings_applies_db_allowed_file_extensions(monkeypatch):
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
pc = ProxyConfig()
await pc._update_general_settings({"allowed_file_extensions": [".jsonl"]})
from litellm.proxy import proxy_server as ps
assert ps.general_settings.get("allowed_file_extensions") == [".jsonl"]
@pytest.mark.asyncio
async def test_ProxyConfig__update_general_settings_yaml_allowed_file_extensions_wins_over_db(monkeypatch):
monkeypatch.setattr(
"litellm.proxy.proxy_server.general_settings",
{"allowed_file_extensions": [".pdf"]},
)
pc = ProxyConfig()
await pc._update_general_settings({"allowed_file_extensions": [".jsonl"]})
from litellm.proxy import proxy_server as ps
assert ps.general_settings.get("allowed_file_extensions") == [".pdf"]
@pytest.mark.asyncio
async def test_ProxyConfig__update_general_settings_applies_db_transcribe_media_buckets(monkeypatch):
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
pc = ProxyConfig()
await pc._update_general_settings({"transcribe_media_buckets": ["team-audio"]})
from litellm.proxy import proxy_server as ps
assert ps.general_settings.get("transcribe_media_buckets") == ["team-audio"]
@pytest.mark.asyncio
async def test_ProxyConfig__update_general_settings_yaml_transcribe_media_buckets_wins_over_db(monkeypatch):
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"transcribe_media_buckets": ["yaml-audio"]})
pc = ProxyConfig()
pc._yaml_general_settings_keys = {"transcribe_media_buckets"}
await pc._update_general_settings({"transcribe_media_buckets": ["team-audio"]})
from litellm.proxy import proxy_server as ps
assert ps.general_settings.get("transcribe_media_buckets") == ["yaml-audio"]
@pytest.mark.asyncio
async def test_ProxyConfig__update_general_settings_none_input_noop():
pc = ProxyConfig()
# None input returns early.
result = await pc._update_general_settings(db_general_settings=None)
assert result is None
# Error-style: dict() will fail on non-mapping non-None input.
with pytest.raises(TypeError):
await pc._update_general_settings(db_general_settings=12345) # type: ignore[arg-type]
@pytest.mark.asyncio
async def test_ProxyConfig__update_general_settings_skips_redundant_retention_reschedule(monkeypatch):
from litellm.proxy import proxy_server
pc = ProxyConfig()
reschedule: Final = AsyncMock()
monkeypatch.setattr(proxy_server, "general_settings", {})
monkeypatch.setattr(proxy_server, "scheduler", MagicMock())
monkeypatch.setattr(pc, "_reschedule_spend_log_cleanup_job", reschedule)
await pc._update_general_settings({"maximum_health_check_retention_period": "30d"})
reschedule.assert_awaited_once()
reschedule.reset_mock()
await pc._update_general_settings({"maximum_health_check_retention_period": "30d"})
reschedule.assert_not_awaited()
@pytest.mark.asyncio
async def test_ProxyConfig__update_general_settings_reschedules_after_retention_key_deletion(monkeypatch):
from litellm.proxy import proxy_server
pc = ProxyConfig()
reschedule: Final = AsyncMock()
monkeypatch.setattr(proxy_server, "general_settings", {})
monkeypatch.setattr(proxy_server, "scheduler", MagicMock(**{"get_job.return_value": None}))
monkeypatch.setattr(pc, "_reschedule_spend_log_cleanup_job", reschedule)
await pc._update_general_settings({"maximum_health_check_retention_period": "30d"})
reschedule.reset_mock()
await pc._update_general_settings({})
reschedule.assert_awaited_once()
@pytest.mark.asyncio
async def test_ProxyConfig__update_general_settings_dispatches_every_side_effect_handler(monkeypatch):
pc = ProxyConfig()
handlers: Final = (
("_apply_alerting_settings", AsyncMock()),
("_apply_pass_through_settings", AsyncMock()),
("_apply_boolean_settings", AsyncMock()),
("_apply_store_model_in_db_setting", AsyncMock()),
("_apply_retention_settings", AsyncMock()),
("_apply_ssrf_settings", AsyncMock()),
("_apply_cache_size_setting", AsyncMock()),
)
for name, handler in handlers:
monkeypatch.setattr(pc, name, handler)
await pc._apply_general_settings_side_effects({}, False, ())
for name, handler in handlers:
if name == "_apply_cache_size_setting":
handler.assert_awaited_once_with({}, cache_size_was_db=False)
elif name == "_apply_retention_settings":
handler.assert_awaited_once_with({}, previous_cleanup_schedule=())
else:
handler.assert_awaited_once_with({})
@pytest.mark.asyncio
async def test_ProxyConfig__update_general_settings_unrelated_value_fires_no_runtime_effect(monkeypatch):
from litellm.proxy import proxy_server
pc = ProxyConfig()
initialize_endpoints: Final = AsyncMock()
reschedule: Final = AsyncMock()
cache: Final = MagicMock()
proxy_logging: Final = MagicMock()
monkeypatch.setattr(proxy_server, "general_settings", {})
monkeypatch.setattr(proxy_server, "initialize_pass_through_endpoints", initialize_endpoints)
monkeypatch.setattr(proxy_server, "user_api_key_cache", cache)
monkeypatch.setattr(proxy_server, "proxy_logging_obj", proxy_logging)
monkeypatch.setattr(pc, "_reschedule_spend_log_cleanup_job", reschedule)
await pc._update_general_settings({"unrelated": "value"})
initialize_endpoints.assert_not_awaited()
reschedule.assert_not_awaited()
cache.update_in_memory_max_size.assert_not_called()
proxy_logging.update_values.assert_not_called()
proxy_logging.slack_alerting_instance.update_values.assert_not_called()
@pytest.mark.asyncio
async def test_ProxyConfig__update_config_from_db_resolves_through_settings_stores(monkeypatch):
pc = ProxyConfig()
config = {
"general_settings": {
"max_file_size_mb": 7,
"max_parallel_requests": 3,
"alerting": ["config"],
"pass_through_endpoints": [{"path": "/config"}],
"maximum_spend_logs_cleanup_batch_size": 10,
},
"router_settings": {"fallbacks": ["config"], "num_retries": 1},
}
db_values = {
"general_settings": {
"max_file_size_mb": 9,
"max_parallel_requests": 11,
"alerting": ["db"],
"pass_through_endpoints": [{"path": "/db"}],
"maximum_spend_logs_cleanup_batch_size": None,
},
"router_settings": {"fallbacks": [], "num_retries": 2},
}
async def get_config_param(_, param_name):
value = db_values.get(param_name)
return SimpleNamespace(param_name=param_name, param_value=value) if value is not None else None
monkeypatch.setattr("litellm.proxy.proxy_server.get_config_param", get_config_param)
pc._load_yaml_settings_stores(config)
resolved = await pc._update_config_from_db(MagicMock(), config, store_model_in_db=True)
assert resolved["general_settings"] == {
"max_file_size_mb": 7,
"max_parallel_requests": 3,
"alerting": ["config"],
"pass_through_endpoints": [{"path": "/db"}, {"path": "/config"}],
"maximum_spend_logs_cleanup_batch_size": 10,
}
assert resolved["router_settings"] == {"fallbacks": ["config"], "num_retries": 1}
assert pc.settings.source("max_file_size_mb") == "config"
assert pc.settings.source("max_parallel_requests") == "config"
@pytest.mark.asyncio
async def test_ProxyConfig__update_config_from_db_keeps_keys_the_config_file_omits(monkeypatch):
pc = ProxyConfig()
config = {"general_settings": {"max_file_size_mb": 7}, "router_settings": {"num_retries": 1}}
db_values = {
"general_settings": {"max_file_size_mb": 9, "max_parallel_requests": 11},
"router_settings": {"fallbacks": ["db"], "num_retries": 2},
}
async def get_config_param(_, param_name):
value = db_values.get(param_name)
return SimpleNamespace(param_name=param_name, param_value=value) if value is not None else None
monkeypatch.setattr("litellm.proxy.proxy_server.get_config_param", get_config_param)
pc._load_yaml_settings_stores(config)
resolved = await pc._update_config_from_db(MagicMock(), config, store_model_in_db=True)
assert resolved["general_settings"] == {"max_file_size_mb": 7, "max_parallel_requests": 11}
assert resolved["router_settings"] == {"num_retries": 1, "fallbacks": ["db"]}
assert pc.settings.source("max_parallel_requests") == "db"
@pytest.mark.asyncio
async def test_ProxyConfig_add_deployment_continues_after_null_pass_through_endpoints(monkeypatch):
from litellm.proxy import proxy_server
pc = ProxyConfig()
non_llm_initialization = AsyncMock()
monkeypatch.setattr(proxy_server, "general_settings", {})
monkeypatch.setattr(proxy_server, "prefetch_config_params", AsyncMock())
monkeypatch.setattr(
proxy_server,
"get_config_param",
AsyncMock(return_value=SimpleNamespace(param_value={"pass_through_endpoints": None})),
)
monkeypatch.setattr(proxy_server, "sync_ui_settings_to_general_settings", AsyncMock())
monkeypatch.setattr(pc, "_should_load_db_object", lambda *, object_type: False)
monkeypatch.setattr(pc, "get_credentials", AsyncMock())
monkeypatch.setattr(pc, "_init_non_llm_objects_in_db", non_llm_initialization)
await pc.add_deployment(prisma_client=MagicMock(), proxy_logging_obj=MagicMock())
non_llm_initialization.assert_awaited_once()
# ---------------------------------------------------------------------------
# ProxyConfig._update_config_from_db
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_ProxyConfig__update_config_from_db_does_not_log_general_settings_secrets(
monkeypatch,
):
"""Regression for LIT-4152 on the store_model_in_db path.
``_update_config_from_db`` logged each DB ``param_value`` verbatim at DEBUG;
for ``general_settings`` that value is the whole dict, leaking ``master_key``
and ``database_url`` the same way the startup config load did. The value now
routes through the recursive redactor. Asserted with the module regex
scrubber (``_ENABLE_SECRET_REDACTION``) disabled so the caller itself must
not build the leaky string. The merge into the returned config must still
carry the raw values, proving only the log record is redacted.
"""
import logging
import litellm._logging as _logging_module
from litellm._logging import verbose_proxy_logger
monkeypatch.setattr(_logging_module, "_ENABLE_SECRET_REDACTION", False)
monkeypatch.delenv("DATABASE_URL", raising=False)
def _fake_decrypt_value_helper(value, key, **_kwargs):
return value
monkeypatch.setattr("litellm.proxy.proxy_server.decrypt_value_helper", _fake_decrypt_value_helper)
master_key_secret = "sk-lit4152-db-path-master-key-abcdef1234567890"
db_url_secret = "postgresql://leak_user:leak_password_9090@leak-host.internal:5432/leak_db"
env_db_url_secret = "postgresql://env_leak_user:env_leak_password_9090@env-leak-host.internal:5432/env_leak_db"
nested_webhook_secret = "https://hooks.slack.com/services/T0/B0/db-path-webhook-secret"
responses = {
"general_settings": SimpleNamespace(
param_name="general_settings",
param_value={
"master_key": master_key_secret,
"database_url": db_url_secret,
"alert_to_webhook_url": {"budget_alerts": nested_webhook_secret},
},
),
"router_settings": None,
"litellm_settings": None,
"environment_variables": SimpleNamespace(
param_name="environment_variables",
param_value={"DATABASE_URL": env_db_url_secret},
),
}
async def _fake_get_config_param(prisma_client, key):
return responses[key]
monkeypatch.setattr("litellm.proxy.proxy_server.get_config_param", _fake_get_config_param)
class LogRecordHandler(logging.Handler):
def __init__(self) -> None:
super().__init__()
self.records: list[logging.LogRecord] = []
def emit(self, record: logging.LogRecord) -> None:
self.records.append(record)
handler = LogRecordHandler()
handler.setLevel(logging.DEBUG)
original_level = verbose_proxy_logger.level
verbose_proxy_logger.setLevel(logging.DEBUG)
verbose_proxy_logger.addHandler(handler)
try:
merged = await ProxyConfig()._update_config_from_db(
prisma_client=MagicMock(),
config={"general_settings": {}},
store_model_in_db=True,
)
rendered = " ".join(record.getMessage() for record in handler.records)
finally:
verbose_proxy_logger.removeHandler(handler)
verbose_proxy_logger.setLevel(original_level)
for secret in (
master_key_secret,
db_url_secret,
env_db_url_secret,
nested_webhook_secret,
"leak_password_9090",
"env_leak_password_9090",
):
assert secret not in rendered, f"leak: {secret} in {rendered!r}"
assert merged["general_settings"]["master_key"] == master_key_secret
assert merged["general_settings"]["database_url"] == db_url_secret
assert merged["environment_variables"]["DATABASE_URL"] == env_db_url_secret
@pytest.mark.asyncio
async def test_ProxyConfig_load_config_redacts_secret_litellm_setting_keeps_plain(tmp_path, monkeypatch):
"""Regression for LIT-4152 on the ``litellm_settings`` apply loop.
``load_config`` logged ``setting litellm.<key>=<value>`` verbatim at DEBUG,
so a secret-bearing setting such as ``api_key`` leaked in cleartext. The
value now routes through ``_redact_general_setting_value``. Crucially the
redaction must be surgical: a secret-named key is masked, but a plain
operational setting like ``num_retries`` must still log its real value, so
the debug line keeps its signal. Asserted with the module regex scrubber
(``_ENABLE_SECRET_REDACTION``) disabled.
"""
import logging
import litellm._logging as _logging_module
from litellm._logging import verbose_proxy_logger
monkeypatch.setattr(_logging_module, "_ENABLE_SECRET_REDACTION", False)
api_key_secret = "sk-lit4152-litellm-settings-secret-abcdef1234567890"
f = tmp_path / "c.yaml"
f.write_text(
f"model_list: []\ngeneral_settings: {{}}\nlitellm_settings:\n api_key: {api_key_secret}\n num_retries: 7\n"
)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False)
monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False)
class LogRecordHandler(logging.Handler):
def __init__(self) -> None:
super().__init__()
self.records: list[logging.LogRecord] = []
def emit(self, record: logging.LogRecord) -> None:
self.records.append(record)
handler = LogRecordHandler()
handler.setLevel(logging.DEBUG)
original_level = verbose_proxy_logger.level
original_api_key = getattr(litellm, "api_key", None)
original_num_retries = getattr(litellm, "num_retries", None)
verbose_proxy_logger.setLevel(logging.DEBUG)
verbose_proxy_logger.addHandler(handler)
try:
await ProxyConfig().load_config(router=None, config_file_path=str(f))
rendered = " ".join(record.getMessage() for record in handler.records)
finally:
verbose_proxy_logger.removeHandler(handler)
verbose_proxy_logger.setLevel(original_level)
litellm.api_key = original_api_key
litellm.num_retries = original_num_retries
assert api_key_secret not in rendered, f"api_key leaked in logs: {rendered!r}"
assert "num_retries=7" in rendered, (
f"non-secret num_retries value was over-redacted; expected it visible in {rendered!r}"
)
# ---------------------------------------------------------------------------
# ProxyConfig agents from config.yaml
# ---------------------------------------------------------------------------
@pytest.fixture
def clean_agent_registry():
from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry
original_agents = list(global_agent_registry.agent_list)
original_config_agents = getattr(global_agent_registry, "config_agents", ())
global_agent_registry.agent_list = []
global_agent_registry.config_agents = ()
try:
yield global_agent_registry
finally:
global_agent_registry.agent_list = original_agents
global_agent_registry.config_agents = original_config_agents
def _config_agent(agent_name: str) -> Dict[str, Any]:
return {
"agent_name": agent_name,
"agent_card_params": {
"name": "Config Agent",
"url": "http://localhost:10001",
"protocolVersion": "1.0",
},
}
def _agent_db_row(agent_id: str, agent_name: str):
import json
from datetime import datetime, timezone
from prisma.models import LiteLLM_AgentsTable
return LiteLLM_AgentsTable(
agent_id=agent_id,
agent_name=agent_name,
agent_card_params=json.dumps({"name": agent_name, "url": "http://db-agent"}),
extra_headers=[],
agent_access_groups=[],
access_group_ids=[],
spend=0.0,
identity_managed=False,
enabled=True,
execution_mode="autonomous",
created_at=datetime.now(timezone.utc),
updated_at=datetime.now(timezone.utc),
created_by="admin",
updated_by="admin",
)
@pytest.mark.asyncio
@pytest.mark.parametrize("config_key", ["agents", "agent_list"])
async def test_ProxyConfig__init_non_llm_configs_registers_agents_from_config(clean_agent_registry, config_key):
"""The documented ``agents:`` key must register agents, as must the legacy ``agent_list:``."""
await ProxyConfig()._init_non_llm_configs(
config={config_key: [_config_agent("config-agent")]},
config_file_path=None,
)
assert [agent.agent_name for agent in clean_agent_registry.get_agent_list()] == ["config-agent"]
@pytest.mark.asyncio
async def test_ProxyConfig__init_agents_in_db_keeps_config_defined_agents(clean_agent_registry):
"""A DB reload rebuilds the registry; config-defined agents must survive it alongside DB rows."""
await ProxyConfig()._init_non_llm_configs(
config={"agents": [_config_agent("config-agent")]},
config_file_path=None,
)
prisma_client = MagicMock()
prisma_client.db.litellm_agentstable.find_many = AsyncMock(return_value=[_agent_db_row("db-id", "db-agent")])
await ProxyConfig()._init_agents_in_db(prisma_client=prisma_client)
assert sorted(agent.agent_name for agent in clean_agent_registry.get_agent_list()) == [
"config-agent",
"db-agent",
]
@pytest.mark.asyncio
@pytest.mark.parametrize("agents_source", ["config", "db", "api"])
async def test_ProxyStartupEvent_jwt_auth_resolves_agent_claims_against_live_registry(
clean_agent_registry, agents_source
):
"""A JWT agent claim must resolve against every agent the proxy knows, including ones created after startup."""
from litellm.proxy import proxy_server
from litellm.proxy._types import LiteLLM_JWTAuth
from litellm.proxy.auth.handle_jwt import JWTAuthManager
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.types.agents import AgentResponse
original_lookup = proxy_server.jwt_handler.agent_lookup
try:
proxy_server.ProxyStartupEvent._initialize_jwt_auth(
general_settings={"litellm_jwtauth": {"agent_id_jwt_field": "appid"}},
prisma_client=None,
user_api_key_cache=UserApiKeyCache(),
)
if agents_source == "config":
await ProxyConfig()._init_non_llm_configs(
config={"agents": [_config_agent("loaded-agent")]},
config_file_path=None,
)
elif agents_source == "db":
prisma_client = MagicMock()
prisma_client.db.litellm_agentstable.find_many = AsyncMock(
return_value=[_agent_db_row("db-id", "loaded-agent")]
)
await ProxyConfig()._init_agents_in_db(prisma_client=prisma_client)
else:
clean_agent_registry.register_agent(
agent_config=AgentResponse(agent_id="api-id", **_config_agent("loaded-agent"))
)
resolved = JWTAuthManager.resolve_agent_id(
jwt_handler=proxy_server.jwt_handler,
jwt_valid_token={"appid": "loaded-agent"},
agent_registry=proxy_server.jwt_handler.agent_lookup,
)
finally:
proxy_server.jwt_handler.bind_agent_lookup(original_lookup)
proxy_server.jwt_handler.update_environment(
prisma_client=None, user_api_key_cache=UserApiKeyCache(), litellm_jwtauth=LiteLLM_JWTAuth()
)
assert resolved == clean_agent_registry.get_agent_by_name(agent_name="loaded-agent").agent_id
@pytest.mark.asyncio
@pytest.mark.parametrize(
"config, expected_agent_names",
[
({"agents": [], "agent_list": [_config_agent("legacy-agent")]}, []),
(
{
"agents": [_config_agent("documented-agent")],
"agent_list": [_config_agent("legacy-agent")],
},
["documented-agent"],
),
({"agent_list": [_config_agent("legacy-agent")]}, ["legacy-agent"]),
],
ids=["empty-agents-wins", "populated-agents-wins", "agent_list-alone-still-works"],
)
async def test_ProxyConfig__init_non_llm_configs_prefers_agents_key_by_presence(
clean_agent_registry, config, expected_agent_names
):
"""
``agents`` outranks the legacy ``agent_list`` whenever the key is present.
Selecting on truthiness instead would silently register the legacy entries
for a config that spells out ``agents: []``.
"""
await ProxyConfig()._init_non_llm_configs(config=config, config_file_path=None)
assert [agent.agent_name for agent in clean_agent_registry.get_agent_list()] == expected_agent_names
@pytest.mark.asyncio
async def test_ProxyConfig__init_non_llm_configs_empty_agents_key_clears_remembered_agents(clean_agent_registry):
"""
An explicitly empty ``agents:`` must reach the registry, not be skipped as falsy.
Skipping it leaves the previously remembered agents in place, so the next DB
rebuild replays agents the operator deleted from config.yaml.
"""
clean_agent_registry.load_agents_from_config([_config_agent("stale-agent")])
assert clean_agent_registry.config_agents != ()
await ProxyConfig()._init_non_llm_configs(config={"agents": []}, config_file_path=None)
assert clean_agent_registry.config_agents == ()
clean_agent_registry.load_agents_from_db_and_config(db_agents=None)
assert clean_agent_registry.get_agent_list() == ()
# ---------------------------------------------------------------------------
# _init_guardrails_in_db
# ---------------------------------------------------------------------------
def _db_guardrail_row(guardrail_id: str, guardrail_type: str) -> dict[str, object]:
return {
"guardrail_id": guardrail_id,
"guardrail_name": f"name-{guardrail_id}",
"litellm_params": {"guardrail": guardrail_type, "mode": "pre_call"},
"guardrail_info": None,
"team_id": None,
}
@pytest.mark.asyncio
async def test_ProxyConfig__init_guardrails_in_db_skips_only_the_unloadable_row(monkeypatch):
"""
A single DB row that fails to initialize used to abort the whole loop, so one
typo'd guardrail type left the proxy running with zero guardrails loaded.
The failing row's id must still reach reconcile_db_guardrails so that eviction
pass cannot treat a row that is alive in the DB as one that was deleted.
"""
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.proxy.guardrails import guardrail_registry as registry_module
from litellm.types.guardrails import Guardrail, GuardrailEventHooks, LitellmParams
class _RecordingHandler(registry_module.InMemoryGuardrailHandler):
def __init__(self) -> None:
super().__init__()
self.reconciled_with: list[set[str]] = []
def reconcile_db_guardrails(self, db_guardrail_ids: set[str]) -> list[str]:
self.reconciled_with.append(set(db_guardrail_ids))
return super().reconcile_db_guardrails(db_guardrail_ids)
handler = _RecordingHandler()
monkeypatch.setattr(registry_module, "IN_MEMORY_GUARDRAIL_HANDLER", handler)
def _initializer(litellm_params: LitellmParams, guardrail: Guardrail) -> CustomGuardrail:
return CustomGuardrail(
guardrail_name=guardrail["guardrail_name"],
event_hook=GuardrailEventHooks.pre_call,
default_on=False,
)
monkeypatch.setitem(registry_module.guardrail_initializer_registry, "lit5367_ok", _initializer)
prisma_client = MagicMock()
prisma_client.db.litellm_guardrailstable.find_many = AsyncMock(
return_value=[
_db_guardrail_row("first", "lit5367_ok"),
_db_guardrail_row("broken", "litellm_tool_permission"),
_db_guardrail_row("last", "lit5367_ok"),
]
)
await ProxyConfig()._init_guardrails_in_db(prisma_client=prisma_client)
assert sorted(handler.IN_MEMORY_GUARDRAILS) == ["first", "last"]
assert handler.reconciled_with == [{"first", "broken", "last"}]
# ---------------------------------------------------------------------------
# add_deployment: UI settings convergence
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_add_deployment_re_reads_ui_settings_so_other_pods_converge(monkeypatch):
"""The periodic config reload picks up a UI setting written through another pod.
Startup used to be the only read, so a proxy admin flipping a runtime flag reached the pod
that served the PATCH and nowhere else until every other pod restarted.
"""
general_settings: Dict[str, Any] = {"allow_agents_for_team_admins": False}
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", general_settings)
prisma_client = MagicMock()
prisma_client.db.litellm_config.find_many = AsyncMock(return_value=[])
prisma_client.db.litellm_config.find_first = AsyncMock(return_value=None)
prisma_client.db.litellm_credentialstable.find_many = AsyncMock(return_value=[])
prisma_client.db.litellm_uisettings.find_unique = AsyncMock(
return_value=SimpleNamespace(
ui_settings=json.dumps({"allow_agents_for_team_admins": True, "enable_chat_ui": False})
)
)
config = ProxyConfig()
config._should_load_db_object = MagicMock(return_value=False)
config._init_non_llm_objects_in_db = AsyncMock()
await config.add_deployment(prisma_client=prisma_client, proxy_logging_obj=MagicMock())
prisma_client.db.litellm_uisettings.find_unique.assert_awaited_once_with(where={"id": "ui_settings"})
assert general_settings["allow_agents_for_team_admins"] is True
assert "enable_chat_ui" not in general_settings
@pytest.mark.asyncio
async def test_add_deployment_syncs_ui_settings_even_when_the_model_reconcile_fails(monkeypatch):
"""A broken model reconcile must not strand every pod on stale settings."""
general_settings: Dict[str, Any] = {"allow_agents_for_team_admins": False}
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", general_settings)
prisma_client = MagicMock()
prisma_client.db.litellm_uisettings.find_unique = AsyncMock(
return_value=SimpleNamespace(ui_settings={"allow_agents_for_team_admins": True})
)
config = ProxyConfig()
config._should_load_db_object = MagicMock(side_effect=RuntimeError("db down"))
await config.add_deployment(prisma_client=prisma_client, proxy_logging_obj=MagicMock())
assert general_settings["allow_agents_for_team_admins"] is True
def _websearch_logger_cls():
from litellm.integrations.websearch_interception.handler import (
WebSearchInterceptionLogger,
)
return WebSearchInterceptionLogger
def _run_websearch_init(monkeypatch, stored_params, starting_callbacks):
pc = ProxyConfig()
monkeypatch.setattr(litellm, "callbacks", list(starting_callbacks))
monkeypatch.setattr(
"litellm.proxy.proxy_server.get_config_param",
AsyncMock(return_value=SimpleNamespace(param_value={"websearch_interception_params": stored_params}))
if stored_params is not None
else AsyncMock(return_value=SimpleNamespace(param_value={})),
)
asyncio.run(pc.init_websearch_interception_settings_in_db(prisma_client=MagicMock()))
return pc
def _poll_websearch_init(pc, monkeypatch, stored_params):
monkeypatch.setattr(
"litellm.proxy.proxy_server.get_config_param",
AsyncMock(return_value=SimpleNamespace(param_value={"websearch_interception_params": stored_params})),
)
asyncio.run(pc.init_websearch_interception_settings_in_db(prisma_client=MagicMock()))
def test_init_websearch_interception_resyncs_after_a_write_drops_the_enabled_flag(monkeypatch):
logger_cls = _websearch_logger_cls()
pc = ProxyConfig()
monkeypatch.setattr(litellm, "callbacks", [])
_poll_websearch_init(pc, monkeypatch, {"enabled": True, "search_tool_name": "old-tool"})
_poll_websearch_init(pc, monkeypatch, {"search_tool_name": "new-tool"})
registered = [cb for cb in litellm.callbacks if isinstance(cb, logger_cls)]
assert len(registered) == 1
assert registered[0].search_tool_name == "new-tool"
def test_init_websearch_interception_ignores_a_non_list_providers_value(monkeypatch):
logger_cls = _websearch_logger_cls()
_run_websearch_init(
monkeypatch,
stored_params={"enabled": True, "enabled_providers": "bedrock", "search_tool_name": "stored-tool"},
starting_callbacks=[],
)
registered = [cb for cb in litellm.callbacks if isinstance(cb, logger_cls)]
assert len(registered) == 1
assert registered[0].enabled_providers == ["bedrock"]
def test_init_websearch_interception_absent_key_leaves_callbacks_untouched(monkeypatch):
logger_cls = _websearch_logger_cls()
config_registered = logger_cls(search_tool_name="from-config-yaml")
_run_websearch_init(monkeypatch, stored_params=None, starting_callbacks=[config_registered])
assert litellm.callbacks == [config_registered]
def test_init_websearch_interception_without_enabled_key_leaves_callbacks_untouched(monkeypatch):
logger_cls = _websearch_logger_cls()
config_registered = logger_cls(search_tool_name="from-config-yaml")
_run_websearch_init(
monkeypatch,
stored_params={"search_tool_name": "stored-tool"},
starting_callbacks=[config_registered],
)
assert litellm.callbacks == [config_registered]
def test_init_websearch_interception_registers_when_explicitly_enabled(monkeypatch):
logger_cls = _websearch_logger_cls()
_run_websearch_init(
monkeypatch,
stored_params={"enabled": True, "search_tool_name": "stored-tool"},
starting_callbacks=[],
)
registered = [cb for cb in litellm.callbacks if isinstance(cb, logger_cls)]
assert len(registered) == 1
assert registered[0].search_tool_name == "stored-tool"
def test_init_websearch_interception_treats_string_false_as_disabled(monkeypatch):
logger_cls = _websearch_logger_cls()
existing = logger_cls(search_tool_name="stored-tool")
_run_websearch_init(
monkeypatch,
stored_params={"enabled": "false", "search_tool_name": "stored-tool"},
starting_callbacks=[existing],
)
assert [cb for cb in litellm.callbacks if isinstance(cb, logger_cls)] == []
def test_init_websearch_interception_empty_providers_falls_back_to_handler_default(monkeypatch):
logger_cls = _websearch_logger_cls()
_run_websearch_init(
monkeypatch,
stored_params={"enabled": True, "enabled_providers": [], "search_tool_name": "stored-tool"},
starting_callbacks=[],
)
registered = [cb for cb in litellm.callbacks if isinstance(cb, logger_cls)]
assert len(registered) == 1
assert registered[0].enabled_providers == ["bedrock"]
def test_init_websearch_interception_keeps_working_callback_when_new_one_cannot_be_built(monkeypatch):
logger_cls = _websearch_logger_cls()
working = logger_cls(search_tool_name="stored-tool", max_agentic_loops=3)
_run_websearch_init(
monkeypatch,
stored_params={"enabled": True, "search_tool_name": "stored-tool", "max_agentic_loops": 0},
starting_callbacks=[working],
)
assert litellm.callbacks == [working]
def test_init_websearch_interception_disabled_removes_the_callback(monkeypatch):
logger_cls = _websearch_logger_cls()
existing = logger_cls(search_tool_name="stored-tool")
_run_websearch_init(
monkeypatch,
stored_params={"enabled": False, "search_tool_name": "stored-tool"},
starting_callbacks=[existing],
)
assert [cb for cb in litellm.callbacks if isinstance(cb, logger_cls)] == []
def test_init_websearch_interception_replaces_stale_instance_on_param_change(monkeypatch):
logger_cls = _websearch_logger_cls()
stale = logger_cls(search_tool_name="old-tool", max_agentic_loops=2)
_run_websearch_init(
monkeypatch,
stored_params={"enabled": True, "search_tool_name": "new-tool", "max_agentic_loops": 7},
starting_callbacks=[stale],
)
registered = [cb for cb in litellm.callbacks if isinstance(cb, logger_cls)]
assert len(registered) == 1
assert (registered[0].search_tool_name, registered[0].max_agentic_loops) == ("new-tool", 7)
def test_init_websearch_interception_honors_enabled_providers(monkeypatch):
logger_cls = _websearch_logger_cls()
_run_websearch_init(
monkeypatch,
stored_params={"enabled": True, "enabled_providers": ["bedrock", "vertex_ai"]},
starting_callbacks=[],
)
registered = [cb for cb in litellm.callbacks if isinstance(cb, logger_cls)]
assert len(registered) == 1
assert registered[0].enabled_providers == ["bedrock", "vertex_ai"]
def test_websearch_interception_settings_can_be_named_in_supported_db_objects(monkeypatch):
from litellm.proxy import proxy_server
from litellm.proxy._types import ConfigGeneralSettings
allowlist = ConfigGeneralSettings(supported_db_objects=["websearch_interception_settings"]).supported_db_objects
assert allowlist
monkeypatch.setattr(proxy_server, "general_settings", {"supported_db_objects": allowlist})
assert proxy_server.should_load_db_object(object_type="websearch_interception_settings") is True
monkeypatch.setattr(proxy_server, "general_settings", {"supported_db_objects": ["models"]})
assert proxy_server.should_load_db_object(object_type="websearch_interception_settings") is False
def test_validate_deployment_access_windows_rejects_malformed_time():
model = {
"model_name": "gpt-4o-shared",
"litellm_params": {"model": "gpt-4o"},
"model_info": {
"access_windows": [{"start": "25:00", "end": "06:00", "timezone": "America/New_York", "team_ids": ["t"]}]
},
}
with pytest.raises(ValueError, match="access_windows") as exc_info:
validate_deployment_access_windows(model)
assert "gpt-4o-shared" in str(exc_info.value)
def test_validate_deployment_access_windows_rejects_unknown_timezone():
model = {
"model_name": "gpt-4o-shared",
"litellm_params": {"model": "gpt-4o"},
"model_info": {
"access_windows": [{"start": "22:00", "end": "06:00", "timezone": "Mars/Olympus", "team_ids": ["t"]}]
},
}
with pytest.raises(ValueError, match="Mars/Olympus"):
validate_deployment_access_windows(model)
def test_validate_deployment_access_windows_accepts_valid_and_absent():
assert (
validate_deployment_access_windows(
{
"model_name": "gpt-4o-shared",
"litellm_params": {"model": "gpt-4o"},
"model_info": {
"access_windows": [
{"start": "22:00", "end": "06:00", "timezone": "America/New_York", "team_ids": ["t"]}
]
},
}
)
is None
)
assert validate_deployment_access_windows({"model_name": "m", "litellm_params": {"model": "m"}}) is None
assert (
validate_deployment_access_windows(
{"model_name": "m", "litellm_params": {"model": "m"}, "model_info": {"id": "x"}}
)
is None
)
@pytest.mark.asyncio
async def test_model_refresh_updates_availability_catalog_and_retains_it_on_db_failure():
pc = ProxyConfig()
row = SimpleNamespace(
model_id="gated",
created_by="owner",
model_info={},
litellm_params={
"model": "auto_router/complexity_router",
"complexity_router_config": {"classifier_type": "heuristic_v2"},
},
)
find_many = AsyncMock(side_effect=[[row], RuntimeError("database unavailable"), []])
client = SimpleNamespace(db=SimpleNamespace(litellm_proxymodeltable=SimpleNamespace(find_many=find_many)))
assert pc.auto_router_db_catalog is None
assert await pc._get_models_from_db(client) == [row]
loaded = pc.auto_router_db_catalog
assert loaded is not None and loaded[0].model_id == "gated"
assert await pc._get_models_from_db(client) is None
assert pc.auto_router_db_catalog == loaded
assert await pc._get_models_from_db(client) == []
assert pc.auto_router_db_catalog == ()
assert find_many.await_count == 3
@pytest.mark.asyncio
@pytest.mark.parametrize("versions", [None, ["2024-11-05"], [], ["2026-07-28"], ["unknown"]])
async def test_proxy_config_validates_advertised_mcp_versions_at_load(tmp_path, monkeypatch, versions):
config = tmp_path / "mcp-versions.yaml"
config.write_text(json.dumps({"model_list": [], "general_settings": {"mcp_advertised_versions": versions}}))
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None)
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False)
monkeypatch.delenv("LITELLM_CONFIG_BUCKET_NAME", raising=False)
if versions is None or versions == ["2024-11-05"]:
_, _, settings = await ProxyConfig().load_config(router=None, config_file_path=str(config))
assert settings["mcp_advertised_versions"] == versions
return
with pytest.raises(ValidationError):
await ProxyConfig().load_config(router=None, config_file_path=str(config))
Reference in a new issue View git blame Copy permalink