From 6062ed7ff19f3c3747391298cf99f464790dc993 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 31 Jul 2026 19:57:12 +0000 Subject: [PATCH 01/40] fix(batches): encode public model group on background-created output file ids CheckBatchCost built unified output file ids with the provider model name, so key model-access checks resolved the file to e.g. gpt-5.5 and every GET /v1/files/{output_file_id}/content failed. Resolve the model group from the batch's managed input file, falling back to the deployment's model_name. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/common_utils/check_batch_cost.py | 33 +++- .../proxy_unit_tests/test_check_batch_cost.py | 183 +++++++++++++++++- 2 files changed, 213 insertions(+), 3 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py index f209ab54f64..22f9f40ecd8 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -17,6 +17,7 @@ if TYPE_CHECKING: from litellm.proxy._types import LiteLLM_ManagedObjectTable from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.router import Router + from litellm.types.router import Deployment from litellm.types.utils import LiteLLMBatch @@ -281,6 +282,32 @@ class CheckBatchCost: return deployment_id return None + @classmethod + def _get_managed_file_model_name( + cls, + job: "LiteLLM_ManagedObjectTable", + deployment_info: "Deployment", + ) -> Optional[str]: + """ + Public model group name to encode as ``target_model_names`` on unified output file ids. + + Key model-access checks resolve a managed file id back to a model via its + ``target_model_names``, so this must be the model group the caller requested, never the + underlying provider model (e.g. ``gpt-5.5``), which no key is allowed to call. + """ + from litellm.proxy.openai_files_endpoints.common_utils import ( + convert_b64_uid_to_unified_uid, + get_models_from_unified_file_id, + ) + + input_file_id = cls._get_input_file_id(job) + target_model_names = ( + get_models_from_unified_file_id(convert_b64_uid_to_unified_uid(input_file_id)) if input_file_id else [] + ) + if target_model_names: + return ",".join(target_model_names) + return deployment_info.model_name or None + @staticmethod def _get_input_file_id(job: "LiteLLM_ManagedObjectTable") -> Optional[str]: import json @@ -406,6 +433,10 @@ class CheckBatchCost: managed_files_hook = self.proxy_logging_obj.get_proxy_hook("managed_files") if managed_files_hook is not None: from litellm.proxy._types import UserAPIKeyAuth + + managed_file_model_name = self._get_managed_file_model_name( + job=job, deployment_info=deployment_info + ) _minimal_auth = UserAPIKeyAuth( user_id=job.created_by or "default-user-id", team_id=getattr(job, "team_id", None), @@ -417,7 +448,7 @@ class CheckBatchCost: _unified_file_id = managed_files_hook.get_unified_output_file_id( output_file_id=_raw_file_id, model_id=model_id, - model_name=str(model_name) if model_name else deployment_info.model_name or None, + model_name=managed_file_model_name, ) await managed_files_hook.store_unified_file_id( file_id=_unified_file_id, diff --git a/tests/proxy_unit_tests/test_check_batch_cost.py b/tests/proxy_unit_tests/test_check_batch_cost.py index b822799fb40..a15abd023d8 100644 --- a/tests/proxy_unit_tests/test_check_batch_cost.py +++ b/tests/proxy_unit_tests/test_check_batch_cost.py @@ -674,12 +674,12 @@ class TestCheckBatchCost: mock_hook.get_unified_output_file_id.assert_any_call( output_file_id=raw_output_file_id, model_id="model-123", - model_name="gpt-5-mini", + model_name="gpt-5-batch", ) mock_hook.get_unified_output_file_id.assert_any_call( output_file_id=raw_error_file_id, model_id="model-123", - model_name="gpt-5-mini", + model_name="gpt-5-batch", ) assert mock_hook.store_unified_file_id.await_count == 2 # {raw_file_id: managed_file_id} for each store call @@ -1221,3 +1221,182 @@ class TestUnmanagedBatchCostFlagIsGeneralized: assert vertex_result == ("deploy-vertex", "8823717160934178816") assert bedrock_result == ("deploy-bedrock", TestUnmanagedBedrockRouting._ARN) + + +class TestManagedOutputFileIdEncodesPublicModelGroup: + """LIT-4964 regression: the unified output file id created by the background poller must + encode the public model group as ``target_model_names``, not the provider model. + + Key model-access checks resolve a managed file id back to a model via ``target_model_names``, + so encoding the provider model (e.g. ``gpt-5.5``) makes + ``GET /v1/files/{output_file_id}/content`` fail for every key. + """ + + _PUBLIC_MODEL_GROUP = "gpt-5-batch" + _RAW_OUTPUT_FILE_ID = "file-batch-output-abc123" + + @staticmethod + def _managed_input_file_id(model_group: str) -> str: + import base64 + + unified_id = ( + "litellm_proxy:application/octet-stream;unified_id,c4843482-b176-4901-8292-7523fd0f2c6e;" + f"target_model_names,{model_group}" + ) + return base64.urlsafe_b64encode(unified_id.encode()).decode().rstrip("=") + + def _job(self, input_file_id: str) -> MagicMock: + from litellm.types.utils import LiteLLMBatch + + job = MagicMock() + job.id = "job-lit-4964" + job.unified_object_id = "dW5pZmllZF9iYXRjaF9pZA==" + job.created_by = "user-1" + job.team_id = None + job.file_object = LiteLLMBatch( + id="batch-456", + completion_window="24h", + created_at=1, + endpoint="/v1/chat/completions", + input_file_id=input_file_id, + object="batch", + status="completed", + ).model_dump_json() + return job + + async def _run(self, job: MagicMock) -> str: + from litellm_enterprise.proxy.common_utils.check_batch_cost import ( + CheckBatchCost, + ) + from litellm.types.utils import LiteLLMBatch + from enterprise.litellm_enterprise.proxy.hooks.managed_files import ( + _PROXY_LiteLLMManagedFiles, + ) + + router = MagicMock() + router.get_deployment_credentials_with_provider = MagicMock( + return_value={"api_key": "sk-test"} + ) + deployment = MagicMock() + deployment.litellm_params.custom_llm_provider = "azure" + deployment.litellm_params.model = "azure/gpt-5.5" + deployment.model_name = self._PUBLIC_MODEL_GROUP + deployment.model_info.model_dump.return_value = {} + router.get_deployment = MagicMock(return_value=deployment) + + hook = MagicMock() + hook.get_unified_output_file_id = ( + lambda output_file_id, model_id, model_name: _PROXY_LiteLLMManagedFiles.get_unified_output_file_id( + None, output_file_id=output_file_id, model_id=model_id, model_name=model_name + ) + ) + hook.store_unified_file_id = AsyncMock() + proxy_logging_obj = MagicMock() + proxy_logging_obj.get_proxy_hook.return_value = hook + + prisma_client = MagicMock() + prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=None) + + instance = CheckBatchCost( + proxy_logging_obj=proxy_logging_obj, + prisma_client=prisma_client, + llm_router=router, + ) + + response = LiteLLMBatch( + id="batch-456", + completion_window="24h", + created_at=1, + endpoint="/v1/chat/completions", + input_file_id=job.file_object, + object="batch", + status="completed", + ) + response.output_file_id = self._RAW_OUTPUT_FILE_ID + + file_content = MagicMock() + file_content.content = b'{"id":"req-1"}' + + with ( + patch( + "litellm.files.main.afile_content", + new_callable=AsyncMock, + return_value=file_content, + ), + patch( + "litellm.batches.batch_utils._get_file_content_as_dictionary", + return_value=[{"id": "req-1"}], + ), + patch( + "litellm.batches.batch_utils.calculate_batch_cost_and_usage", + new_callable=AsyncMock, + return_value=(0.01, {"prompt_tokens": 10}, ["gpt-5.5"]), + ), + patch("litellm.litellm_core_utils.litellm_logging.Logging") as logging_cls, + ): + logging_obj = MagicMock() + logging_obj.async_success_handler = AsyncMock() + logging_cls.return_value = logging_obj + + await instance._track_completed_batch_cost( + job=job, + response=response, + model_id="model-123", + batch_id="batch-456", + prom_logger=None, + ) + + return response.output_file_id + + @pytest.mark.asyncio + async def test_target_model_names_comes_from_input_file_not_provider_model(self): + from litellm.proxy.openai_files_endpoints.common_utils import ( + _is_base64_encoded_unified_file_id, + get_models_from_unified_file_id, + ) + + output_file_id = await self._run( + self._job(self._managed_input_file_id(self._PUBLIC_MODEL_GROUP)) + ) + + decoded = _is_base64_encoded_unified_file_id(output_file_id) + assert get_models_from_unified_file_id(decoded) == [self._PUBLIC_MODEL_GROUP] + assert "gpt-5.5" not in decoded + + @pytest.mark.asyncio + async def test_key_scoped_to_model_group_can_read_the_output_file(self): + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.auth.auth_checks import can_key_call_model + from litellm.proxy.auth.auth_utils import ( + _extract_models_from_managed_resource_id, + ) + + output_file_id = await self._run( + self._job(self._managed_input_file_id(self._PUBLIC_MODEL_GROUP)) + ) + + models = _extract_models_from_managed_resource_id(output_file_id, "file_id", None) + assert models == [self._PUBLIC_MODEL_GROUP] + assert ( + await can_key_call_model( + model=models[0], + llm_model_list=None, + valid_token=UserAPIKeyAuth( + api_key="sk-test", models=[self._PUBLIC_MODEL_GROUP] + ), + llm_router=None, + ) + is True + ) + + @pytest.mark.asyncio + async def test_falls_back_to_deployment_model_group_without_managed_input_file(self): + from litellm.proxy.openai_files_endpoints.common_utils import ( + _is_base64_encoded_unified_file_id, + get_models_from_unified_file_id, + ) + + output_file_id = await self._run(self._job("file-raw-provider-input")) + + decoded = _is_base64_encoded_unified_file_id(output_file_id) + assert get_models_from_unified_file_id(decoded) == [self._PUBLIC_MODEL_GROUP] From d640ace6d88a08d4e1bef6dc68e9fbce54d6c035 Mon Sep 17 00:00:00 2001 From: Shivam Rawat Date: Fri, 31 Jul 2026 18:05:50 -0700 Subject: [PATCH 02/40] fix(rate-limit): enforce token limits when the pre-call increment is zero The atomic check-and-increment path skipped any counter whose increment was <= 0. The dynamic rate limiter always passes a zero token increment pre-call because usage lands on the counters post-response, so on a model configured with only tpm the limiter evaluated no counters at all: no model-wide TPM cap and no priority reservation, in either generous or strict mode. Regressed in dd57ae6691 when the pre-call flow moved off the read-only should_rate_limit check, which did evaluate token limits. Keep zero-increment counters in the payload so they act as a pure check (current + 0 > limit), matching the pre-regression semantics in both the Lua and in-memory paths. Adds unit regressions at the primitive and hook level plus a live e2e covering the priority_generous/priority_strict registry rows. --- .../hooks/parallel_request_limiter_v3.py | 2 +- tests/e2e/models.py | 3 + .../test_dynamic_rate_limit_priority_e2e.py | 242 ++++++++++++++++++ .../hooks/test_dynamic_rate_limiter_v3.py | 92 +++++++ .../hooks/test_parallel_request_limiter_v3.py | 57 +++++ 5 files changed, 395 insertions(+), 1 deletion(-) create mode 100644 tests/e2e/quota_management/ratelimit/test_dynamic_rate_limit_priority_e2e.py diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index b04ef5f7087..486e3cf88a4 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -1341,7 +1341,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): else: limit_value = rate_limit.get("tokens_per_unit") inc_amount = int(increment_amounts.get("tokens", 0) or 0) - if limit_value is None or inc_amount <= 0: + if limit_value is None: continue counter_key = self.create_rate_limit_keys(descriptor_key, descriptor_value, rlt) # Counter-key TTL and window_size are conceptually distinct diff --git a/tests/e2e/models.py b/tests/e2e/models.py index af695acaa5e..f1c0ede0e85 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -43,6 +43,7 @@ class KeyLoggingCallback(BaseModel): class KeyMetadata(BaseModel): logging: list[KeyLoggingCallback] | None = None + priority: str | None = None class ObjectPermission(BaseModel): @@ -97,6 +98,7 @@ class LiteLLMBudgetTable(BaseModel): class KeyInfo(BaseModel): key_alias: str | None = None + metadata: KeyMetadata | None = None models: list[str] = [] tpm_limit: int | None = None rpm_limit: int | None = None @@ -694,6 +696,7 @@ class LiteLLMParamsBody(BaseModel): complexity_router_config: dict[str, object] | None = None mock_response: str | None = None timeout: float | None = None + tpm: int | None = None ModelMode = Literal["batch", "realtime", "image_generation"] diff --git a/tests/e2e/quota_management/ratelimit/test_dynamic_rate_limit_priority_e2e.py b/tests/e2e/quota_management/ratelimit/test_dynamic_rate_limit_priority_e2e.py new file mode 100644 index 00000000000..06b4aa1a0b5 --- /dev/null +++ b/tests/e2e/quota_management/ratelimit/test_dynamic_rate_limit_priority_e2e.py @@ -0,0 +1,242 @@ +"""Live e2e: the v3 dynamic rate limiter's saturation-aware priority reservation. + +Covers quota_management.ratelimit.priority_generous / priority_strict: with +`dynamic_rate_limiter_v3` enabled, a model's TPM capacity is split into priority +reservations, but a reservation is only enforced once the model is saturated. + +- Generous mode (recorded usage below the saturation threshold): a key whose + priority reserves 25% of capacity keeps serving past its reservation, + borrowing the idle capacity (priority_generous.picks_under_tpm) +- Strict mode (recorded usage at/over the threshold): the over-reservation key + is blocked with the priority-flavored 429 while a key of a different priority, + still inside its own reservation, is served (priority_strict.picks_under_tpm) + +The proxy under test must run with this config (and LITELLM_LICENSE set, since +priority reservation is a premium feature): + + litellm_settings: + callbacks: ["dynamic_rate_limiter_v3"] + priority_reservation: + prod: 0.5 + dev: 0.25 + priority_reservation_settings: + saturation_threshold: 0.5 + saturation_check_cache_ttl: 1 + +The constants below mirror those values; if the proxy runs different ones the +tests fail with a message naming the required config rather than skipping. + +The limiter counts a request against the model-wide window pre-call, but tokens +only land on the counters after each response completes (there is no pre-call +token reservation at the model level), so recorded saturation always trails the +traffic that produced it. The tests therefore drive spend by summing each +body's usage.total_tokens (the counter can never be ahead of that sum) and poll +for the strict-mode block instead of expecting it on an exact call. Each test +creates its own /model/new deployment so its 60s rate-limit window and counters +are isolated from concurrent runs. +""" + +from __future__ import annotations + +import time +from dataclasses import dataclass + +import pytest +from pydantic import BaseModel, ConfigDict, ValidationError + +from e2e_config import unique_marker +from e2e_http import StreamingResponse, require_successful_call +from lifecycle import ResourceManager +from models import KeyGenerateBody, KeyMetadata, LiteLLMParamsBody +from quota_client import QuotaClient + +pytestmark = pytest.mark.e2e + +BACKEND = "anthropic/claude-haiku-4-5-20251001" +MODEL_TPM = 400 +DEV_PRIORITY = "dev" +PROD_PRIORITY = "prod" +DEV_RESERVED_TOKENS = int(MODEL_TPM * 0.25) +SATURATION_TOKENS = int(MODEL_TPM * 0.5) +CHAT_MAX_TOKENS = 16 +WINDOW_SECONDS = 60 +WINDOW_MARGIN_SECONDS = 10 +STRICT_POLL_SPEND_CEILING = int(MODEL_TPM * 0.7) + +REQUIRED_CONFIG_HINT = ( + "the proxy must run litellm_settings.callbacks=['dynamic_rate_limiter_v3'] with " + "priority_reservation {prod: 0.5, dev: 0.25} and priority_reservation_settings " + "{saturation_threshold: 0.5, saturation_check_cache_ttl: 1}; see this module's docstring" +) + + +class _ChatUsage(BaseModel): + model_config = ConfigDict(extra="ignore") + + total_tokens: int + + +class _ChatBodyWithUsage(BaseModel): + model_config = ConfigDict(extra="ignore") + + usage: _ChatUsage + + +def _total_tokens(outcome: StreamingResponse) -> int: + try: + return _ChatBodyWithUsage.model_validate_json(outcome.body).usage.total_tokens + except ValidationError: + pytest.fail(f"successful chat body must report usage.total_tokens, got: {outcome.body[:300]}") + + +@dataclass(frozen=True, slots=True) +class _Fixture: + model: str + dev_key: str + prod_key: str + + +def _dynamic_limited_model(client: QuotaClient, resources: ResourceManager, label: str) -> _Fixture: + model = f"e2e-dynpri-{label}-{unique_marker()}" + model_id = client.proxy.create_model( + model, + LiteLLMParamsBody(model=BACKEND, api_key="os.environ/ANTHROPIC_API_KEY", tpm=MODEL_TPM), + ) + resources.defer(lambda: client.proxy.delete_model(model_id)) + + def _priority_key(priority: str) -> str: + key = client.proxy.generate_key( + KeyGenerateBody( + models=[model], + metadata=KeyMetadata(priority=priority), + key_alias=f"e2e-dynpri-{label}-{priority}-{unique_marker()}", + ) + ) + resources.defer(lambda: client.proxy.delete_key(key)) + return key + + return _Fixture(model=model, dev_key=_priority_key(DEV_PRIORITY), prod_key=_priority_key(PROD_PRIORITY)) + + +def _chat(client: QuotaClient, key: str, model: str) -> StreamingResponse: + return client.chat(key, model, f"reply with one word {unique_marker()}", max_tokens=CHAT_MAX_TOKENS) + + +@dataclass(frozen=True, slots=True) +class _FirstOk: + sent_at: float + response: StreamingResponse + + +def _first_ok(client: QuotaClient, key: str, model: str) -> _FirstOk: + """First successful call on a fresh key opens the model's rate-limit window; + `sent_at` (captured before the winning send) is a lower bound on the window + start. A fresh key may briefly 401 until the auth cache picks it up, so + retry 401s to a deadline; a 401 never reaches the limiter, so only the + successful call consumes budget.""" + deadline = time.monotonic() + client.proxy.poll_timeout + while True: + sent_at = time.monotonic() + outcome = _chat(client, key, model) + if outcome.ok: + return _FirstOk(sent_at=sent_at, response=outcome) + if outcome.status_code != 401 or time.monotonic() >= deadline: + require_successful_call(outcome) + time.sleep(client.proxy.poll_interval) + + +def _window_guard(first: _FirstOk, spent: int) -> None: + assert time.monotonic() < first.sent_at + WINDOW_SECONDS - WINDOW_MARGIN_SECONDS, ( + f"only {spent} tokens of spend landed before the {WINDOW_SECONDS}s rate-limit window could " + "roll; this test needs every call inside one window" + ) + + +class TestDynamicRateLimitPriority: + @pytest.mark.covers( + "quota_management.ratelimit.priority_generous.picks_under_tpm", + exercised_on=["chat_completions"], + ) + def test_generous_mode_lets_priority_borrow_past_reservation( + self, client: QuotaClient, resources: ResourceManager + ) -> None: + fixture = _dynamic_limited_model(client, resources, "generous") + + info = client.proxy.key_info(fixture.dev_key) + assert info.metadata is not None and info.metadata.priority == DEV_PRIORITY, ( + f"/key/info must echo the key's priority metadata, got {info.metadata}" + ) + + first = _first_ok(client, fixture.dev_key, fixture.model) + spent = _total_tokens(first.response) + while spent <= DEV_RESERVED_TOKENS: + _window_guard(first, spent) + assert spent < SATURATION_TOKENS, ( + f"spend reached the saturation threshold ({spent} of {SATURATION_TOKENS}) before " + f"crossing the dev reservation ({DEV_RESERVED_TOKENS}); shrink per-call spend to " + "keep the borrowing claim observable" + ) + outcome = _chat(client, fixture.dev_key, fixture.model) + assert outcome.status_code != 429, ( + f"dev key was blocked at {spent} recorded tokens, under the saturation threshold " + f"({SATURATION_TOKENS} of {MODEL_TPM}); generous mode must let it borrow past its " + f"{DEV_RESERVED_TOKENS}-token reservation. If the limiter is missing entirely, " + f"{REQUIRED_CONFIG_HINT}. 429 body: {outcome.body[:300]}" + ) + require_successful_call(outcome) + spent += _total_tokens(outcome) + + assert spent > DEV_RESERVED_TOKENS + + @pytest.mark.covers( + "quota_management.ratelimit.priority_strict.picks_under_tpm", + exercised_on=["chat_completions"], + ) + def test_strict_mode_blocks_saturated_priority_but_serves_the_other( + self, client: QuotaClient, resources: ResourceManager + ) -> None: + fixture = _dynamic_limited_model(client, resources, "strict") + + prod_warmup = _first_ok(client, fixture.prod_key, fixture.model) + first = _first_ok(client, fixture.dev_key, fixture.model) + prod_spent = _total_tokens(prod_warmup.response) + dev_spent = _total_tokens(first.response) + + while prod_spent + dev_spent < SATURATION_TOKENS: + _window_guard(prod_warmup, prod_spent + dev_spent) + outcome = _chat(client, fixture.dev_key, fixture.model) + assert outcome.status_code != 429, ( + f"dev key was blocked at {prod_spent + dev_spent} recorded tokens, before the " + f"saturation threshold ({SATURATION_TOKENS} of {MODEL_TPM}); strict enforcement " + f"must not engage early. 429 body: {outcome.body[:300]}" + ) + require_successful_call(outcome) + dev_spent += _total_tokens(outcome) + + while True: + _window_guard(prod_warmup, prod_spent + dev_spent) + assert prod_spent + dev_spent < STRICT_POLL_SPEND_CEILING, ( + f"dev key was still served at {prod_spent + dev_spent} tokens, past the saturation " + f"threshold ({SATURATION_TOKENS}) and {DEV_RESERVED_TOKENS}-token dev reservation; " + f"strict priority enforcement never engaged. Check that {REQUIRED_CONFIG_HINT}" + ) + outcome = _chat(client, fixture.dev_key, fixture.model) + if outcome.status_code == 429: + assert "Priority-based rate limit exceeded" in outcome.body, ( + f"the saturated dev key must get the priority-flavored 429, got: {outcome.body[:300]}" + ) + assert outcome.headers.get("x-litellm-priority") == DEV_PRIORITY, ( + f"the 429 must attribute the blocked priority, headers: " + f"{ {k: v for k, v in outcome.headers.items() if 'litellm' in k} }" + ) + break + require_successful_call(outcome) + dev_spent += _total_tokens(outcome) + + prod_outcome = _chat(client, fixture.prod_key, fixture.model) + require_successful_call(prod_outcome) + prod_spent += _total_tokens(prod_outcome) + assert prod_spent < int(MODEL_TPM * 0.5), ( + f"the prod fairness claim needs prod spend ({prod_spent}) inside its reservation " + f"({int(MODEL_TPM * 0.5)}); shrink per-call spend" + ) diff --git a/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py index c8176ca6337..6c717d6f71c 100644 --- a/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter_v3.py @@ -1771,3 +1771,95 @@ async def test_priority_429_includes_model_name_and_configured_limits(): assert "Priority: prod" in error_msg, error_msg assert "Rate limit type: tokens" in error_msg, error_msg assert "Model saturation:" in error_msg, error_msg + + +@pytest.mark.asyncio +async def test_tpm_only_model_enforces_priority_and_model_capacity(): + """Regression: a model configured with ONLY tpm (no rpm) must still be + rate limited. + + The atomic check-and-increment path used to drop any counter whose + pre-call increment was zero. Token increments are always zero pre-call + (usage lands on the counters post-call), so on a TPM-only model the + limiter evaluated no counters at all: no model-wide cap, no priority + reservation, in either mode. This test drives the real pre-call -> + log-success -> pre-call flow with no limiter internals mocked. + """ + from fastapi import HTTPException + + from litellm.types.utils import ModelResponse, Usage + + os.environ["LITELLM_LICENSE"] = "test-license-key" + litellm.priority_reservation = {"dev": 0.25, "prod": 0.5} + + dual_cache = DualCache() + handler = DynamicRateLimitHandler(internal_usage_cache=dual_cache) + + model = "tpm-only-model" + llm_router = Router( + model_list=[ + { + "model_name": model, + "litellm_params": { + "model": "gpt-3.5-turbo", + "api_key": "test-key", + "api_base": "test-base", + "tpm": 400, + }, + } + ] + ) + handler.update_variables(llm_router=llm_router) + + dev_user = UserAPIKeyAuth() + dev_user.metadata = {"priority": "dev"} + prod_user = UserAPIKeyAuth() + prod_user.metadata = {"priority": "prod"} + + async def record_usage(priority: str, total_tokens: int) -> None: + await handler.async_log_success_event( + kwargs={ + "standard_logging_object": { + "metadata": {"user_api_key_auth_metadata": {"priority": priority}}, + }, + "litellm_params": {"metadata": {"model_group": model}}, + }, + response_obj=ModelResponse( + model=model, + usage=Usage(prompt_tokens=0, completion_tokens=total_tokens, total_tokens=total_tokens), + ), + start_time=None, + end_time=None, + ) + + assert ( + await handler.async_pre_call_hook( + user_api_key_dict=dev_user, cache=dual_cache, data={"model": model}, call_type="completion" + ) + is None + ) + + await record_usage("dev", 250) + + with pytest.raises(HTTPException) as dev_blocked: + await handler.async_pre_call_hook( + user_api_key_dict=dev_user, cache=dual_cache, data={"model": model}, call_type="completion" + ) + assert dev_blocked.value.status_code == 429 + assert "Priority-based rate limit exceeded" in dev_blocked.value.detail["error"] + + assert ( + await handler.async_pre_call_hook( + user_api_key_dict=prod_user, cache=dual_cache, data={"model": model}, call_type="completion" + ) + is None + ) + + await record_usage("prod", 200) + + with pytest.raises(HTTPException) as capacity_blocked: + await handler.async_pre_call_hook( + user_api_key_dict=prod_user, cache=dual_cache, data={"model": model}, call_type="completion" + ) + assert capacity_blocked.value.status_code == 429 + assert "Model capacity reached" in capacity_blocked.value.detail["error"] diff --git a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py index 56bfd1829b5..84e41711176 100644 --- a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py @@ -4998,3 +4998,60 @@ async def test_split_usage_still_respects_the_configured_limit_type(monkeypatch) token_operations = [op for op in captured_operations if op["key"].endswith(":tokens")] assert token_operations assert all(op["increment_value"] == 7 for op in token_operations) + + +@pytest.mark.asyncio +async def test_atomic_check_with_zero_increment_still_enforces_token_limit(): + """Regression: a zero token increment must still CHECK the token limit. + + The dynamic rate limiter calls atomic_check_and_increment_by_n with + {"requests": 1, "tokens": 0} because tokens land on the counter post-call. + The payload builder used to skip any counter whose increment was <= 0, so a + TPM-only descriptor produced zero counters to evaluate and the call + returned OK with empty statuses; TPM limits were never enforced at all. + """ + from litellm.proxy.hooks.parallel_request_limiter_v3 import RateLimitDescriptor + + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(DualCache()) + ) + descriptor = RateLimitDescriptor( + key="model_saturation_check", + value="tpm-only-model", + rate_limit={"tokens_per_unit": 100, "window_size": 60}, + ) + zero_token_increment: Dict[str, int] = {"requests": 1, "tokens": 0} + + under_limit = await handler.atomic_check_and_increment_by_n( + descriptors=[descriptor], + increments=[zero_token_increment], + ) + assert under_limit["overall_code"] == "OK" + assert [s["rate_limit_type"] for s in under_limit["statuses"]] == ["tokens"] + + counter_key = handler.create_rate_limit_keys( + "model_saturation_check", "tpm-only-model", "tokens" + ) + await handler.async_increment_tokens_with_ttl_preservation( + pipeline_operations=[ + RedisPipelineIncrementOperation( + key=counter_key, increment_value=150, ttl=60 + ) + ], + ) + + over_limit = await handler.atomic_check_and_increment_by_n( + descriptors=[descriptor], + increments=[zero_token_increment], + ) + assert over_limit["overall_code"] == "OVER_LIMIT" + blocked = over_limit["statuses"][0] + assert blocked["rate_limit_type"] == "tokens" + assert blocked["current_limit"] == 100 + + assert ( + await handler.internal_usage_cache.async_get_cache( + key=counter_key, litellm_parent_otel_span=None, local_only=True + ) + == 150 + ) From 836bd927b68a7fbabfb1afcdcbdc6b31dd13b8fc Mon Sep 17 00:00:00 2001 From: Shivam Rawat Date: Fri, 31 Jul 2026 18:05:50 -0700 Subject: [PATCH 03/40] fix(rate-limit): skip negative increments in the atomic payload builder Review feedback: the relaxed predicate admitted negative increments, which both atomic backends would apply as decrements. Restrict the new behavior to zero-valued pure checks and assert negatives neither check nor mutate counters. --- litellm/proxy/hooks/parallel_request_limiter_v3.py | 2 +- .../hooks/test_parallel_request_limiter_v3.py | 14 ++++++++++++++ 2 files changed, 15 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 486e3cf88a4..89559ff72ae 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -1341,7 +1341,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): else: limit_value = rate_limit.get("tokens_per_unit") inc_amount = int(increment_amounts.get("tokens", 0) or 0) - if limit_value is None: + if limit_value is None or inc_amount < 0: continue counter_key = self.create_rate_limit_keys(descriptor_key, descriptor_value, rlt) # Counter-key TTL and window_size are conceptually distinct diff --git a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py index 84e41711176..7b3f00a55a8 100644 --- a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py @@ -5055,3 +5055,17 @@ async def test_atomic_check_with_zero_increment_still_enforces_token_limit(): ) == 150 ) + + negative_increment: Dict[str, int] = {"requests": -1, "tokens": -50} + refund_attempt = await handler.atomic_check_and_increment_by_n( + descriptors=[descriptor], + increments=[negative_increment], + ) + assert refund_attempt["overall_code"] == "OK" + assert refund_attempt["statuses"] == [] + assert ( + await handler.internal_usage_cache.async_get_cache( + key=counter_key, litellm_parent_otel_span=None, local_only=True + ) + == 150 + ) From e8e2e07ef67010685839f79c0b0ad18dd9ec7ced Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 31 Jul 2026 19:57:52 -0700 Subject: [PATCH 04/40] fix(proxy): align team member add with existing user provisioning rules Adding a team member by a user_id with no user row created that row as a side effect for any caller permitted to add members, while creating users directly is restricted to proxy admins. Restrict that path to proxy admins too; adding an existing user, and inviting a new one by user_email (where the user_id is allocated server-side), are unchanged. Also record the membership change, and any user row it creates, in the audit log, matching /team/update, /user/new and /key/*. --- .../management_endpoints/team_endpoints.py | 147 +++++++++++ .../test_team_endpoints.py | 236 ++++++++++++++++++ 2 files changed, 383 insertions(+) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index fa01f43d049..7b5ce33cb8b 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -13,6 +13,7 @@ import asyncio import json import math import traceback +from collections.abc import Sequence from datetime import datetime, timezone from typing import Annotated, Any, Dict, List, Mapping, Optional, Tuple, Union, cast @@ -2426,6 +2427,123 @@ def _emit_team_members_metric(team: LiteLLM_TeamTable) -> None: verbose_proxy_logger.debug("Prometheus: failed to emit team members metric: %s", str(e)) +async def _resolve_existing_member_user_ids( + members: Sequence[Member], + prisma_client: PrismaClient, +) -> frozenset[str]: + """Return the caller-supplied user_ids that already have a user row.""" + user_repository = UserRepository(prisma_client) + found = await asyncio.gather( + *(user_repository.find_by_id(member.user_id) for member in members if member.user_id is not None) + ) + return frozenset(user.user_id for user in found if user is not None and user.user_id is not None) + + +def _pre_existing_user_ids( + members: Sequence[Member], + caller_supplied_user_ids: frozenset[str], + existing_user_ids: frozenset[str], +) -> frozenset[str]: + """Return the user_ids that already had a user row before this request. + + Combines the caller-supplied ids that resolved to a user with the ids + ``_validate_and_populate_member_user_info`` filled in, which it only does + from a matched user row. Deriving it that way keeps this in step with the + email matching that resolution performs, rather than repeating it here. + """ + populated_user_ids = frozenset( + member.user_id + for member in members + if member.user_id is not None and member.user_id not in caller_supplied_user_ids + ) + return existing_user_ids | populated_user_ids + + +def _validate_member_user_id_provisioning( + members: Sequence[Member], + existing_user_ids: frozenset[str], + user_api_key_dict: UserAPIKeyAuth, +) -> None: + """Restrict adding a caller-chosen user_id that has no user row yet to proxy admins. + + Team and org admins keep the ability to add users that already exist and to + invite new ones by user_email, where the user_id is allocated server-side. + """ + if user_api_key_dict.user_role in ( + LitellmUserRoles.PROXY_ADMIN, + LitellmUserRoles.PROXY_ADMIN.value, + ): + return + + unknown_user_ids = tuple( + member.user_id for member in members if member.user_id is not None and member.user_id not in existing_user_ids + ) + if not unknown_user_ids: + return + + raise HTTPException( + status_code=403, + detail={ # mutable-ok: HTTPException detail must be a plain mapping to keep this route's {"error": ...} response shape + "error": ( + "Only proxy admins can add a user_id that does not exist yet: {}. " + "Add the member by user_email to invite a new user, or ask a proxy admin " + "to create the user first.".format(", ".join(unknown_user_ids)) + ) + }, + ) + + +def _members_audit_value(members: Sequence[Member]) -> str: + """Serialize a team's member list for an audit-log value. + + The audit-log columns hold a JSON object, so the member list is nested + under a key rather than serialized as a top-level array. + """ + return safe_dumps( + { # mutable-ok: the audit-log JSON column rejects a top-level array, so this value must be an object + "members_with_roles": tuple(member.model_dump() for member in members) + } + ) + + +async def _create_team_member_add_audit_logs( + team_id: str, + updated_users: Sequence[LiteLLM_UserTable], + existing_user_ids: frozenset[str], + before_members: Sequence[Member], + after_members: Sequence[Member], + user_api_key_dict: UserAPIKeyAuth, + litellm_proxy_admin_name: str, +) -> None: + """Record the membership change, and any user row it created, in the audit log.""" + from litellm.proxy.management_helpers.audit_logs import create_object_audit_log + + for user in updated_users: + if user.user_id is None or user.user_id in existing_user_ids: + continue + await create_object_audit_log( + object_id=user.user_id, + action="created", + litellm_changed_by=None, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=litellm_proxy_admin_name, + table_name=LitellmTableNames.USER_TABLE_NAME, + before_value=None, + after_value=safe_dumps(user.model_dump(exclude_none=True)), + ) + + await create_object_audit_log( + object_id=team_id, + action="updated", + litellm_changed_by=None, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=litellm_proxy_admin_name, + table_name=LitellmTableNames.TEAM_TABLE_NAME, + before_value=_members_audit_value(before_members), + after_value=_members_audit_value(after_members), + ) + + async def _validate_and_populate_member_user_info( member: Member, prisma_client: PrismaClient, @@ -2606,6 +2724,19 @@ async def team_member_add( data=data, ) + requested_members = tuple(data.member) if isinstance(data.member, list) else (data.member,) + caller_supplied_user_ids = frozenset(member.user_id for member in requested_members if member.user_id is not None) + existing_user_ids = await _resolve_existing_member_user_ids( + members=requested_members, + prisma_client=prisma_client, + ) + _validate_member_user_id_provisioning( + members=requested_members, + existing_user_ids=existing_user_ids, + user_api_key_dict=user_api_key_dict, + ) + members_before_add = tuple(complete_team_data.members_with_roles) + # Validate and populate user_email/user_id for members before processing if isinstance(data.member, Member): await _validate_and_populate_member_user_info( @@ -2619,6 +2750,12 @@ async def team_member_add( prisma_client=prisma_client, ) + pre_existing_user_ids = _pre_existing_user_ids( + members=requested_members, + caller_supplied_user_ids=caller_supplied_user_ids, + existing_user_ids=existing_user_ids, + ) + ( updated_team, updated_users, @@ -2637,6 +2774,16 @@ async def team_member_add( _emit_team_members_metric(complete_team_data) + await _create_team_member_add_audit_logs( + team_id=data.team_id, + updated_users=updated_users, + existing_user_ids=pre_existing_user_ids, + before_members=members_before_add, + after_members=tuple(complete_team_data.members_with_roles), + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=litellm_proxy_admin_name, + ) + return TeamAddMemberResponse.model_validate( { **updated_team.model_dump(), diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index 25a0ff644f8..b1438447e4f 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -10338,3 +10338,239 @@ async def test_list_available_teams_filters_joined_and_validates_rows(monkeypatc assert result[0].team_alias == "open team" find_many_kwargs = mock_prisma_client.db.litellm_teamtable.find_many.call_args.kwargs assert find_many_kwargs["where"] == {"team_id": {"in": ["team-open"]}} + + +def _provisioning_caller(role: LitellmUserRoles) -> UserAPIKeyAuth: + return UserAPIKeyAuth(user_id="caller-1", user_role=role) + + +def test_validate_member_user_id_provisioning_allows_proxy_admin(): + """Proxy admins may add a user_id that has no user row yet.""" + from litellm.proxy.management_endpoints.team_endpoints import ( + _validate_member_user_id_provisioning, + ) + + _validate_member_user_id_provisioning( + members=[Member(user_id="brand-new", role="user")], + existing_user_ids=frozenset(), + user_api_key_dict=_provisioning_caller(LitellmUserRoles.PROXY_ADMIN), + ) + + +def test_validate_member_user_id_provisioning_rejects_unknown_user_id_for_non_proxy_admin(): + """A non-proxy-admin cannot add a user_id that has no user row yet.""" + from litellm.proxy.management_endpoints.team_endpoints import ( + _validate_member_user_id_provisioning, + ) + + with pytest.raises(HTTPException) as exc_info: + _validate_member_user_id_provisioning( + members=[Member(user_id="brand-new", role="user")], + existing_user_ids=frozenset(), + user_api_key_dict=_provisioning_caller(LitellmUserRoles.INTERNAL_USER), + ) + + assert exc_info.value.status_code == 403 + assert "brand-new" in str(exc_info.value.detail) + + +def test_validate_member_user_id_provisioning_allows_existing_user_id_for_non_proxy_admin(): + """A non-proxy-admin may still add a user that already exists.""" + from litellm.proxy.management_endpoints.team_endpoints import ( + _validate_member_user_id_provisioning, + ) + + _validate_member_user_id_provisioning( + members=[Member(user_id="already-here", role="user")], + existing_user_ids=frozenset({"already-here"}), + user_api_key_dict=_provisioning_caller(LitellmUserRoles.INTERNAL_USER), + ) + + +def test_validate_member_user_id_provisioning_allows_email_only_member_for_non_proxy_admin(): + """Inviting by user_email stays open to non-proxy-admins; the user_id is server-allocated.""" + from litellm.proxy.management_endpoints.team_endpoints import ( + _validate_member_user_id_provisioning, + ) + + _validate_member_user_id_provisioning( + members=[Member(user_email="invitee@example.com", role="user")], + existing_user_ids=frozenset(), + user_api_key_dict=_provisioning_caller(LitellmUserRoles.INTERNAL_USER), + ) + + +def test_validate_member_user_id_provisioning_rejects_unknown_user_id_paired_with_email(): + """Supplying a user_email alongside an unknown user_id does not lift the restriction.""" + from litellm.proxy.management_endpoints.team_endpoints import ( + _validate_member_user_id_provisioning, + ) + + with pytest.raises(HTTPException) as exc_info: + _validate_member_user_id_provisioning( + members=[Member(user_id="chosen-id", user_email="invitee@example.com", role="user")], + existing_user_ids=frozenset(), + user_api_key_dict=_provisioning_caller(LitellmUserRoles.INTERNAL_USER), + ) + + assert exc_info.value.status_code == 403 + + +def test_validate_member_user_id_provisioning_reports_every_unknown_member(): + """A bulk add names each unknown user_id rather than only the first.""" + from litellm.proxy.management_endpoints.team_endpoints import ( + _validate_member_user_id_provisioning, + ) + + with pytest.raises(HTTPException) as exc_info: + _validate_member_user_id_provisioning( + members=[ + Member(user_id="known", role="user"), + Member(user_id="unknown-a", role="user"), + Member(user_id="unknown-b", role="user"), + ], + existing_user_ids=frozenset({"known"}), + user_api_key_dict=_provisioning_caller(LitellmUserRoles.INTERNAL_USER), + ) + + detail = str(exc_info.value.detail) + assert "unknown-a" in detail + assert "unknown-b" in detail + + +@pytest.mark.asyncio +async def test_resolve_existing_member_user_ids_matches_caller_supplied_user_ids(): + """Only caller-supplied user_ids are looked up; unknown ones resolve to nothing.""" + from litellm.proxy.management_endpoints.team_endpoints import ( + _resolve_existing_member_user_ids, + ) + + prisma_client = MagicMock() + + async def find_by_id(user_id): + if user_id == "by-id": + return LiteLLM_UserTable(user_id="by-id", max_budget=None, spend=0.0, user_email=None, models=[]) + return None + + with patch("litellm.proxy.management_endpoints.team_endpoints.UserRepository") as repo: + repo.return_value.find_by_id = AsyncMock(side_effect=find_by_id) + + resolved = await _resolve_existing_member_user_ids( + members=[ + Member(user_id="by-id", role="user"), + Member(user_id="missing", role="user"), + Member(user_email="someone@example.com", role="user"), + ], + prisma_client=prisma_client, + ) + + assert resolved == frozenset({"by-id"}) + + +def test_pre_existing_user_ids_counts_ids_filled_in_by_member_resolution(): + """An id the member-resolution step filled in came from a matched row, so it pre-existed. + + This is what keeps a case-variant email invite of an existing user from being + recorded as a newly created user. + """ + from litellm.proxy.management_endpoints.team_endpoints import _pre_existing_user_ids + + # member arrived email-only; resolution matched an existing row and filled in the id + resolved_member = Member(user_id="matched-existing", user_email="Someone@Example.com", role="user") + + assert _pre_existing_user_ids( + members=[resolved_member], + caller_supplied_user_ids=frozenset(), + existing_user_ids=frozenset(), + ) == frozenset({"matched-existing"}) + + +def test_pre_existing_user_ids_excludes_caller_supplied_ids_that_do_not_exist(): + """A caller-supplied id that resolved to nothing is genuinely new, so it stays out.""" + from litellm.proxy.management_endpoints.team_endpoints import _pre_existing_user_ids + + assert _pre_existing_user_ids( + members=[Member(user_id="brand-new", role="user"), Member(user_id="already-here", role="user")], + caller_supplied_user_ids=frozenset({"brand-new", "already-here"}), + existing_user_ids=frozenset({"already-here"}), + ) == frozenset({"already-here"}) + + +def test_members_audit_value_serializes_to_a_json_object(): + """The audit-log columns hold a JSON object; a top-level array is rejected by the DB.""" + from litellm.proxy.management_endpoints.team_endpoints import _members_audit_value + + payload = json.loads(_members_audit_value([Member(user_id="u1", role="admin"), Member(user_id="u2", role="user")])) + + assert isinstance(payload, dict) + assert [m["user_id"] for m in payload["members_with_roles"]] == ["u1", "u2"] + + +@pytest.mark.asyncio +async def test_team_member_add_audits_a_user_created_from_a_list_payload(monkeypatch): + """A user created by a list payload must still be reported as newly created. + + For a list payload the member-list reconciliation back-fills the caller's own + Member objects with the ids of users this request just created. The set of + pre-existing ids therefore has to be captured before that runs, otherwise a + freshly created user looks like it was already there and no creation is recorded. + """ + from litellm.proxy._types import TeamMemberAddRequest + from litellm.proxy.management_endpoints.team_endpoints import team_member_add + + team_id = "team-list-audit" + created_user_id = "generated-uuid-for-new-invitee" + member = Member(user_email="invitee@example.com", role="user") + + mock_prisma_client = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + monkeypatch.setattr("litellm.proxy.proxy_server.litellm_proxy_admin_name", "default_user_id") + + team_row = LiteLLM_TeamTable(team_id=team_id, members_with_roles=[]) + created_user = LiteLLM_UserTable( + user_id=created_user_id, user_email="invitee@example.com", max_budget=None, spend=0.0, models=[] + ) + updated_team = MagicMock() + updated_team.model_dump.return_value = {"team_id": team_id, "members_with_roles": []} + + async def fake_add_team_members_to_team(**kwargs): + # mirrors _update_team_members_list: the list branch mutates the caller's Member in place + member.user_id = created_user_id + return updated_team, [created_user], [] + + with ( + patch( + "litellm.proxy.management_endpoints.team_endpoints.get_team_object", + new_callable=AsyncMock, + return_value=team_row, + ), + patch( + "litellm.proxy.management_endpoints.team_endpoints._validate_team_member_add_permissions", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.management_endpoints.team_endpoints._validate_and_populate_member_user_info", + new_callable=AsyncMock, + ), + patch( + "litellm.proxy.management_endpoints.team_endpoints._resolve_existing_member_user_ids", + new_callable=AsyncMock, + return_value=frozenset(), + ), + patch( + "litellm.proxy.management_endpoints.team_endpoints._add_team_members_to_team", + side_effect=fake_add_team_members_to_team, + ), + patch( + "litellm.proxy.management_endpoints.team_endpoints._create_team_member_add_audit_logs", + new_callable=AsyncMock, + ) as mock_audit, + ): + await team_member_add( + data=TeamMemberAddRequest(team_id=team_id, member=[member]), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin-1"), + ) + + mock_audit.assert_called_once() + assert created_user_id not in mock_audit.call_args.kwargs["existing_user_ids"] From 629d58443ea77bd6156db09f0dfa1280acd6fd20 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 31 Jul 2026 20:08:02 -0700 Subject: [PATCH 05/40] feat(proxy): push config sync to pods via redis pub/sub After any management write to a DB-backed config table, publish an invalidation event on the coordination Redis; every pod runs a subscriber that debounces, jitters, and triggers an immediate add_deployment plus get_credentials resync. The interval polls stay as slow reconciliation fallback and behavior without Redis is unchanged since publish and subscribe both no-op. --- .../proxy/common_utils/config_sync_pubsub.py | 258 ++++++++ .../key_management_endpoints.py | 5 + .../model_management_endpoints.py | 7 + litellm/proxy/proxy_server.py | 27 +- .../proxy_setting_endpoints.py | 2 + litellm/proxy/utils.py | 11 +- .../repositories/credentials_repository.py | 6 +- litellm/repositories/model_repository.py | 6 +- litellm/repositories/table_repositories.py | 7 +- ruff.toml | 2 +- .../common_utils/test_config_sync_pubsub.py | 600 ++++++++++++++++++ .../repositories/test_repositories.py | 58 +- 12 files changed, 978 insertions(+), 11 deletions(-) create mode 100644 litellm/proxy/common_utils/config_sync_pubsub.py create mode 100644 tests/test_litellm/proxy/common_utils/test_config_sync_pubsub.py diff --git a/litellm/proxy/common_utils/config_sync_pubsub.py b/litellm/proxy/common_utils/config_sync_pubsub.py new file mode 100644 index 00000000000..4292c67c8c2 --- /dev/null +++ b/litellm/proxy/common_utils/config_sync_pubsub.py @@ -0,0 +1,258 @@ +import asyncio +import json +import random +from collections.abc import Awaitable, Callable +from dataclasses import asdict, dataclass +from typing import TYPE_CHECKING, Protocol, cast # noqa: TID251 # untyped prisma/redis boundary needs cast + +from litellm._logging import verbose_proxy_logger + +if TYPE_CHECKING: + from litellm.caching.redis_cache import RedisCache + + +class _ConfigSyncPubSub(Protocol): + def subscribe(self, *channels: str) -> Awaitable[object]: ... + + def get_message(self, *, ignore_subscribe_messages: bool, timeout: float) -> Awaitable[object]: ... + + def aclose(self) -> Awaitable[object]: ... + + +class _ConfigSyncPubSubClient(Protocol): + def publish(self, channel: str, message: str) -> Awaitable[int]: ... + + def pubsub(self) -> _ConfigSyncPubSub: ... + + +CONFIG_SYNC_CHANNEL = "litellm_proxy.config_change" +CONFIG_SYNC_DEBOUNCE_SECONDS = 1.0 +CONFIG_SYNC_JITTER_MAX_SECONDS = 5.0 +_POLL_TIMEOUT_SECONDS = 1.0 +_BACKOFF_INITIAL_SECONDS = 5.0 +_BACKOFF_MAX_SECONDS = 60.0 + +_WRITE_ACTION_NAMES: frozenset[str] = frozenset( + {"create", "create_many", "update", "update_many", "upsert", "delete", "delete_many"} +) + +_CONFIG_SYNCED_TABLE_NAMES: frozenset[str] = frozenset( + { + "litellm_proxymodeltable", + "litellm_credentialstable", + "litellm_guardrailstable", + "litellm_policytable", + "litellm_policyattachmenttable", + "litellm_managedvectorstorestable", + "litellm_managedvectorstoreindextable", + "litellm_mcpservertable", + "litellm_agentstable", + "litellm_prompttable", + "litellm_searchtoolstable", + "litellm_ssoconfig", + "litellm_cacheconfig", + "litellm_configoverrides", + } +) + + +def coordination_redis_cache() -> "RedisCache | None": + from litellm.proxy.proxy_server import redis_usage_cache + + return redis_usage_cache + + +def config_sync_channel(redis_cache: "RedisCache") -> str: + if redis_cache.namespace is None: + return CONFIG_SYNC_CHANNEL + return f"{redis_cache.namespace}:{CONFIG_SYNC_CHANNEL}" + + +def _raw_async_client(redis_cache: "RedisCache") -> object: + return cast( # cast-ok: redis-py generics leave the client type partially unknown + object, + redis_cache.init_async_client(), # pyright: ignore[reportUnknownMemberType] # redis generics + ) + + +def _pubsub_capable_client(redis_cache: "RedisCache") -> _ConfigSyncPubSubClient | None: + from redis.asyncio import Redis + + client = _raw_async_client(redis_cache) + if isinstance(client, Redis): + return cast(_ConfigSyncPubSubClient, client) # cast-ok: protocol view of the standalone redis client + return None + + +@dataclass(frozen=True, slots=True) +class _ConfigChangeMessage: + object_type: str + + +def _config_change_message_json(object_type: str) -> str: + return json.dumps(asdict(_ConfigChangeMessage(object_type=object_type))) + + +async def publish_config_change(redis_cache: "RedisCache | None", object_type: str) -> None: + if redis_cache is None: + return + try: + client = _pubsub_capable_client(redis_cache) + if client is None: + verbose_proxy_logger.debug( + "config sync publish for %s skipped: cluster redis client has no pub/sub support", + object_type, + ) + return + await client.publish(config_sync_channel(redis_cache), _config_change_message_json(object_type)) + except Exception as e: # noqa: BLE001 # best-effort publish; writes must never fail on redis errors + verbose_proxy_logger.warning("config sync publish for %s failed: %s", object_type, e) + + +async def publish_config_change_for_object_type(object_type: str) -> None: + await publish_config_change(redis_cache=coordination_redis_cache(), object_type=object_type) + + +class _PublishOnWriteActions: + __slots__ = ("_actions", "_object_type", "_publish") + + def __init__(self, actions: object, object_type: str, publish: Callable[[str], Awaitable[None]]) -> None: + self._actions = actions + self._object_type = object_type + self._publish = publish + + def __getattr__(self, name: str) -> object: + attribute = cast(object, getattr(self._actions, name)) # cast-ok: getattr on dynamic prisma actions + if name not in _WRITE_ACTION_NAMES: + return attribute + write_action = cast(Callable[..., Awaitable[object]], attribute) # cast-ok: prisma actions are untyped + object_type = self._object_type + publish = self._publish + + async def _write_then_publish( + *args: object, + **kwargs: object, # kwargs-ok: transparent passthrough to untyped prisma action + ) -> object: + result = await write_action(*args, **kwargs) + await publish(object_type) + return result + + return _write_then_publish + + +def wrap_table_actions_for_config_sync( + actions: object, + table_name: str, + publish: Callable[[str], Awaitable[None]] = publish_config_change_for_object_type, +) -> object: + if table_name not in _CONFIG_SYNCED_TABLE_NAMES: + return actions + return _PublishOnWriteActions(actions=actions, object_type=table_name, publish=publish) + + +class ConfigSyncSubscriber: + __slots__ = ( + "_backoff_initial_seconds", + "_backoff_max_seconds", + "_debounce_seconds", + "_jitter_max_seconds", + "_redis_cache", + "_resync_callbacks", + "_rng", + "_sleep", + "_task", + ) + + def __init__( + self, + redis_cache: "RedisCache", + resync_callbacks: tuple[Callable[[], Awaitable[None]], ...], + debounce_seconds: float = CONFIG_SYNC_DEBOUNCE_SECONDS, + jitter_max_seconds: float = CONFIG_SYNC_JITTER_MAX_SECONDS, + backoff_initial_seconds: float = _BACKOFF_INITIAL_SECONDS, + backoff_max_seconds: float = _BACKOFF_MAX_SECONDS, + rng: random.Random | None = None, + sleep: Callable[[float], Awaitable[None]] = asyncio.sleep, + ) -> None: + self._redis_cache = redis_cache + self._resync_callbacks = resync_callbacks + self._debounce_seconds = debounce_seconds + self._jitter_max_seconds = jitter_max_seconds + self._backoff_initial_seconds = backoff_initial_seconds + self._backoff_max_seconds = backoff_max_seconds + self._rng = rng if rng is not None else random.Random() + self._sleep = sleep + self._task: asyncio.Task[None] | None = None + + def start(self) -> None: + if self._task is not None: + return + self._task = asyncio.create_task(self._run()) + + async def stop(self) -> None: + task = self._task + if task is None: + return + self._task = None + _ = task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + async def _run(self) -> None: + backoff_seconds = self._backoff_initial_seconds + while True: + try: + client = _pubsub_capable_client(self._redis_cache) + if client is None: + verbose_proxy_logger.warning( + "config sync subscriber disabled: cluster redis client has no pub/sub support; " + "interval polling remains the only sync mechanism" + ) + return + pubsub = client.pubsub() + try: + await pubsub.subscribe(config_sync_channel(self._redis_cache)) + backoff_seconds = self._backoff_initial_seconds + await self._consume(pubsub) + finally: + await self._close_pubsub(pubsub) + except asyncio.CancelledError: + raise + except Exception as e: # noqa: BLE001 # any redis failure falls through to backoff and reconnect + verbose_proxy_logger.warning( + "config sync subscriber redis error: %s; reconnecting in %.0fs", + e, + backoff_seconds, + ) + await self._sleep(backoff_seconds) + backoff_seconds = min(backoff_seconds * 2, self._backoff_max_seconds) + + async def _consume(self, pubsub: _ConfigSyncPubSub) -> None: + while True: + message = await pubsub.get_message(ignore_subscribe_messages=True, timeout=_POLL_TIMEOUT_SECONDS) + if message is None: + continue + await self._sleep(self._debounce_seconds + self._rng.uniform(0.0, self._jitter_max_seconds)) + await self._drain_pending(pubsub) + await self._run_resync_callbacks() + + @staticmethod + async def _drain_pending(pubsub: _ConfigSyncPubSub) -> None: + while await pubsub.get_message(ignore_subscribe_messages=True, timeout=0) is not None: + pass + + async def _run_resync_callbacks(self) -> None: + for callback in self._resync_callbacks: + try: + await callback() + except Exception as e: # noqa: BLE001 # one failing resync callback must not kill the subscriber + verbose_proxy_logger.warning("config sync resync callback failed: %s", e) + + @staticmethod + async def _close_pubsub(pubsub: _ConfigSyncPubSub) -> None: + try: + await pubsub.aclose() + except Exception as e: # noqa: BLE001 # best-effort close of a possibly-broken connection + verbose_proxy_logger.debug("config sync pubsub close failed: %s", e) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index a94a75fdfa3..1fbfda8bf8a 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -60,6 +60,10 @@ from litellm.proxy.common_utils.callback_utils import ( decrypt_callback_vars, encrypt_callback_vars, ) +from litellm.proxy.common_utils.config_sync_pubsub import ( + coordination_redis_cache, + publish_config_change, +) from litellm.proxy.common_utils.rbac_utils import check_org_admin_can_generate_keys from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache @@ -4240,6 +4244,7 @@ async def _rotate_master_key( await tx.litellm_proxymodeltable.create_many( data=new_models, ) + await publish_config_change(redis_cache=coordination_redis_cache(), object_type="litellm_proxymodeltable") # 3. process config table try: config = await ConfigRepository(prisma_client).table.find_many() diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index b6422d7f5ae..21395658dac 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -38,6 +38,10 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.config_sync_pubsub import ( + coordination_redis_cache, + publish_config_change, +) from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper from litellm.proxy.management_endpoints.common_utils import _is_user_team_admin from litellm.proxy.management_endpoints.team_endpoints import ( @@ -809,6 +813,9 @@ async def delete_team_models( await tx.litellm_proxymodeltable.delete_many(where={"model_id": {"in": model_ids}}) deleted_model_ids.extend(model_ids) + if deleted_model_ids: + await publish_config_change(redis_cache=coordination_redis_cache(), object_type="litellm_proxymodeltable") + if llm_router is not None: for model_id in deleted_model_ids: llm_router.delete_deployment(id=model_id) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index a60ea2da019..6bb0032b38c 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -301,6 +301,7 @@ from litellm.proxy.common_request_processing import ( create_response, ) from litellm.proxy.common_utils.callback_utils import initialize_callbacks_on_proxy +from litellm.proxy.common_utils.config_sync_pubsub import ConfigSyncSubscriber from litellm.proxy.common_utils.debug_utils import init_verbose_loggers from litellm.proxy.common_utils.debug_utils import router as debugging_endpoints_router from litellm.proxy.common_utils.encrypt_decrypt_utils import ( @@ -546,6 +547,7 @@ from litellm.proxy.utils import ( _get_redoc_url, _is_projected_spend_over_limit, _is_valid_team_configs, + evict_config_param, get_config_param, get_custom_url, get_error_message_str, @@ -1151,6 +1153,12 @@ async def proxy_startup_event(app: FastAPI): except Exception as e: verbose_proxy_logger.error(f"Error stopping DB health watchdog task: {e}") + if proxy_config.config_sync_subscriber is not None: + try: + await proxy_config.config_sync_subscriber.stop() + except Exception as e: + verbose_proxy_logger.error(f"Error stopping config sync subscriber: {e}") + await proxy_shutdown_event() # type: ignore[reportGeneralTypeIssues] @@ -3837,6 +3845,7 @@ class ProxyConfig: self._last_semantic_filter_config: Optional[Dict[str, Any]] = None self._last_hashicorp_vault_config: Optional[Dict[str, Any]] = None self.worker_registry: List["WorkerRegistryEntry"] = [] + self.config_sync_subscriber: ConfigSyncSubscriber | None = None def is_yaml(self, config_file_path: str) -> bool: if not os.path.isfile(config_file_path): @@ -6495,7 +6504,7 @@ class ProxyConfig: }, }, ) - await invalidate_config_param("model_cost_map_reload_config") + await evict_config_param("model_cost_map_reload_config") verbose_proxy_logger.info( f"Model cost map reloaded successfully. Models count: {len(new_model_cost_map) if new_model_cost_map else 0}" @@ -6590,7 +6599,7 @@ class ProxyConfig: }, }, ) - await invalidate_config_param("anthropic_beta_headers_reload_config") + await evict_config_param("anthropic_beta_headers_reload_config") # Count providers in config provider_count = sum(1 for k in new_config.keys() if k != "provider_aliases" and k != "description") @@ -8165,6 +8174,20 @@ class ProxyStartupEvent: ) await proxy_config.get_credentials(prisma_client=prisma_client) + if redis_usage_cache is not None and proxy_config.config_sync_subscriber is None: + + async def _resync_config_from_db() -> None: + await proxy_config.add_deployment(prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj) + + async def _resync_credentials_from_db() -> None: + await proxy_config.get_credentials(prisma_client=prisma_client) + + proxy_config.config_sync_subscriber = ConfigSyncSubscriber( + redis_cache=redis_usage_cache, + resync_callbacks=(_resync_config_from_db, _resync_credentials_from_db), + ) + proxy_config.config_sync_subscriber.start() + if store_model_in_db is not True: await proxy_config.init_mcp_servers_from_db() if prisma_client is not None: diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index 8f3f8ad1bfc..e4872a4b6b5 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -21,6 +21,7 @@ from litellm.proxy.config_resolvers.sso import ( SSO_SECRET_FIELDS, resolve_sso_config, ) +from litellm.proxy.utils import invalidate_config_param from litellm.repositories.config_repository import ConfigRepository from litellm.repositories.table_repositories import ( SSOConfigRepository, @@ -971,6 +972,7 @@ async def update_sso_settings( "param_value": json.dumps(filtered_env_vars, default=str), }, ) + await invalidate_config_param("environment_variables") except Exception as e: raise HTTPException( status_code=500, diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 2ca251a3211..fbcca73e779 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -119,6 +119,10 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.auth.route_checks import RouteChecks +from litellm.proxy.common_utils.config_sync_pubsub import ( + coordination_redis_cache, + publish_config_change, +) from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.db.create_views import ( create_missing_views, @@ -2971,9 +2975,14 @@ async def get_config_param(prisma_client: Any, param_name: str) -> Optional[Any] return row +async def evict_config_param(param_name: str) -> None: + await litellm_config_cache.async_delete_cache(_config_cache_key(param_name)) + + async def invalidate_config_param(param_name: str) -> None: """Evict from both cache layers; call after every LiteLLM_Config write.""" - await litellm_config_cache.async_delete_cache(_config_cache_key(param_name)) + await evict_config_param(param_name) + await publish_config_change(redis_cache=coordination_redis_cache(), object_type=param_name) async def prefetch_config_params(prisma_client: Any, param_names: List[str]) -> None: diff --git a/litellm/repositories/credentials_repository.py b/litellm/repositories/credentials_repository.py index b5a315d233c..8e4b9ac0be7 100644 --- a/litellm/repositories/credentials_repository.py +++ b/litellm/repositories/credentials_repository.py @@ -9,6 +9,7 @@ so reads return the stored values verbatim. from typing import Any, Dict, Optional from litellm.models.credentials import CredentialItem +from litellm.proxy.common_utils.config_sync_pubsub import wrap_table_actions_for_config_sync class CredentialsRepository: @@ -25,7 +26,10 @@ class CredentialsRepository: @property def table(self) -> Any: - return self.prisma_client.db.litellm_credentialstable + return wrap_table_actions_for_config_sync( + actions=self.prisma_client.db.litellm_credentialstable, + table_name="litellm_credentialstable", + ) @staticmethod def _to_model(record: Any) -> Optional[CredentialItem]: diff --git a/litellm/repositories/model_repository.py b/litellm/repositories/model_repository.py index 0da51519964..50a50cc60d6 100644 --- a/litellm/repositories/model_repository.py +++ b/litellm/repositories/model_repository.py @@ -7,6 +7,7 @@ from typing import Any, Dict, List, Optional, Type from litellm.models.model import LiteLLM_ProxyModelTable from litellm.repositories.base_repository import BaseRepository +from litellm.proxy.common_utils.config_sync_pubsub import wrap_table_actions_for_config_sync from litellm.proxy.common_utils.encrypt_decrypt_utils import ( decrypt_value_helper, encrypt_value_helper, @@ -22,7 +23,10 @@ class ModelRepository(BaseRepository[LiteLLM_ProxyModelTable]): @property def table(self) -> Any: - return self.prisma_client.db.litellm_proxymodeltable + return wrap_table_actions_for_config_sync( + actions=self.prisma_client.db.litellm_proxymodeltable, + table_name="litellm_proxymodeltable", + ) @property def model_class(self) -> Type[LiteLLM_ProxyModelTable]: diff --git a/litellm/repositories/table_repositories.py b/litellm/repositories/table_repositories.py index 54008c0950c..af8be986831 100644 --- a/litellm/repositories/table_repositories.py +++ b/litellm/repositories/table_repositories.py @@ -9,6 +9,8 @@ methods; richer repositories live in their own modules. from typing import Any +from litellm.proxy.common_utils.config_sync_pubsub import wrap_table_actions_for_config_sync + class PrismaTableRepository: """Base for repositories that expose a single Prisma table.""" @@ -26,7 +28,10 @@ class PrismaTableRepository: @property def table(self) -> Any: - return getattr(self.prisma_client.db, self.table_name) + return wrap_table_actions_for_config_sync( + actions=getattr(self.prisma_client.db, self.table_name), + table_name=self.table_name, + ) class PolicyRepository(PrismaTableRepository): diff --git a/ruff.toml b/ruff.toml index 2ea9d7260fb..b652e206f41 100644 --- a/ruff.toml +++ b/ruff.toml @@ -6,7 +6,7 @@ lint.extend-select = ["T20", "PGH004", "RUF008", "RUF009", "RUF100"] # litellm's own ruff config both rely on suppressions this config can't see. lint.external = [ # Enforced by the strict-rule gate (scripts/ruff_strict_gate.py + ruff-strict.toml) - "C901", + "C901", "TID251", # Enforced by upstream litellm's ruff config, but not run in this repo's CI "PLC0415", "E402", "BLE001", "ARG002", "S102", "S324", "S606", "D401", "F403", "F405", ] diff --git a/tests/test_litellm/proxy/common_utils/test_config_sync_pubsub.py b/tests/test_litellm/proxy/common_utils/test_config_sync_pubsub.py new file mode 100644 index 00000000000..8a8ced8bc41 --- /dev/null +++ b/tests/test_litellm/proxy/common_utils/test_config_sync_pubsub.py @@ -0,0 +1,600 @@ +import asyncio +import json +import random +from typing import Callable, Coroutine, Iterable, List, Optional, Tuple +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from redis.asyncio import Redis + +import litellm +from litellm.proxy.common_utils.config_sync_pubsub import ( + CONFIG_SYNC_CHANNEL, + CONFIG_SYNC_JITTER_MAX_SECONDS, + ConfigSyncSubscriber, + _CONFIG_SYNCED_TABLE_NAMES, + _PublishOnWriteActions, + _WRITE_ACTION_NAMES, + publish_config_change, + wrap_table_actions_for_config_sync, +) + +_EXPECTED_WRITE_ACTION_NAMES = ( + "create", + "create_many", + "delete", + "delete_many", + "update", + "update_many", + "upsert", +) + +_EXPECTED_CONFIG_SYNCED_TABLE_NAMES = frozenset( + { + "litellm_agentstable", + "litellm_cacheconfig", + "litellm_configoverrides", + "litellm_credentialstable", + "litellm_guardrailstable", + "litellm_managedvectorstoreindextable", + "litellm_managedvectorstorestable", + "litellm_mcpservertable", + "litellm_policyattachmenttable", + "litellm_policytable", + "litellm_prompttable", + "litellm_proxymodeltable", + "litellm_searchtoolstable", + "litellm_ssoconfig", + } +) + + +class _RecordingRedisClient(Redis): + def __init__(self) -> None: + self.published: List[Tuple[str, str]] = [] + + async def publish(self, channel: str, message: str) -> int: + self.published.append((channel, message)) + return 1 + + +class _FailingPublishRedisClient(Redis): + def __init__(self) -> None: + pass + + async def publish(self, channel: str, message: str) -> int: + raise ConnectionError("redis down") + + +class _NotRedisClient: + def __init__(self) -> None: + self.published: List[Tuple[str, str]] = [] + + async def publish(self, channel: str, message: str) -> int: + self.published.append((channel, message)) + return 1 + + +class _QueuePubSub: + def __init__(self, initial_messages: Iterable[str] = ()) -> None: + self.queue: "asyncio.Queue[str]" = asyncio.Queue() + for message in initial_messages: + self.queue.put_nowait(message) + self.subscribed_channels: List[str] = [] + self.closed = False + + async def subscribe(self, *channels: str) -> None: + self.subscribed_channels.extend(channels) + + async def get_message(self, *, ignore_subscribe_messages: bool, timeout: float) -> Optional[str]: + if timeout == 0: + try: + return self.queue.get_nowait() + except asyncio.QueueEmpty: + return None + try: + return await asyncio.wait_for(self.queue.get(), timeout) + except asyncio.TimeoutError: + return None + + async def aclose(self) -> None: + self.closed = True + + +class _BrokenPubSub(_QueuePubSub): + async def get_message(self, *, ignore_subscribe_messages: bool, timeout: float) -> Optional[str]: + raise ConnectionError("connection lost") + + +class _ScriptedPubSubRedisClient(Redis): + def __init__(self, pubsubs: Iterable[_QueuePubSub]) -> None: + self._scripted_pubsubs = iter(pubsubs) + + def pubsub(self) -> _QueuePubSub: + return next(self._scripted_pubsubs) + + +class _FakeRedisCache: + def __init__(self, client: object, namespace: Optional[str] = None) -> None: + self._client = client + self.namespace = namespace + + def init_async_client(self) -> object: + return self._client + + +class _ExplodingRedisCache: + namespace: Optional[str] = None + + def init_async_client(self) -> object: + raise ConnectionError("cannot connect") + + +def _recording_callback( + events: List[str], name: str, fired: asyncio.Event +) -> Callable[[], Coroutine[None, None, None]]: + async def callback() -> None: + events.append(name) + fired.set() + + return callback + + +async def test_publish_noops_when_redis_cache_is_none() -> None: + await publish_config_change(redis_cache=None, object_type="litellm_proxymodeltable") + + +async def test_publish_sends_object_type_json_on_channel() -> None: + client = _RecordingRedisClient() + cache = _FakeRedisCache(client) + + await publish_config_change(redis_cache=cache, object_type="litellm_proxymodeltable") + + assert len(client.published) == 1 + channel, message = client.published[0] + assert channel == "litellm_proxy.config_change" + assert json.loads(message) == {"object_type": "litellm_proxymodeltable"} + + +async def test_publish_uses_namespaced_channel() -> None: + client = _RecordingRedisClient() + cache = _FakeRedisCache(client, namespace="prod-eu") + + await publish_config_change(redis_cache=cache, object_type="litellm_credentialstable") + + assert client.published[0][0] == "prod-eu:litellm_proxy.config_change" + + +async def test_publish_swallows_redis_publish_errors() -> None: + cache = _FakeRedisCache(_FailingPublishRedisClient()) + + await publish_config_change(redis_cache=cache, object_type="litellm_proxymodeltable") + + +async def test_publish_swallows_client_init_errors() -> None: + await publish_config_change(redis_cache=_ExplodingRedisCache(), object_type="litellm_proxymodeltable") + + +async def test_publish_skips_clients_without_pubsub_support() -> None: + client = _NotRedisClient() + cache = _FakeRedisCache(client) + + await publish_config_change(redis_cache=cache, object_type="litellm_proxymodeltable") + + assert client.published == [] + + +async def test_subscriber_runs_injected_callbacks_in_order_on_message() -> None: + pubsub = _QueuePubSub() + cache = _FakeRedisCache(_ScriptedPubSubRedisClient([pubsub])) + events: List[str] = [] + fired = asyncio.Event() + subscriber = ConfigSyncSubscriber( + redis_cache=cache, + resync_callbacks=( + _recording_callback(events, "add_deployment", asyncio.Event()), + _recording_callback(events, "get_credentials", fired), + ), + debounce_seconds=0.01, + jitter_max_seconds=0.0, + ) + + subscriber.start() + pubsub.queue.put_nowait(json.dumps({"object_type": "litellm_proxymodeltable"})) + await asyncio.wait_for(fired.wait(), timeout=5) + await subscriber.stop() + + assert events == ["add_deployment", "get_credentials"] + assert pubsub.subscribed_channels == [CONFIG_SYNC_CHANNEL] + assert pubsub.closed is True + + +async def test_burst_within_debounce_window_coalesces_into_one_resync() -> None: + burst = [json.dumps({"object_type": "litellm_proxymodeltable"}) for _ in range(5)] + pubsub = _QueuePubSub(initial_messages=burst) + cache = _FakeRedisCache(_ScriptedPubSubRedisClient([pubsub])) + resyncs: List[str] = [] + fired = asyncio.Event() + subscriber = ConfigSyncSubscriber( + redis_cache=cache, + resync_callbacks=(_recording_callback(resyncs, "resync", fired),), + debounce_seconds=0.05, + jitter_max_seconds=0.0, + ) + + subscriber.start() + await asyncio.wait_for(fired.wait(), timeout=5) + await asyncio.sleep(0.3) + await subscriber.stop() + + assert resyncs == ["resync"] + assert pubsub.queue.empty() + + +async def test_subscriber_subscribes_on_namespaced_channel_and_resyncs() -> None: + pubsub = _QueuePubSub() + cache = _FakeRedisCache(_ScriptedPubSubRedisClient([pubsub]), namespace="prod-eu") + resyncs: List[str] = [] + fired = asyncio.Event() + subscriber = ConfigSyncSubscriber( + redis_cache=cache, + resync_callbacks=(_recording_callback(resyncs, "resync", fired),), + debounce_seconds=0.01, + jitter_max_seconds=0.0, + ) + + subscriber.start() + pubsub.queue.put_nowait(json.dumps({"object_type": "litellm_proxymodeltable"})) + await asyncio.wait_for(fired.wait(), timeout=5) + await subscriber.stop() + + assert pubsub.subscribed_channels == ["prod-eu:litellm_proxy.config_change"] + assert resyncs == ["resync"] + + +class _MaxJitterRandom(random.Random): + def uniform(self, a: float, b: float) -> float: + return b + + +async def test_debounce_sleep_adds_jitter_from_injected_rng() -> None: + pubsub = _QueuePubSub(initial_messages=[json.dumps({"object_type": "litellm_proxymodeltable"})]) + cache = _FakeRedisCache(_ScriptedPubSubRedisClient([pubsub])) + sleeps: List[float] = [] + fired = asyncio.Event() + + async def recording_sleep(seconds: float) -> None: + sleeps.append(seconds) + + subscriber = ConfigSyncSubscriber( + redis_cache=cache, + resync_callbacks=(_recording_callback([], "resync", fired),), + debounce_seconds=1.0, + jitter_max_seconds=4.0, + rng=_MaxJitterRandom(), + sleep=recording_sleep, + ) + + subscriber.start() + await asyncio.wait_for(fired.wait(), timeout=5) + await subscriber.stop() + + assert sleeps == [5.0] + + +def test_default_jitter_window_is_nonzero() -> None: + assert CONFIG_SYNC_JITTER_MAX_SECONDS > 0 + + +async def test_redis_error_leads_to_backoff_and_resubscribe() -> None: + broken = _BrokenPubSub() + healthy = _QueuePubSub(initial_messages=[json.dumps({"object_type": "litellm_credentialstable"})]) + cache = _FakeRedisCache(_ScriptedPubSubRedisClient([broken, healthy])) + resyncs: List[str] = [] + fired = asyncio.Event() + subscriber = ConfigSyncSubscriber( + redis_cache=cache, + resync_callbacks=(_recording_callback(resyncs, "resync", fired),), + debounce_seconds=0.01, + jitter_max_seconds=0.0, + backoff_initial_seconds=0.02, + backoff_max_seconds=0.05, + ) + + subscriber.start() + await asyncio.wait_for(fired.wait(), timeout=5) + task = subscriber._task + assert task is not None + assert task.done() is False + await subscriber.stop() + + assert broken.subscribed_channels == [CONFIG_SYNC_CHANNEL] + assert broken.closed is True + assert healthy.subscribed_channels == [CONFIG_SYNC_CHANNEL] + assert resyncs == ["resync"] + + +async def test_failing_resync_callback_does_not_kill_subscriber() -> None: + pubsub = _QueuePubSub() + cache = _FakeRedisCache(_ScriptedPubSubRedisClient([pubsub])) + resyncs: List[str] = [] + fired = asyncio.Event() + + async def failing_callback() -> None: + raise RuntimeError("resync exploded") + + subscriber = ConfigSyncSubscriber( + redis_cache=cache, + resync_callbacks=(failing_callback, _recording_callback(resyncs, "resync", fired)), + debounce_seconds=0.01, + jitter_max_seconds=0.0, + ) + + subscriber.start() + pubsub.queue.put_nowait("change") + await asyncio.wait_for(fired.wait(), timeout=5) + fired.clear() + pubsub.queue.put_nowait("change") + await asyncio.wait_for(fired.wait(), timeout=5) + await subscriber.stop() + + assert resyncs == ["resync", "resync"] + + +async def test_stop_cancels_subscriber_cleanly() -> None: + pubsub = _QueuePubSub() + cache = _FakeRedisCache(_ScriptedPubSubRedisClient([pubsub])) + subscriber = ConfigSyncSubscriber(redis_cache=cache, resync_callbacks=(), debounce_seconds=0.01) + + subscriber.start() + await asyncio.sleep(0.05) + task = subscriber._task + assert task is not None + await subscriber.stop() + + assert task.done() is True + assert subscriber._task is None + assert pubsub.closed is True + await subscriber.stop() + + +async def test_stop_before_start_is_a_noop() -> None: + cache = _FakeRedisCache(_ScriptedPubSubRedisClient([])) + subscriber = ConfigSyncSubscriber(redis_cache=cache, resync_callbacks=()) + + await subscriber.stop() + + +async def test_subscriber_exits_without_callbacks_when_client_lacks_pubsub() -> None: + cache = _FakeRedisCache(_NotRedisClient()) + resyncs: List[str] = [] + subscriber = ConfigSyncSubscriber( + redis_cache=cache, + resync_callbacks=(_recording_callback(resyncs, "resync", asyncio.Event()),), + ) + + subscriber.start() + task = subscriber._task + assert task is not None + await asyncio.wait_for(task, timeout=5) + + assert resyncs == [] + + +class _FakeTableActions: + def __init__(self, calls: List[Tuple[str, str]]) -> None: + self._calls = calls + + async def create(self, **kwargs: object) -> object: + self._calls.append(("write", "create")) + return {"id": "m-1"} + + async def find_many(self, **kwargs: object) -> object: + self._calls.append(("read", "find_many")) + return [] + + +class _AllWritesTableActions: + def __init__(self, calls: List[str]) -> None: + self._calls = calls + + def __getattr__(self, name: str) -> Callable[..., Coroutine[None, None, str]]: + async def action(*args: object, **kwargs: object) -> str: + self._calls.append(name) + return name + + return action + + +def _recording_publish(calls: List[Tuple[str, str]]) -> Callable[[str], Coroutine[None, None, None]]: + async def publish(object_type: str) -> None: + calls.append(("publish", object_type)) + + return publish + + +def test_wrapper_passes_through_unsynced_tables() -> None: + actions = object() + + wrapped = wrap_table_actions_for_config_sync(actions=actions, table_name="litellm_spendlogs") + + assert wrapped is actions + + +async def test_wrapper_publishes_table_name_after_write() -> None: + calls: List[Tuple[str, str]] = [] + wrapped = wrap_table_actions_for_config_sync( + actions=_FakeTableActions(calls), + table_name="litellm_proxymodeltable", + publish=_recording_publish(calls), + ) + + result = await wrapped.create(data={"model_name": "gpt-5.2"}) + + assert result == {"id": "m-1"} + assert calls == [("write", "create"), ("publish", "litellm_proxymodeltable")] + + +async def test_wrapper_does_not_publish_on_reads() -> None: + calls: List[Tuple[str, str]] = [] + wrapped = wrap_table_actions_for_config_sync( + actions=_FakeTableActions(calls), + table_name="litellm_proxymodeltable", + publish=_recording_publish(calls), + ) + + result = await wrapped.find_many(where={}) + + assert result == [] + assert calls == [("read", "find_many")] + + +def test_write_action_names_are_pinned() -> None: + assert _WRITE_ACTION_NAMES == frozenset(_EXPECTED_WRITE_ACTION_NAMES) + + +def test_config_synced_table_membership_is_pinned() -> None: + assert _CONFIG_SYNCED_TABLE_NAMES == _EXPECTED_CONFIG_SYNCED_TABLE_NAMES + + +def test_tool_telemetry_table_writes_pass_through_unwrapped() -> None: + actions = object() + + wrapped = wrap_table_actions_for_config_sync(actions=actions, table_name="litellm_tooltable") + + assert wrapped is actions + + +@pytest.mark.parametrize("action_name", _EXPECTED_WRITE_ACTION_NAMES) +async def test_wrapper_publishes_for_every_write_action(action_name: str) -> None: + write_calls: List[str] = [] + publish_calls: List[Tuple[str, str]] = [] + wrapped = wrap_table_actions_for_config_sync( + actions=_AllWritesTableActions(write_calls), + table_name="litellm_guardrailstable", + publish=_recording_publish(publish_calls), + ) + + result = await getattr(wrapped, action_name)(data={}) + + assert result == action_name + assert write_calls == [action_name] + assert publish_calls == [("publish", "litellm_guardrailstable")] + + +async def test_model_repository_write_publishes_via_live_coordination_cache() -> None: + from litellm.proxy import proxy_server + from litellm.proxy.proxy_server import _set_redis_usage_cache + from litellm.repositories.model_repository import ModelRepository + + client = _RecordingRedisClient() + prisma_client = MagicMock() + prisma_client.db.litellm_proxymodeltable.update = AsyncMock(return_value={"model_id": "m-1"}) + repository = ModelRepository(prisma_client) + table = repository.table + assert isinstance(table, _PublishOnWriteActions) + + previous_cache = proxy_server.redis_usage_cache + _set_redis_usage_cache(_FakeRedisCache(client)) + try: + await table.update(where={"model_id": "m-1"}, data={"model_name": "gpt-5.2"}) + finally: + _set_redis_usage_cache(previous_cache) + + prisma_client.db.litellm_proxymodeltable.update.assert_awaited_once_with( + where={"model_id": "m-1"}, data={"model_name": "gpt-5.2"} + ) + assert len(client.published) == 1 + channel, message = client.published[0] + assert channel == CONFIG_SYNC_CHANNEL + assert json.loads(message) == {"object_type": "litellm_proxymodeltable"} + + +async def test_invalidate_config_param_publishes_param_name() -> None: + from litellm.proxy import proxy_server + from litellm.proxy.proxy_server import _set_redis_usage_cache + from litellm.proxy.utils import invalidate_config_param + + client = _RecordingRedisClient() + previous_cache = proxy_server.redis_usage_cache + _set_redis_usage_cache(_FakeRedisCache(client)) + try: + await invalidate_config_param("environment_variables") + finally: + _set_redis_usage_cache(previous_cache) + + assert len(client.published) == 1 + channel, message = client.published[0] + assert channel == CONFIG_SYNC_CHANNEL + assert json.loads(message) == {"object_type": "environment_variables"} + + +async def test_evict_config_param_does_not_publish() -> None: + from litellm.proxy import proxy_server + from litellm.proxy.proxy_server import _set_redis_usage_cache + from litellm.proxy.utils import evict_config_param + + client = _RecordingRedisClient() + previous_cache = proxy_server.redis_usage_cache + _set_redis_usage_cache(_FakeRedisCache(client)) + try: + await evict_config_param("model_cost_map_reload_config") + finally: + _set_redis_usage_cache(previous_cache) + + assert client.published == [] + + +def _reload_config_prisma_client() -> MagicMock: + config_record = MagicMock() + config_record.param_value = {"interval_hours": 6, "force_reload": True} + prisma_client = MagicMock() + prisma_client.get_generic_data = AsyncMock(return_value=config_record) + prisma_client.db.litellm_config.upsert = AsyncMock(return_value=None) + return prisma_client + + +async def test_model_cost_map_reload_does_not_publish_config_change() -> None: + from litellm.proxy import proxy_server + from litellm.proxy.proxy_server import ProxyConfig, _set_redis_usage_cache + from litellm.proxy.utils import litellm_config_cache + from litellm.utils import _invalidate_model_cost_lowercase_map + + litellm_config_cache.flush_cache() + prisma_client = _reload_config_prisma_client() + client = _RecordingRedisClient() + previous_cache = proxy_server.redis_usage_cache + original_model_cost = litellm.model_cost.copy() + _set_redis_usage_cache(_FakeRedisCache(client)) + try: + with patch("litellm.litellm_core_utils.get_model_cost_map.get_model_cost_map") as mock_get_map: + mock_get_map.return_value = {"gpt-5.2": {"input_cost_per_token": 0.001}} + await ProxyConfig()._check_and_reload_model_cost_map(prisma_client=prisma_client) + finally: + litellm.model_cost = original_model_cost + _invalidate_model_cost_lowercase_map() + _set_redis_usage_cache(previous_cache) + + prisma_client.db.litellm_config.upsert.assert_awaited_once() + assert client.published == [] + + +async def test_anthropic_beta_headers_reload_does_not_publish_config_change() -> None: + from litellm.proxy import proxy_server + from litellm.proxy.proxy_server import ProxyConfig, _set_redis_usage_cache + from litellm.proxy.utils import litellm_config_cache + + litellm_config_cache.flush_cache() + prisma_client = _reload_config_prisma_client() + client = _RecordingRedisClient() + previous_cache = proxy_server.redis_usage_cache + _set_redis_usage_cache(_FakeRedisCache(client)) + try: + with patch("litellm.anthropic_beta_headers_manager.reload_beta_headers_config") as mock_reload: + mock_reload.return_value = {} + await ProxyConfig()._check_and_reload_anthropic_beta_headers(prisma_client=prisma_client) + finally: + _set_redis_usage_cache(previous_cache) + + prisma_client.db.litellm_config.upsert.assert_awaited_once() + assert client.published == [] diff --git a/tests/test_litellm/repositories/test_repositories.py b/tests/test_litellm/repositories/test_repositories.py index c923b722991..994e73d33f5 100644 --- a/tests/test_litellm/repositories/test_repositories.py +++ b/tests/test_litellm/repositories/test_repositories.py @@ -307,6 +307,15 @@ class TestModelRepository: client = MockPrismaClient() return ModelRepository(client) + def test_table_is_wrapped_for_config_sync(self, repo): + from litellm.proxy.common_utils.config_sync_pubsub import ( + _PublishOnWriteActions, + ) + + table = repo.table + assert isinstance(table, _PublishOnWriteActions) + assert table._actions is repo.prisma_client.db.litellm_proxymodeltable + @pytest.mark.asyncio @patch( "litellm.repositories.model_repository.encrypt_value_helper", @@ -1313,6 +1322,15 @@ class TestCredentialsRepository: client = MockPrismaClient() return CredentialsRepository(client) + def test_table_is_wrapped_for_config_sync(self, repo): + from litellm.proxy.common_utils.config_sync_pubsub import ( + _PublishOnWriteActions, + ) + + table = repo.table + assert isinstance(table, _PublishOnWriteActions) + assert table._actions is repo.prisma_client.db.litellm_credentialstable + @pytest.mark.asyncio async def test_create(self, repo): record = await repo.create( @@ -2188,6 +2206,9 @@ class TestConfigRepositoryDeepCopy: class TestPrismaTableRepository: def test_table_property_returns_named_delegate(self): + from litellm.proxy.common_utils.config_sync_pubsub import ( + _PublishOnWriteActions, + ) from litellm.repositories.table_repositories import ( AgentsRepository, PolicyRepository, @@ -2197,9 +2218,11 @@ class TestPrismaTableRepository: agents = AgentsRepository(prisma_client) policy = PolicyRepository(prisma_client) - assert agents.table is prisma_client.db.litellm_agentstable - assert policy.table is prisma_client.db.litellm_policytable - assert agents.table is not policy.table + assert isinstance(agents.table, _PublishOnWriteActions) + assert isinstance(policy.table, _PublishOnWriteActions) + assert agents.table._actions is prisma_client.db.litellm_agentstable + assert policy.table._actions is prisma_client.db.litellm_policytable + assert agents.table._actions is not policy.table._actions def test_table_access_raises_without_db(self): from litellm.repositories.table_repositories import SpendLogsRepository @@ -2208,8 +2231,28 @@ class TestPrismaTableRepository: with pytest.raises(RuntimeError, match="No DB Connected"): _ = repo.table + CONFIG_SYNCED_TABLE_NAMES = frozenset( + { + "litellm_agentstable", + "litellm_cacheconfig", + "litellm_configoverrides", + "litellm_guardrailstable", + "litellm_managedvectorstoreindextable", + "litellm_managedvectorstorestable", + "litellm_mcpservertable", + "litellm_policyattachmenttable", + "litellm_policytable", + "litellm_prompttable", + "litellm_searchtoolstable", + "litellm_ssoconfig", + } + ) + def test_each_repository_binds_its_own_table_name(self): import litellm.repositories.table_repositories as tr + from litellm.proxy.common_utils.config_sync_pubsub import ( + _PublishOnWriteActions, + ) prisma_client = MagicMock() repos = [ @@ -2226,7 +2269,14 @@ class TestPrismaTableRepository: assert name.startswith("litellm_") assert name not in seen, f"duplicate table_name {name}" seen.add(name) - assert repo_cls(prisma_client).table is getattr(prisma_client.db, name) + table = repo_cls(prisma_client).table + raw_actions = getattr(prisma_client.db, name) + if name in self.CONFIG_SYNCED_TABLE_NAMES: + assert isinstance(table, _PublishOnWriteActions), name + assert table._actions is raw_actions + else: + assert table is raw_actions, name + assert self.CONFIG_SYNCED_TABLE_NAMES <= seen def _json_path_equals( From a8018f7500fe5ed801803eada2fbe1c219f30af9 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 31 Jul 2026 21:08:32 -0700 Subject: [PATCH 06/40] fix(proxy): throttle pub/sub resyncs and stop publishing startup-only config params Caps fleet-wide reload rate at one resync per 10s per pod so a burst of authenticated writes cannot amplify into continuous cross-pod reloads, and skips publishing config params (environment_variables, router_settings) that no resync callback applies outside proxy startup --- .../proxy/common_utils/config_sync_pubsub.py | 43 +++ litellm/proxy/proxy_server.py | 56 ++-- litellm/proxy/utils.py | 7 +- .../common_utils/test_config_sync_pubsub.py | 302 +++++++++++++++++- 4 files changed, 380 insertions(+), 28 deletions(-) diff --git a/litellm/proxy/common_utils/config_sync_pubsub.py b/litellm/proxy/common_utils/config_sync_pubsub.py index 4292c67c8c2..e4521c397b7 100644 --- a/litellm/proxy/common_utils/config_sync_pubsub.py +++ b/litellm/proxy/common_utils/config_sync_pubsub.py @@ -1,6 +1,7 @@ import asyncio import json import random +import time from collections.abc import Awaitable, Callable from dataclasses import asdict, dataclass from typing import TYPE_CHECKING, Protocol, cast # noqa: TID251 # untyped prisma/redis boundary needs cast @@ -28,6 +29,7 @@ class _ConfigSyncPubSubClient(Protocol): CONFIG_SYNC_CHANNEL = "litellm_proxy.config_change" CONFIG_SYNC_DEBOUNCE_SECONDS = 1.0 CONFIG_SYNC_JITTER_MAX_SECONDS = 5.0 +CONFIG_SYNC_MIN_RESYNC_INTERVAL_SECONDS = 10.0 _POLL_TIMEOUT_SECONDS = 1.0 _BACKOFF_INITIAL_SECONDS = 5.0 _BACKOFF_MAX_SECONDS = 60.0 @@ -55,6 +57,15 @@ _CONFIG_SYNCED_TABLE_NAMES: frozenset[str] = frozenset( } ) +_RESYNC_APPLIED_CONFIG_PARAM_NAMES: frozenset[str] = frozenset( + { + "general_settings", + "litellm_settings", + "model_cost_map_reload_config", + "anthropic_beta_headers_reload_config", + } +) + def coordination_redis_cache() -> "RedisCache | None": from litellm.proxy.proxy_server import redis_usage_cache @@ -113,6 +124,16 @@ async def publish_config_change_for_object_type(object_type: str) -> None: await publish_config_change(redis_cache=coordination_redis_cache(), object_type=object_type) +async def publish_config_param_change(param_name: str) -> None: + if param_name not in _RESYNC_APPLIED_CONFIG_PARAM_NAMES: + verbose_proxy_logger.debug( + "config sync publish for %s skipped: no resync callback applies this param outside proxy startup", + param_name, + ) + return + await publish_config_change_for_object_type(param_name) + + class _PublishOnWriteActions: __slots__ = ("_actions", "_object_type", "_publish") @@ -156,6 +177,9 @@ class ConfigSyncSubscriber: "_backoff_max_seconds", "_debounce_seconds", "_jitter_max_seconds", + "_last_resync_at", + "_min_resync_interval_seconds", + "_monotonic", "_redis_cache", "_resync_callbacks", "_rng", @@ -169,20 +193,25 @@ class ConfigSyncSubscriber: resync_callbacks: tuple[Callable[[], Awaitable[None]], ...], debounce_seconds: float = CONFIG_SYNC_DEBOUNCE_SECONDS, jitter_max_seconds: float = CONFIG_SYNC_JITTER_MAX_SECONDS, + min_resync_interval_seconds: float = CONFIG_SYNC_MIN_RESYNC_INTERVAL_SECONDS, backoff_initial_seconds: float = _BACKOFF_INITIAL_SECONDS, backoff_max_seconds: float = _BACKOFF_MAX_SECONDS, rng: random.Random | None = None, sleep: Callable[[float], Awaitable[None]] = asyncio.sleep, + monotonic: Callable[[], float] = time.monotonic, ) -> None: self._redis_cache = redis_cache self._resync_callbacks = resync_callbacks self._debounce_seconds = debounce_seconds self._jitter_max_seconds = jitter_max_seconds + self._min_resync_interval_seconds = min_resync_interval_seconds self._backoff_initial_seconds = backoff_initial_seconds self._backoff_max_seconds = backoff_max_seconds self._rng = rng if rng is not None else random.Random() self._sleep = sleep + self._monotonic = monotonic self._task: asyncio.Task[None] | None = None + self._last_resync_at: float | None = None def start(self) -> None: if self._task is not None: @@ -235,8 +264,22 @@ class ConfigSyncSubscriber: if message is None: continue await self._sleep(self._debounce_seconds + self._rng.uniform(0.0, self._jitter_max_seconds)) + await self._wait_for_min_resync_interval() await self._drain_pending(pubsub) await self._run_resync_callbacks() + self._last_resync_at = self._monotonic() + + async def _wait_for_min_resync_interval(self) -> None: + if self._last_resync_at is None: + return + seconds_until_next_resync = self._min_resync_interval_seconds - (self._monotonic() - self._last_resync_at) + if seconds_until_next_resync <= 0: + return + verbose_proxy_logger.debug( + "config sync resync throttled for %.1fs to cap fleet-wide reload rate", + seconds_until_next_resync, + ) + await self._sleep(seconds_until_next_resync) @staticmethod async def _drain_pending(pubsub: _ConfigSyncPubSub) -> None: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 6bb0032b38c..7d3e58bfe52 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1153,11 +1153,7 @@ async def proxy_startup_event(app: FastAPI): except Exception as e: verbose_proxy_logger.error(f"Error stopping DB health watchdog task: {e}") - if proxy_config.config_sync_subscriber is not None: - try: - await proxy_config.config_sync_subscriber.stop() - except Exception as e: - verbose_proxy_logger.error(f"Error stopping config sync subscriber: {e}") + await proxy_config.stop_config_sync_subscriber() await proxy_shutdown_event() # type: ignore[reportGeneralTypeIssues] @@ -6213,6 +6209,38 @@ class ProxyConfig: "litellm.proxy.proxy_server.py::ProxyConfig:add_deployment - {}".format(str(e)) ) + def start_config_sync_subscriber( + self, + prisma_client: PrismaClient, + proxy_logging_obj: ProxyLogging, + redis_cache: Optional[RedisCache], + ) -> None: + if redis_cache is None or self.config_sync_subscriber is not None: + return + + async def _resync_config_from_db() -> None: + await self.add_deployment(prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj) + + async def _resync_credentials_from_db() -> None: + await self.get_credentials(prisma_client=prisma_client) + + subscriber = ConfigSyncSubscriber( + redis_cache=redis_cache, + resync_callbacks=(_resync_config_from_db, _resync_credentials_from_db), + ) + self.config_sync_subscriber = subscriber + subscriber.start() + + async def stop_config_sync_subscriber(self) -> None: + subscriber = self.config_sync_subscriber + if subscriber is None: + return + self.config_sync_subscriber = None + try: + await subscriber.stop() + except Exception as e: + verbose_proxy_logger.error(f"Error stopping config sync subscriber: {e}") + async def _init_non_llm_objects_in_db(self, prisma_client: PrismaClient): """ Use this to read non-llm objects from the db and initialize them @@ -8174,19 +8202,11 @@ class ProxyStartupEvent: ) await proxy_config.get_credentials(prisma_client=prisma_client) - if redis_usage_cache is not None and proxy_config.config_sync_subscriber is None: - - async def _resync_config_from_db() -> None: - await proxy_config.add_deployment(prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj) - - async def _resync_credentials_from_db() -> None: - await proxy_config.get_credentials(prisma_client=prisma_client) - - proxy_config.config_sync_subscriber = ConfigSyncSubscriber( - redis_cache=redis_usage_cache, - resync_callbacks=(_resync_config_from_db, _resync_credentials_from_db), - ) - proxy_config.config_sync_subscriber.start() + proxy_config.start_config_sync_subscriber( + prisma_client=prisma_client, + proxy_logging_obj=proxy_logging_obj, + redis_cache=redis_usage_cache, + ) if store_model_in_db is not True: await proxy_config.init_mcp_servers_from_db() diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index fbcca73e779..5d2d29efa12 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -119,10 +119,7 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.auth.route_checks import RouteChecks -from litellm.proxy.common_utils.config_sync_pubsub import ( - coordination_redis_cache, - publish_config_change, -) +from litellm.proxy.common_utils.config_sync_pubsub import publish_config_param_change from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.db.create_views import ( create_missing_views, @@ -2982,7 +2979,7 @@ async def evict_config_param(param_name: str) -> None: async def invalidate_config_param(param_name: str) -> None: """Evict from both cache layers; call after every LiteLLM_Config write.""" await evict_config_param(param_name) - await publish_config_change(redis_cache=coordination_redis_cache(), object_type=param_name) + await publish_config_param_change(param_name) async def prefetch_config_params(prisma_client: Any, param_names: List[str]) -> None: diff --git a/tests/test_litellm/proxy/common_utils/test_config_sync_pubsub.py b/tests/test_litellm/proxy/common_utils/test_config_sync_pubsub.py index 8a8ced8bc41..f50eef4f1cc 100644 --- a/tests/test_litellm/proxy/common_utils/test_config_sync_pubsub.py +++ b/tests/test_litellm/proxy/common_utils/test_config_sync_pubsub.py @@ -11,9 +11,11 @@ import litellm from litellm.proxy.common_utils.config_sync_pubsub import ( CONFIG_SYNC_CHANNEL, CONFIG_SYNC_JITTER_MAX_SECONDS, + CONFIG_SYNC_MIN_RESYNC_INTERVAL_SECONDS, ConfigSyncSubscriber, _CONFIG_SYNCED_TABLE_NAMES, _PublishOnWriteActions, + _RESYNC_APPLIED_CONFIG_PARAM_NAMES, _WRITE_ACTION_NAMES, publish_config_change, wrap_table_actions_for_config_sync, @@ -48,6 +50,17 @@ _EXPECTED_CONFIG_SYNCED_TABLE_NAMES = frozenset( } ) +_EXPECTED_RESYNC_APPLIED_CONFIG_PARAM_NAMES = frozenset( + { + "anthropic_beta_headers_reload_config", + "general_settings", + "litellm_settings", + "model_cost_map_reload_config", + } +) + +_STARTUP_ONLY_CONFIG_PARAM_NAMES = ("environment_variables", "router_settings") + class _RecordingRedisClient(Redis): def __init__(self) -> None: @@ -106,6 +119,31 @@ class _BrokenPubSub(_QueuePubSub): raise ConnectionError("connection lost") +class _CloseFailingBrokenPubSub(_BrokenPubSub): + async def aclose(self) -> None: + raise ConnectionError("close failed") + + +class _EmptyPollsThenMessagePubSub(_QueuePubSub): + def __init__(self, empty_polls: int, initial_messages: Iterable[str] = ()) -> None: + super().__init__(initial_messages=initial_messages) + self.remaining_empty_polls = empty_polls + + async def get_message(self, *, ignore_subscribe_messages: bool, timeout: float) -> Optional[str]: + if timeout != 0 and self.remaining_empty_polls > 0: + self.remaining_empty_polls -= 1 + return None + return await super().get_message(ignore_subscribe_messages=ignore_subscribe_messages, timeout=timeout) + + +class _FakeClock: + def __init__(self, now: float = 1000.0) -> None: + self.now = now + + def __call__(self) -> float: + return self.now + + class _ScriptedPubSubRedisClient(Redis): def __init__(self, pubsubs: Iterable[_QueuePubSub]) -> None: self._scripted_pubsubs = iter(pubsubs) @@ -286,6 +324,158 @@ def test_default_jitter_window_is_nonzero() -> None: assert CONFIG_SYNC_JITTER_MAX_SECONDS > 0 +def test_default_min_resync_interval_caps_reload_rate() -> None: + assert CONFIG_SYNC_MIN_RESYNC_INTERVAL_SECONDS > CONFIG_SYNC_JITTER_MAX_SECONDS + + +def _throttled_subscriber( + cache: object, + events: List[str], + fired: asyncio.Event, + clock: _FakeClock, + min_resync_interval_seconds: float = 10.0, +) -> ConfigSyncSubscriber: + async def recording_sleep(seconds: float) -> None: + events.append(f"sleep:{seconds}") + await asyncio.sleep(0) + + async def resync() -> None: + events.append("resync") + fired.set() + + return ConfigSyncSubscriber( + redis_cache=cache, + resync_callbacks=(resync,), + debounce_seconds=0.0, + jitter_max_seconds=0.0, + min_resync_interval_seconds=min_resync_interval_seconds, + sleep=recording_sleep, + monotonic=clock, + ) + + +async def test_resync_arriving_inside_min_interval_waits_out_the_remainder() -> None: + pubsub = _QueuePubSub() + cache = _FakeRedisCache(_ScriptedPubSubRedisClient([pubsub])) + events: List[str] = [] + fired = asyncio.Event() + clock = _FakeClock() + subscriber = _throttled_subscriber(cache=cache, events=events, fired=fired, clock=clock) + + subscriber.start() + pubsub.queue.put_nowait("change") + await asyncio.wait_for(fired.wait(), timeout=5) + fired.clear() + clock.now += 4.0 + pubsub.queue.put_nowait("change") + await asyncio.wait_for(fired.wait(), timeout=5) + await subscriber.stop() + + assert events == ["sleep:0.0", "resync", "sleep:0.0", "sleep:6.0", "resync"] + + +async def test_resync_after_min_interval_elapsed_is_not_throttled() -> None: + pubsub = _QueuePubSub() + cache = _FakeRedisCache(_ScriptedPubSubRedisClient([pubsub])) + events: List[str] = [] + fired = asyncio.Event() + clock = _FakeClock() + subscriber = _throttled_subscriber(cache=cache, events=events, fired=fired, clock=clock) + + subscriber.start() + pubsub.queue.put_nowait("change") + await asyncio.wait_for(fired.wait(), timeout=5) + fired.clear() + clock.now += 30.0 + pubsub.queue.put_nowait("change") + await asyncio.wait_for(fired.wait(), timeout=5) + await subscriber.stop() + + assert events == ["sleep:0.0", "resync", "sleep:0.0", "resync"] + + +async def test_writes_during_the_throttle_wait_collapse_into_the_next_resync() -> None: + pubsub = _QueuePubSub() + cache = _FakeRedisCache(_ScriptedPubSubRedisClient([pubsub])) + events: List[str] = [] + fired = asyncio.Event() + clock = _FakeClock() + subscriber = _throttled_subscriber(cache=cache, events=events, fired=fired, clock=clock) + + subscriber.start() + pubsub.queue.put_nowait("change") + await asyncio.wait_for(fired.wait(), timeout=5) + fired.clear() + for _ in range(5): + pubsub.queue.put_nowait("change") + await asyncio.wait_for(fired.wait(), timeout=5) + await asyncio.sleep(0.1) + await subscriber.stop() + + assert events.count("resync") == 2 + assert pubsub.queue.empty() + + +async def test_polls_without_messages_do_not_trigger_resyncs() -> None: + pubsub = _EmptyPollsThenMessagePubSub(empty_polls=3, initial_messages=["change"]) + cache = _FakeRedisCache(_ScriptedPubSubRedisClient([pubsub])) + resyncs: List[str] = [] + fired = asyncio.Event() + subscriber = ConfigSyncSubscriber( + redis_cache=cache, + resync_callbacks=(_recording_callback(resyncs, "resync", fired),), + debounce_seconds=0.01, + jitter_max_seconds=0.0, + ) + + subscriber.start() + await asyncio.wait_for(fired.wait(), timeout=5) + await asyncio.sleep(0.1) + await subscriber.stop() + + assert pubsub.remaining_empty_polls == 0 + assert resyncs == ["resync"] + + +async def test_failing_pubsub_close_still_reconnects() -> None: + broken = _CloseFailingBrokenPubSub() + healthy = _QueuePubSub(initial_messages=["change"]) + cache = _FakeRedisCache(_ScriptedPubSubRedisClient([broken, healthy])) + resyncs: List[str] = [] + fired = asyncio.Event() + subscriber = ConfigSyncSubscriber( + redis_cache=cache, + resync_callbacks=(_recording_callback(resyncs, "resync", fired),), + debounce_seconds=0.01, + jitter_max_seconds=0.0, + backoff_initial_seconds=0.02, + backoff_max_seconds=0.05, + ) + + subscriber.start() + await asyncio.wait_for(fired.wait(), timeout=5) + await subscriber.stop() + + assert healthy.subscribed_channels == [CONFIG_SYNC_CHANNEL] + assert resyncs == ["resync"] + + +async def test_second_start_does_not_open_a_second_subscription() -> None: + pubsub = _QueuePubSub() + cache = _FakeRedisCache(_ScriptedPubSubRedisClient([pubsub])) + subscriber = ConfigSyncSubscriber(redis_cache=cache, resync_callbacks=(), debounce_seconds=0.01) + + subscriber.start() + task = subscriber._task + subscriber.start() + assert task is not None + assert subscriber._task is task + await asyncio.sleep(0.05) + await subscriber.stop() + + assert pubsub.subscribed_channels == [CONFIG_SYNC_CHANNEL] + + async def test_redis_error_leads_to_backoff_and_resubscribe() -> None: broken = _BrokenPubSub() healthy = _QueuePubSub(initial_messages=[json.dumps({"object_type": "litellm_credentialstable"})]) @@ -328,6 +518,7 @@ async def test_failing_resync_callback_does_not_kill_subscriber() -> None: resync_callbacks=(failing_callback, _recording_callback(resyncs, "resync", fired)), debounce_seconds=0.01, jitter_max_seconds=0.0, + min_resync_interval_seconds=0.0, ) subscriber.start() @@ -510,7 +701,7 @@ async def test_model_repository_write_publishes_via_live_coordination_cache() -> assert json.loads(message) == {"object_type": "litellm_proxymodeltable"} -async def test_invalidate_config_param_publishes_param_name() -> None: +async def _publish_calls_for_invalidated_param(param_name: str) -> List[Tuple[str, str]]: from litellm.proxy import proxy_server from litellm.proxy.proxy_server import _set_redis_usage_cache from litellm.proxy.utils import invalidate_config_param @@ -519,14 +710,30 @@ async def test_invalidate_config_param_publishes_param_name() -> None: previous_cache = proxy_server.redis_usage_cache _set_redis_usage_cache(_FakeRedisCache(client)) try: - await invalidate_config_param("environment_variables") + await invalidate_config_param(param_name) finally: _set_redis_usage_cache(previous_cache) + return client.published - assert len(client.published) == 1 - channel, message = client.published[0] + +async def test_invalidate_config_param_publishes_params_a_resync_applies() -> None: + published = await _publish_calls_for_invalidated_param("general_settings") + + assert len(published) == 1 + channel, message = published[0] assert channel == CONFIG_SYNC_CHANNEL - assert json.loads(message) == {"object_type": "environment_variables"} + assert json.loads(message) == {"object_type": "general_settings"} + + +@pytest.mark.parametrize("param_name", _STARTUP_ONLY_CONFIG_PARAM_NAMES) +async def test_invalidate_config_param_does_not_publish_startup_only_params(param_name: str) -> None: + published = await _publish_calls_for_invalidated_param(param_name) + + assert published == [] + + +def test_resync_applied_config_param_membership_is_pinned() -> None: + assert _RESYNC_APPLIED_CONFIG_PARAM_NAMES == _EXPECTED_RESYNC_APPLIED_CONFIG_PARAM_NAMES async def test_evict_config_param_does_not_publish() -> None: @@ -598,3 +805,88 @@ async def test_anthropic_beta_headers_reload_does_not_publish_config_change() -> prisma_client.db.litellm_config.upsert.assert_awaited_once() assert client.published == [] + + +class _StopFailingSubscriber(ConfigSyncSubscriber): + async def stop(self) -> None: + raise RuntimeError("stop failed") + + +async def test_proxy_config_subscriber_resyncs_deployments_and_credentials() -> None: + from litellm.proxy.proxy_server import ProxyConfig + + cache = _FakeRedisCache(_ScriptedPubSubRedisClient([_QueuePubSub()])) + config = ProxyConfig() + prisma_client = MagicMock() + proxy_logging_obj = MagicMock() + calls: List[Tuple[str, object, object]] = [] + + async def fake_add_deployment(prisma_client: object, proxy_logging_obj: object) -> None: + calls.append(("add_deployment", prisma_client, proxy_logging_obj)) + + async def fake_get_credentials(prisma_client: object) -> None: + calls.append(("get_credentials", prisma_client, None)) + + config.add_deployment = fake_add_deployment + config.get_credentials = fake_get_credentials + config.start_config_sync_subscriber( + prisma_client=prisma_client, + proxy_logging_obj=proxy_logging_obj, + redis_cache=cache, + ) + subscriber = config.config_sync_subscriber + assert subscriber is not None + for callback in subscriber._resync_callbacks: + await callback() + await config.stop_config_sync_subscriber() + + assert calls == [ + ("add_deployment", prisma_client, proxy_logging_obj), + ("get_credentials", prisma_client, None), + ] + assert config.config_sync_subscriber is None + assert subscriber._task is None + + +async def test_proxy_config_does_not_start_subscriber_without_coordination_redis() -> None: + from litellm.proxy.proxy_server import ProxyConfig + + config = ProxyConfig() + + config.start_config_sync_subscriber( + prisma_client=MagicMock(), + proxy_logging_obj=MagicMock(), + redis_cache=None, + ) + + assert config.config_sync_subscriber is None + + +async def test_proxy_config_keeps_the_first_subscriber_on_repeat_start() -> None: + from litellm.proxy.proxy_server import ProxyConfig + + cache = _FakeRedisCache(_ScriptedPubSubRedisClient([_QueuePubSub()])) + config = ProxyConfig() + + config.start_config_sync_subscriber(prisma_client=MagicMock(), proxy_logging_obj=MagicMock(), redis_cache=cache) + first = config.config_sync_subscriber + config.start_config_sync_subscriber(prisma_client=MagicMock(), proxy_logging_obj=MagicMock(), redis_cache=cache) + second = config.config_sync_subscriber + await config.stop_config_sync_subscriber() + + assert first is not None + assert second is first + + +async def test_proxy_config_shutdown_survives_a_failing_subscriber_stop() -> None: + from litellm.proxy.proxy_server import ProxyConfig + + config = ProxyConfig() + config.config_sync_subscriber = _StopFailingSubscriber( + redis_cache=_FakeRedisCache(_ScriptedPubSubRedisClient([])), + resync_callbacks=(), + ) + + await config.stop_config_sync_subscriber() + + assert config.config_sync_subscriber is None From 2a13bbe1cb502ba472ad41fecac95363b077c68f Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Fri, 31 Jul 2026 22:13:57 -0700 Subject: [PATCH 07/40] refactor(proxy): resolve team member lookups in one query and cap the rejection message Resolve the requested member user_ids with a single find_many instead of one lookup per member, so a large member list no longer turns into that many round-trips before the permission check runs. Write the member-add audit entries concurrently rather than one after another, and list at most a few ids in the rejection message instead of echoing the whole request back. Update the team-admin member-add case that covered adding a user_id with no user row, which the endpoint now leaves to proxy admins. --- .../management_endpoints/team_endpoints.py | 49 +++++++++++----- tests/proxy_unit_tests/test_proxy_server.py | 2 +- .../test_team_endpoints.py | 57 ++++++++++++++++--- 3 files changed, 87 insertions(+), 21 deletions(-) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 7b5ce33cb8b..a8c5fcea4df 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -2431,12 +2431,23 @@ async def _resolve_existing_member_user_ids( members: Sequence[Member], prisma_client: PrismaClient, ) -> frozenset[str]: - """Return the caller-supplied user_ids that already have a user row.""" - user_repository = UserRepository(prisma_client) - found = await asyncio.gather( - *(user_repository.find_by_id(member.user_id) for member in members if member.user_id is not None) + """Return the caller-supplied user_ids that already have a user row. + + Resolved with a single query so the number of members in the request does + not translate into that many concurrent connections. + """ + requested_user_ids = frozenset(member.user_id for member in members if member.user_id is not None) + if not requested_user_ids: + return frozenset() + + found = await UserRepository(prisma_client).table.find_many( + where={ # mutable-ok: Prisma query filters are dict-shaped + "user_id": { # mutable-ok: Prisma query filters are dict-shaped + "in": sorted(requested_user_ids) + } + } ) - return frozenset(user.user_id for user in found if user is not None and user.user_id is not None) + return frozenset(user.user_id for user in found or () if user.user_id is not None) def _pre_existing_user_ids( @@ -2459,6 +2470,9 @@ def _pre_existing_user_ids( return existing_user_ids | populated_user_ids +_MAX_REPORTED_UNKNOWN_USER_IDS = 10 + + def _validate_member_user_id_provisioning( members: Sequence[Member], existing_user_ids: frozenset[str], @@ -2481,13 +2495,15 @@ def _validate_member_user_id_provisioning( if not unknown_user_ids: return + listed = ", ".join(unknown_user_ids[:_MAX_REPORTED_UNKNOWN_USER_IDS]) + remaining = len(unknown_user_ids) - _MAX_REPORTED_UNKNOWN_USER_IDS raise HTTPException( status_code=403, detail={ # mutable-ok: HTTPException detail must be a plain mapping to keep this route's {"error": ...} response shape "error": ( - "Only proxy admins can add a user_id that does not exist yet: {}. " + "Only proxy admins can add a user_id that does not exist yet: {}{}. " "Add the member by user_email to invite a new user, or ask a proxy admin " - "to create the user first.".format(", ".join(unknown_user_ids)) + "to create the user first.".format(listed, " and {} more".format(remaining) if remaining > 0 else "") ) }, ) @@ -2515,13 +2531,15 @@ async def _create_team_member_add_audit_logs( user_api_key_dict: UserAPIKeyAuth, litellm_proxy_admin_name: str, ) -> None: - """Record the membership change, and any user row it created, in the audit log.""" + """Record the membership change, and any user row it created, in the audit log. + + The entries are written concurrently so a request adding many members does + not pay for them one after another. + """ from litellm.proxy.management_helpers.audit_logs import create_object_audit_log - for user in updated_users: - if user.user_id is None or user.user_id in existing_user_ids: - continue - await create_object_audit_log( + created_user_entries = tuple( + create_object_audit_log( object_id=user.user_id, action="created", litellm_changed_by=None, @@ -2531,8 +2549,11 @@ async def _create_team_member_add_audit_logs( before_value=None, after_value=safe_dumps(user.model_dump(exclude_none=True)), ) + for user in updated_users + if user.user_id is not None and user.user_id not in existing_user_ids + ) - await create_object_audit_log( + membership_entry = create_object_audit_log( object_id=team_id, action="updated", litellm_changed_by=None, @@ -2543,6 +2564,8 @@ async def _create_team_member_add_audit_logs( after_value=_members_audit_value(after_members), ) + await asyncio.gather(*created_user_entries, membership_entry) + async def _validate_and_populate_member_user_info( member: Member, diff --git a/tests/proxy_unit_tests/test_proxy_server.py b/tests/proxy_unit_tests/test_proxy_server.py index bedd4dd1838..f64994cb3b1 100644 --- a/tests/proxy_unit_tests/test_proxy_server.py +++ b/tests/proxy_unit_tests/test_proxy_server.py @@ -1475,7 +1475,7 @@ async def test_create_team_member_add_team_admin( user_api_key_dict=valid_token, ) except HTTPException as e: - if user_role == "user": + if user_role == "user" or new_member_method == "user_id": assert e.status_code == 403 return else: diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index b1438447e4f..0de2c1ac71d 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -10440,20 +10440,18 @@ def test_validate_member_user_id_provisioning_reports_every_unknown_member(): @pytest.mark.asyncio async def test_resolve_existing_member_user_ids_matches_caller_supplied_user_ids(): - """Only caller-supplied user_ids are looked up; unknown ones resolve to nothing.""" + """Caller-supplied user_ids resolve in one query; unknown ones resolve to nothing.""" from litellm.proxy.management_endpoints.team_endpoints import ( _resolve_existing_member_user_ids, ) prisma_client = MagicMock() - - async def find_by_id(user_id): - if user_id == "by-id": - return LiteLLM_UserTable(user_id="by-id", max_budget=None, spend=0.0, user_email=None, models=[]) - return None + find_many = AsyncMock( + return_value=[LiteLLM_UserTable(user_id="by-id", max_budget=None, spend=0.0, user_email=None, models=[])] + ) with patch("litellm.proxy.management_endpoints.team_endpoints.UserRepository") as repo: - repo.return_value.find_by_id = AsyncMock(side_effect=find_by_id) + repo.return_value.table.find_many = find_many resolved = await _resolve_existing_member_user_ids( members=[ @@ -10465,6 +10463,28 @@ async def test_resolve_existing_member_user_ids_matches_caller_supplied_user_ids ) assert resolved == frozenset({"by-id"}) + # one round-trip, and email-only members contribute no id to look up + find_many.assert_awaited_once() + assert find_many.await_args.kwargs["where"] == {"user_id": {"in": ["by-id", "missing"]}} + + +@pytest.mark.asyncio +async def test_resolve_existing_member_user_ids_skips_the_query_when_no_user_ids(): + """An all-email payload must not hit the database at all.""" + from litellm.proxy.management_endpoints.team_endpoints import ( + _resolve_existing_member_user_ids, + ) + + with patch("litellm.proxy.management_endpoints.team_endpoints.UserRepository") as repo: + repo.return_value.table.find_many = AsyncMock() + + resolved = await _resolve_existing_member_user_ids( + members=[Member(user_email="a@example.com", role="user")], + prisma_client=MagicMock(), + ) + + assert resolved == frozenset() + repo.return_value.table.find_many.assert_not_awaited() def test_pre_existing_user_ids_counts_ids_filled_in_by_member_resolution(): @@ -10574,3 +10594,26 @@ async def test_team_member_add_audits_a_user_created_from_a_list_payload(monkeyp mock_audit.assert_called_once() assert created_user_id not in mock_audit.call_args.kwargs["existing_user_ids"] + + +def test_validate_member_user_id_provisioning_caps_the_ids_it_echoes_back(): + """A large member list must not echo every id back in the error body.""" + from litellm.proxy.management_endpoints.team_endpoints import ( + _MAX_REPORTED_UNKNOWN_USER_IDS, + _validate_member_user_id_provisioning, + ) + + members = [Member(user_id=f"u{i}", role="user") for i in range(500)] + + with pytest.raises(HTTPException) as exc_info: + _validate_member_user_id_provisioning( + members=members, + existing_user_ids=frozenset(), + user_api_key_dict=_provisioning_caller(LitellmUserRoles.INTERNAL_USER), + ) + + detail = str(exc_info.value.detail) + assert "u0" in detail + assert f"u{_MAX_REPORTED_UNKNOWN_USER_IDS}" not in detail + assert f"and {500 - _MAX_REPORTED_UNKNOWN_USER_IDS} more" in detail + assert len(detail) < 1000 From 77e490a69513c44fa2157d41d8a69ed5333b529f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 31 Jul 2026 22:52:00 -0700 Subject: [PATCH 08/40] fix(proxy): publish router_settings changes so peer pods apply them on resync add_deployment already reapplies DB router settings through _update_llm_router, so gating router_settings out of the pub/sub publish set left the push path covering less than the resync actually applies --- .../proxy/common_utils/config_sync_pubsub.py | 1 + .../common_utils/test_config_sync_pubsub.py | 10 +++--- .../proxy/proxy_server/test_proxy_config.py | 36 +++++++++++++++++++ 3 files changed, 43 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/common_utils/config_sync_pubsub.py b/litellm/proxy/common_utils/config_sync_pubsub.py index e4521c397b7..76c3066ae83 100644 --- a/litellm/proxy/common_utils/config_sync_pubsub.py +++ b/litellm/proxy/common_utils/config_sync_pubsub.py @@ -60,6 +60,7 @@ _CONFIG_SYNCED_TABLE_NAMES: frozenset[str] = frozenset( _RESYNC_APPLIED_CONFIG_PARAM_NAMES: frozenset[str] = frozenset( { "general_settings", + "router_settings", "litellm_settings", "model_cost_map_reload_config", "anthropic_beta_headers_reload_config", diff --git a/tests/test_litellm/proxy/common_utils/test_config_sync_pubsub.py b/tests/test_litellm/proxy/common_utils/test_config_sync_pubsub.py index f50eef4f1cc..6872407808c 100644 --- a/tests/test_litellm/proxy/common_utils/test_config_sync_pubsub.py +++ b/tests/test_litellm/proxy/common_utils/test_config_sync_pubsub.py @@ -56,10 +56,11 @@ _EXPECTED_RESYNC_APPLIED_CONFIG_PARAM_NAMES = frozenset( "general_settings", "litellm_settings", "model_cost_map_reload_config", + "router_settings", } ) -_STARTUP_ONLY_CONFIG_PARAM_NAMES = ("environment_variables", "router_settings") +_STARTUP_ONLY_CONFIG_PARAM_NAMES = ("environment_variables",) class _RecordingRedisClient(Redis): @@ -716,13 +717,14 @@ async def _publish_calls_for_invalidated_param(param_name: str) -> List[Tuple[st return client.published -async def test_invalidate_config_param_publishes_params_a_resync_applies() -> None: - published = await _publish_calls_for_invalidated_param("general_settings") +@pytest.mark.parametrize("param_name", sorted(_EXPECTED_RESYNC_APPLIED_CONFIG_PARAM_NAMES)) +async def test_invalidate_config_param_publishes_params_a_resync_applies(param_name: str) -> None: + published = await _publish_calls_for_invalidated_param(param_name) assert len(published) == 1 channel, message = published[0] assert channel == CONFIG_SYNC_CHANNEL - assert json.loads(message) == {"object_type": "general_settings"} + assert json.loads(message) == {"object_type": param_name} @pytest.mark.parametrize("param_name", _STARTUP_ONLY_CONFIG_PARAM_NAMES) diff --git a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py index 8d1d8185e4d..a79bd25b60b 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -2058,6 +2058,42 @@ async def test_ProxyConfig__add_router_settings_from_db_config_none_router_noop( 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") + + # --------------------------------------------------------------------------- # ProxyConfig._add_general_settings_from_db_config # --------------------------------------------------------------------------- From 48b3d1889395bb63c0a3a458167ed65073538b5d Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 1 Aug 2026 09:28:02 -0700 Subject: [PATCH 09/40] feat(ui): note in the add-member modal that search covers existing users only Both fields select from a server-side search over existing accounts, so a typed-in address or id never becomes a value. Say so up front rather than letting the form look like it accepts a new user and fail on submit. Applies to the organization member modal too, which shares this component. --- .../common_components/user_search_modal.test.tsx | 10 ++++++++++ .../components/common_components/user_search_modal.tsx | 10 +++++++++- 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/ui/litellm-dashboard/src/components/common_components/user_search_modal.test.tsx b/ui/litellm-dashboard/src/components/common_components/user_search_modal.test.tsx index 72b0e10e5d6..634c78d28d0 100644 --- a/ui/litellm-dashboard/src/components/common_components/user_search_modal.test.tsx +++ b/ui/litellm-dashboard/src/components/common_components/user_search_modal.test.tsx @@ -65,4 +65,14 @@ describe("UserSearchModal", () => { expect(userFilterUICall).not.toHaveBeenCalled(); }); + + it("tells the user that only existing accounts can be selected", () => { + renderModal(); + + const notice = screen.getByRole("alert"); + expect(notice).toHaveTextContent(/users that already exist/i); + expect(notice).toHaveTextContent(/ask a proxy admin to create their account first/i); + // info, not warning: a warning here would read as an error state on an empty form + expect(notice.className).toMatch(/ant-alert-info/); + }); }); diff --git a/ui/litellm-dashboard/src/components/common_components/user_search_modal.tsx b/ui/litellm-dashboard/src/components/common_components/user_search_modal.tsx index fafafd8e5d5..9c1d64a1f86 100644 --- a/ui/litellm-dashboard/src/components/common_components/user_search_modal.tsx +++ b/ui/litellm-dashboard/src/components/common_components/user_search_modal.tsx @@ -1,5 +1,5 @@ import { useState } from "react"; -import { Modal, Form, Button, Select, Tooltip } from "antd"; +import { Modal, Form, Button, Select, Tooltip, Alert } from "antd"; import { UserAddOutlined } from "@ant-design/icons"; import { useDebouncedCallback } from "@tanstack/react-pacer/debouncer"; import { userFilterUICall } from "@/components/networking"; @@ -140,6 +140,14 @@ const UserSearchModal: React.FC = ({ role: defaultRole, }} > + +