From a558a0b6a983a78fd24a7e1f3760482c079b6255 Mon Sep 17 00:00:00 2001 From: yassin Date: Tue, 8 Sep 2026 23:10:03 +0000 Subject: [PATCH 01/12] fix(proxy): renew budget reservation counter TTL while the request is in flight A reservation lives inside spend counter keys that expire on the Redis idle TTL (60s). A stream that outlives the TTL dropped its reservation, so a concurrent request on any worker was admitted against the DB floor until the stream reconciled. Renew the counter TTL with EXPIRE every ttl/2 while the reservation is open and stop once it is finalized, so an idle counter still expires on its own if the worker dies. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/caching/redis_cache.py | 15 +++++ litellm/proxy/proxy_server.py | 6 ++ .../spend_tracking/budget_reservation.py | 44 +++++++++++++- .../proxy/test_budget_reservation.py | 60 ++++++++++++++++++- 4 files changed, 123 insertions(+), 2 deletions(-) diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index 106c1580110..b1531e299d0 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -76,6 +76,8 @@ class _AsyncRedisCommands(Protocol): def ttl(self, name: str) -> Awaitable[int]: ... + def expire(self, name: str, time: int) -> Awaitable[bool]: ... + def rpush(self, name: str, *values: str | bytes | float) -> Awaitable[int]: ... def lpop(self, name: str, count: int | None = None) -> Awaitable[object]: ... @@ -1795,6 +1797,19 @@ class RedisCache(BaseCache): _record_swallowed_redis_failure(self._circuit_breaker, e) return None + @_redis_circuit_breaker_guard + async def async_refresh_ttl(self, key: str, ttl: int | None = None) -> bool: + """EXPIRE an existing key without touching its value. False when the key is absent or Redis failed.""" + _used_ttl: Final = self.get_ttl(ttl=ttl) + if _used_ttl is None: + return False + try: + return await self._async_commands().expire(self.check_and_fix_namespace(key=key), _used_ttl) + except Exception as e: + verbose_logger.debug("Redis EXPIRE Error: %s", e) + _record_swallowed_redis_failure(self._circuit_breaker, e) + return False + @_redis_circuit_breaker_guard async def async_rpush( self, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index a617aec9f5c..33f099540b4 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -3222,6 +3222,12 @@ async def increment_spend_counter(counter_key: str, increment: float): return await _increment_spend_counter_cache(counter_key=counter_key, increment=increment) +async def refresh_spend_counter_ttl(counter_key: str) -> bool: + if spend_counter_cache.redis_cache is None: + return False + return await spend_counter_cache.redis_cache.async_refresh_ttl(key=counter_key) + + async def _increment_spend_counter_cache(counter_key: str, increment: float): if spend_counter_cache.redis_cache is not None: try: diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index d9f6e33c43c..bd9eb154839 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -2,6 +2,7 @@ from __future__ import annotations import asyncio import json +import time from collections.abc import Mapping, Sequence from dataclasses import dataclass from datetime import datetime, timedelta, timezone @@ -99,6 +100,42 @@ def get_reserved_counter_keys(budget_reservation: dict | None) -> set: } +_lease_renewals: Final[set[asyncio.Task[None]]] = set() # mutable-ok: asyncio only weak-refs pending tasks + + +def _start_reservation_lease_renewal(budget_reservation: Mapping[str, object], counter_keys: frozenset[str]) -> None: + """A reservation lives inside spend counter keys that expire on their Redis TTL. Renew the TTL + while the request is in flight so a request longer than the TTL does not drop its + reservation and admit concurrent requests against the DB floor on any worker.""" + from litellm.proxy.proxy_server import spend_counter_cache + + if spend_counter_cache.redis_cache is None or not counter_keys: + return + task: Final = asyncio.create_task( + _renew_reservation_lease( + budget_reservation=budget_reservation, + counter_keys=counter_keys, + interval=spend_counter_cache.redis_cache.default_ttl / 2, + ) + ) + _lease_renewals.add(task) + task.add_done_callback(_lease_renewals.discard) + + +async def _renew_reservation_lease( + budget_reservation: Mapping[str, object], counter_keys: frozenset[str], interval: float +) -> None: + from litellm.proxy.proxy_server import refresh_spend_counter_ttl + + deadline: Final = time.monotonic() + litellm.request_timeout + while time.monotonic() < deadline: + await asyncio.sleep(interval) + if budget_reservation.get("finalized") is True: + return + for counter_key in counter_keys: + await refresh_spend_counter_ttl(counter_key=counter_key) + + def _key_reservation_should_release_for_throttle(counter_key: str, valid_token: UserAPIKeyAuth | None) -> bool: """ Whether an over-budget key's own ``max_budget`` reservation should be @@ -294,12 +331,17 @@ async def reserve_budget_for_request( llm_router=llm_router, input_token_counts=input_token_counts, ) - return { + budget_reservation: Final = { "reserved_cost": reservation_cost, "entries": applied_entries, "finalized": False, "input_cost": min(float(input_cost or 0.0), reservation_cost), } + _start_reservation_lease_renewal( + budget_reservation=budget_reservation, + counter_keys=frozenset(get_reserved_counter_keys(budget_reservation=budget_reservation)), + ) + return budget_reservation async def reconcile_budget_reservation( diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 6dab054d8ea..2812be2c53e 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -1,5 +1,6 @@ import asyncio import threading +import time from collections.abc import Mapping from datetime import datetime, timedelta, timezone from types import SimpleNamespace @@ -2199,26 +2200,83 @@ async def test_release_non_numeric_counter_reseeds_from_db(spend_counter_state): class _ExpiringRedisCache: - def __init__(self) -> None: + """In-memory stand-in for RedisCache with real wall-clock key expiry.""" + + def __init__(self, default_ttl: float = 60.0) -> None: + self.default_ttl = default_ttl self.store: dict[str, float] = {} + self.expires_at: dict[str, float] = {} + self.refresh_count = 0 + + def _evict_expired(self, key: str) -> None: + if self.expires_at.get(key, float("inf")) <= time.monotonic(): + self.store.pop(key, None) + self.expires_at.pop(key, None) async def async_get_cache(self, key: str, *args: object, **kwargs: object) -> float | None: + self._evict_expired(key) return self.store.get(key) async def async_increment(self, key: str, value: float, **kwargs: object) -> float: + self._evict_expired(key) self.store[key] = self.store.get(key, 0.0) + float(value) + self.expires_at[key] = time.monotonic() + self.default_ttl return self.store[key] async def async_set_max(self, key: str, value: float, **kwargs: object) -> float: + self._evict_expired(key) self.store[key] = max(self.store.get(key, float("-inf")), float(value)) + self.expires_at[key] = time.monotonic() + self.default_ttl return self.store[key] async def async_set_cache(self, key: str, value: float, *args: object, **kwargs: object) -> bool: self.store[key] = float(value) + self.expires_at[key] = time.monotonic() + self.default_ttl return True async def async_delete_cache(self, key: str, *args: object, **kwargs: object) -> None: self.store.pop(key, None) + self.expires_at.pop(key, None) + + async def async_refresh_ttl(self, key: str, ttl: int | None = None) -> bool: + self._evict_expired(key) + if key not in self.store: + return False + self.refresh_count += 1 + self.expires_at[key] = time.monotonic() + (ttl if ttl is not None else self.default_ttl) + return True + + +@pytest.mark.asyncio +async def test_reservation_survives_redis_counter_ttl_while_request_in_flight( + spend_counter_state, +): + """A request that runs longer than the counter TTL must keep its reservation in Redis + (so a concurrent request on any worker still sees it), and renewal must stop once the + reservation is reconciled so an idle counter still expires on its own.""" + counter_cache, key_cache = spend_counter_state + redis_cache = _ExpiringRedisCache(default_ttl=0.2) + counter_cache.redis_cache = redis_cache + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth(token="key-lease", spend=0.0, max_budget=1.0) + counter_key = "spend:key:key-lease" + + reservation = await _reserve(valid_token, 0.6, key_cache, proxy_logging_obj) + assert reservation is not None + + await asyncio.sleep(0.5) + assert await redis_cache.async_get_cache(key=counter_key) == pytest.approx(0.6) + concurrent = await _reserve(valid_token, 0.6, key_cache, proxy_logging_obj) + assert concurrent is not None + assert concurrent["reserved_cost"] == pytest.approx(0.4) + + await release_budget_reservation(reservation) + await release_budget_reservation(concurrent) + await asyncio.sleep(0.15) + refreshes_after_release = redis_cache.refresh_count + await asyncio.sleep(0.35) + assert redis_cache.refresh_count == refreshes_after_release + assert await redis_cache.async_get_cache(key=counter_key) is None @pytest.mark.asyncio From 3cd4a768aef300c62f66792242904c9bc801c00d Mon Sep 17 00:00:00 2001 From: yassin Date: Tue, 8 Sep 2026 23:24:21 +0000 Subject: [PATCH 02/12] fix(proxy): keep the reservation lease alive across a failed Redis EXPIRE Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/caching/redis_cache.py | 9 ++----- litellm/proxy/proxy_server.py | 6 ++++- .../proxy/test_budget_reservation.py | 27 ++++++++++++++++++- 3 files changed, 33 insertions(+), 9 deletions(-) diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index b1531e299d0..5ed739b5d10 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -1799,16 +1799,11 @@ class RedisCache(BaseCache): @_redis_circuit_breaker_guard async def async_refresh_ttl(self, key: str, ttl: int | None = None) -> bool: - """EXPIRE an existing key without touching its value. False when the key is absent or Redis failed.""" + """EXPIRE an existing key without touching its value. False when the key is absent.""" _used_ttl: Final = self.get_ttl(ttl=ttl) if _used_ttl is None: return False - try: - return await self._async_commands().expire(self.check_and_fix_namespace(key=key), _used_ttl) - except Exception as e: - verbose_logger.debug("Redis EXPIRE Error: %s", e) - _record_swallowed_redis_failure(self._circuit_breaker, e) - return False + return await self._async_commands().expire(self.check_and_fix_namespace(key=key), _used_ttl) @_redis_circuit_breaker_guard async def async_rpush( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 33f099540b4..476eda3d081 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -3225,7 +3225,11 @@ async def increment_spend_counter(counter_key: str, increment: float): async def refresh_spend_counter_ttl(counter_key: str) -> bool: if spend_counter_cache.redis_cache is None: return False - return await spend_counter_cache.redis_cache.async_refresh_ttl(key=counter_key) + try: + return await spend_counter_cache.redis_cache.async_refresh_ttl(key=counter_key) + except Exception as e: + verbose_proxy_logger.debug("spend counter TTL refresh skipped for %s: %s", counter_key, e) + return False async def _increment_spend_counter_cache(counter_key: str, increment: float): diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 2812be2c53e..3c33ca11c91 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -2202,11 +2202,13 @@ async def test_release_non_numeric_counter_reseeds_from_db(spend_counter_state): class _ExpiringRedisCache: """In-memory stand-in for RedisCache with real wall-clock key expiry.""" - def __init__(self, default_ttl: float = 60.0) -> None: + def __init__(self, default_ttl: float = 60.0, fail_first_refresh: bool = False) -> None: self.default_ttl = default_ttl self.store: dict[str, float] = {} self.expires_at: dict[str, float] = {} + self.refresh_attempts = 0 self.refresh_count = 0 + self.fail_first_refresh = fail_first_refresh def _evict_expired(self, key: str) -> None: if self.expires_at.get(key, float("inf")) <= time.monotonic(): @@ -2239,6 +2241,9 @@ class _ExpiringRedisCache: self.expires_at.pop(key, None) async def async_refresh_ttl(self, key: str, ttl: int | None = None) -> bool: + self.refresh_attempts += 1 + if self.fail_first_refresh and self.refresh_attempts == 1: + raise ConnectionError("Redis circuit breaker is open") self._evict_expired(key) if key not in self.store: return False @@ -2279,6 +2284,26 @@ async def test_reservation_survives_redis_counter_ttl_while_request_in_flight( assert await redis_cache.async_get_cache(key=counter_key) is None +@pytest.mark.asyncio +async def test_reservation_lease_keeps_renewing_after_transient_redis_failure( + spend_counter_state, +): + """One failed EXPIRE (Redis blip, open circuit breaker) must not end renewal for the + rest of the request.""" + counter_cache, key_cache = spend_counter_state + redis_cache = _ExpiringRedisCache(default_ttl=0.2, fail_first_refresh=True) + counter_cache.redis_cache = redis_cache + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth(token="key-lease-blip", spend=0.0, max_budget=1.0) + + reservation = await _reserve(valid_token, 0.6, key_cache, proxy_logging_obj) + assert reservation is not None + + await asyncio.sleep(0.5) + assert redis_cache.refresh_attempts >= 3 + await release_budget_reservation(reservation) + + @pytest.mark.asyncio async def test_reconcile_after_redis_counter_expiry_keeps_request_cost_enforced( spend_counter_state, From 668e0142664ac6fa6ab6a6c53925c79cf20b1b6e Mon Sep 17 00:00:00 2001 From: yassin Date: Wed, 9 Sep 2026 00:11:05 +0000 Subject: [PATCH 03/12] fix(proxy): stop reservation lease renewal once the request task is gone A streaming /v1/messages client disconnect skips reconciliation, so the lease kept the orphaned reservation alive until request_timeout instead of the plain 60s counter TTL the base branch had. Stop renewing when the request task that took the reservation is done. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../spend_tracking/budget_reservation.py | 10 ++++++-- .../proxy/test_budget_reservation.py | 23 +++++++++++++++++++ 2 files changed, 31 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index bd9eb154839..c0af6685600 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -116,6 +116,7 @@ def _start_reservation_lease_renewal(budget_reservation: Mapping[str, object], c budget_reservation=budget_reservation, counter_keys=counter_keys, interval=spend_counter_cache.redis_cache.default_ttl / 2, + request_task=asyncio.current_task(), ) ) _lease_renewals.add(task) @@ -123,14 +124,19 @@ def _start_reservation_lease_renewal(budget_reservation: Mapping[str, object], c async def _renew_reservation_lease( - budget_reservation: Mapping[str, object], counter_keys: frozenset[str], interval: float + budget_reservation: Mapping[str, object], + counter_keys: frozenset[str], + interval: float, + request_task: asyncio.Task[object] | None, ) -> None: + """Stops on finalization or once the request task that took the reservation is gone, so a + disconnect path that skipped reconciliation falls back to the plain counter TTL.""" from litellm.proxy.proxy_server import refresh_spend_counter_ttl deadline: Final = time.monotonic() + litellm.request_timeout while time.monotonic() < deadline: await asyncio.sleep(interval) - if budget_reservation.get("finalized") is True: + if budget_reservation.get("finalized") is True or (request_task is not None and request_task.done()): return for counter_key in counter_keys: await refresh_spend_counter_ttl(counter_key=counter_key) diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 3c33ca11c91..626c82ea397 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -2304,6 +2304,29 @@ async def test_reservation_lease_keeps_renewing_after_transient_redis_failure( await release_budget_reservation(reservation) +@pytest.mark.asyncio +async def test_reservation_lease_stops_when_request_task_ends_without_reconciling( + spend_counter_state, +): + """A request whose task ends without reconciling (client disconnect path that skips the + cost callbacks) must not keep renewing: the counter falls back to its plain TTL instead of + pinning the reservation until the request timeout.""" + counter_cache, key_cache = spend_counter_state + redis_cache = _ExpiringRedisCache(default_ttl=0.2) + counter_cache.redis_cache = redis_cache + proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache) + valid_token = UserAPIKeyAuth(token="key-lease-orphan", spend=0.0, max_budget=1.0) + counter_key = "spend:key:key-lease-orphan" + + reservation = await asyncio.create_task(_reserve(valid_token, 0.6, key_cache, proxy_logging_obj)) + assert reservation is not None + assert reservation["finalized"] is False + + await asyncio.sleep(0.5) + assert redis_cache.refresh_count == 0 + assert await redis_cache.async_get_cache(key=counter_key) is None + + @pytest.mark.asyncio async def test_reconcile_after_redis_counter_expiry_keeps_request_cost_enforced( spend_counter_state, From b415f2263a6b5cb3b26071c977c55bcc91604a8b Mon Sep 17 00:00:00 2001 From: yassin Date: Sun, 20 Sep 2026 08:59:46 +0000 Subject: [PATCH 04/12] test(proxy): restore pipeline methods on expiring Redis fake after merge Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/test_budget_reservation.py | 20 +++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 6905585fd5e..86bf896188f 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -1,7 +1,7 @@ import asyncio import threading import time -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from datetime import datetime, timedelta, timezone from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch @@ -11,6 +11,7 @@ from fastapi import HTTPException import litellm from litellm.caching.dual_cache import DualCache +from litellm.types.caching import RedisPipelineIncrementOperation from litellm.constants import STREAM_SSE_KEEPALIVE_PING_BYTES from litellm.llms.anthropic.experimental_pass_through.messages.agentic_streaming_iterator import ( AgenticAnthropicStreamingIterator, @@ -2400,6 +2401,14 @@ class _ExpiringRedisCache: self.expires_at[key] = time.monotonic() + (ttl if ttl is not None else self.default_ttl) return True + async def async_increment_pipeline( + self, increment_list: Sequence[RedisPipelineIncrementOperation], **kwargs: object + ) -> list[float]: + return [await self.async_increment(op["key"], op["increment_value"]) for op in increment_list] + + def get_ttl(self, **kwargs: object) -> int | None: + return int(self.default_ttl) + @pytest.mark.asyncio async def test_reservation_survives_redis_counter_ttl_while_request_in_flight( @@ -2475,15 +2484,6 @@ async def test_reservation_lease_stops_when_request_task_ends_without_reconcilin assert redis_cache.refresh_count == 0 assert await redis_cache.async_get_cache(key=counter_key) is None - async def async_increment_pipeline(self, increment_list, **kwargs): - results = [] - for op in increment_list: - results.append(await self.async_increment(op["key"], op["increment_value"])) - return results - - def get_ttl(self, **kwargs) -> None: - return None - class _TeamMembershipFloorDb: """Stands in for `prisma_client.db`: only the team-membership row exists and its spend is the DB floor.""" From 0e1605010052a54b5a2b9e01fee89499e10858d9 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 21 Sep 2026 12:21:50 -0700 Subject: [PATCH 05/12] fix(anthropic): forward Claude Code safeguards and dangerous-tool-use beta to Bedrock Invoke and Vertex on /v1/messages Claude Code's server-side auto-mode classifier sends a `safeguards` body field together with the `dangerous-tool-use-2026-09-03` beta. PR #42152 made the first-party anthropic route pass them through, but the beta header mapping left the other two Claude platforms at null, so Bedrock Invoke dropped both (classifier silently disabled) and Vertex forwarded the body field without the beta, which the platform rejects with "safeguards: Extra inputs are not permitted" (a 400 Claude Code hides by retrying without them). Map the beta for bedrock and vertex_ai in the beta headers config and add `safeguards` to the Bedrock Invoke request allowlist so the pair reaches both platforms unchanged. Nothing is injected: a client that sends `safeguards` without the beta still gets the platform's 400, exactly as api.anthropic.com answers it. --- litellm/anthropic_beta_headers_config.json | 6 ++ litellm/types/llms/bedrock.py | 1 + ...erimental_pass_through_messages_handler.py | 95 +++++++++++++++++++ .../test_anthropic_claude3_transformation.py | 43 +++++++++ .../test_anthropic_beta_headers_filtering.py | 14 +++ 5 files changed, 159 insertions(+) diff --git a/litellm/anthropic_beta_headers_config.json b/litellm/anthropic_beta_headers_config.json index eb31cc17a15..c4c60ae715b 100644 --- a/litellm/anthropic_beta_headers_config.json +++ b/litellm/anthropic_beta_headers_config.json @@ -11,6 +11,7 @@ "computer-use-2025-11-24": "computer-use-2025-11-24", "context-1m-2025-08-07": "context-1m-2025-08-07", "context-management-2025-06-27": "context-management-2025-06-27", + "dangerous-tool-use-2026-09-03": "dangerous-tool-use-2026-09-03", "effort-2025-11-24": "effort-2025-11-24", "fast-mode-2026-02-01": "fast-mode-2026-02-01", "files-api-2025-04-14": "files-api-2025-04-14", @@ -44,6 +45,7 @@ "computer-use-2025-11-24": "computer-use-2025-11-24", "context-1m-2025-08-07": "context-1m-2025-08-07", "context-management-2025-06-27": "context-management-2025-06-27", + "dangerous-tool-use-2026-09-03": null, "effort-2025-11-24": "effort-2025-11-24", "fast-mode-2026-02-01": null, "files-api-2025-04-14": "files-api-2025-04-14", @@ -76,6 +78,7 @@ "computer-use-2025-11-24": "computer-use-2025-11-24", "context-1m-2025-08-07": "context-1m-2025-08-07", "context-management-2025-06-27": null, + "dangerous-tool-use-2026-09-03": null, "effort-2025-11-24": "effort-2025-11-24", "fast-mode-2026-02-01": null, "files-api-2025-04-14": null, @@ -109,6 +112,7 @@ "computer-use-2025-11-24": "computer-use-2025-11-24", "context-1m-2025-08-07": "context-1m-2025-08-07", "context-management-2025-06-27": "context-management-2025-06-27", + "dangerous-tool-use-2026-09-03": "dangerous-tool-use-2026-09-03", "effort-2025-11-24": "effort-2025-11-24", "fast-mode-2026-02-01": null, "files-api-2025-04-14": null, @@ -142,6 +146,7 @@ "computer-use-2025-11-24": "computer-use-2025-11-24", "context-1m-2025-08-07": "context-1m-2025-08-07", "context-management-2025-06-27": "context-management-2025-06-27", + "dangerous-tool-use-2026-09-03": "dangerous-tool-use-2026-09-03", "effort-2025-11-24": null, "fast-mode-2026-02-01": null, "files-api-2025-04-14": null, @@ -175,6 +180,7 @@ "computer-use-2025-11-24": "computer-use-2025-11-24", "context-1m-2025-08-07": "context-1m-2025-08-07", "context-management-2025-06-27": "context-management-2025-06-27", + "dangerous-tool-use-2026-09-03": null, "effort-2025-11-24": "effort-2025-11-24", "fast-mode-2026-02-01": "fast-mode-2026-02-01", "files-api-2025-04-14": "files-api-2025-04-14", diff --git a/litellm/types/llms/bedrock.py b/litellm/types/llms/bedrock.py index 10082cf2373..f7518fefae4 100644 --- a/litellm/types/llms/bedrock.py +++ b/litellm/types/llms/bedrock.py @@ -1236,6 +1236,7 @@ class BedrockInvokeAnthropicMessagesRequest(TypedDict, total=False): thinking: dict metadata: dict output_config: dict + safeguards: list # `context_management` is allowed for Bedrock InvokeModel only when it # carries `compact_20260112` edits paired with the `compact-2026-01-12` diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py index e8bfcb86bf6..8d940e5efad 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py @@ -1544,3 +1544,98 @@ async def test_anthropic_messages_streaming_forwards_safeguards_and_keeps_safegu assert captured["body"]["safeguards"] == safeguards assert events[0]["message"]["safeguard_results"] == safeguard_results assert [e for e in events if e["type"] == "message_delta"][0]["delta"]["safeguard_results"] == safeguard_results + + +def _claude_code_auto_mode_request() -> tuple[list[dict[str, object]], list[dict[str, object]]]: + """Shapes are what Claude Code 2.1.278 sends and Bedrock Invoke / Vertex rawPredict return, captured 2026-09-21.""" + safeguards = [{"type": "dangerous_tool_use", "classifier_context": {"v": 1, "permission_mode": "auto"}}] + tool_verdicts = {"toolu_01": {"type": "evaluated", "outcome": "not_flagged"}} + safeguard_results = [{"type": "dangerous_tool_use", "status": {"type": "available", "tool_uses": tool_verdicts}}] + return safeguards, safeguard_results + + +def _upstream_answering_with(safeguard_results: list[dict[str, object]], captured: dict[str, object]) -> AsyncHTTPHandler: + def upstream_records_the_request(request: httpx.Request) -> httpx.Response: + captured["body"] = json.loads(request.content) + captured["anthropic-beta"] = request.headers.get("anthropic-beta") + return httpx.Response( + 200, + json={ + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5", + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 1, "output_tokens": 1}, + "safeguard_results": safeguard_results, + }, + request=request, + ) + + upstream = AsyncHTTPHandler() + upstream.client = httpx.AsyncClient(transport=httpx.MockTransport(upstream_records_the_request)) + return upstream + + +@pytest.mark.asyncio +async def test_anthropic_messages_forwards_safeguards_and_dangerous_tool_use_beta_to_bedrock_invoke( + local_beta_headers_config, +): + """Bedrock Invoke takes betas in the body's `anthropic_beta` and 400s on `safeguards` without the beta.""" + from litellm.llms.anthropic.experimental_pass_through.messages import handler + + safeguards, safeguard_results = _claude_code_auto_mode_request() + captured: dict[str, object] = {} + + response = await handler.anthropic_messages( + max_tokens=16, + messages=[{"role": "user", "content": "hi"}], + model="bedrock/us.anthropic.claude-sonnet-5", + custom_llm_provider="bedrock", + aws_access_key_id="test-access-key", + aws_secret_access_key="test-secret-key", + aws_region_name="us-east-1", + client=_upstream_answering_with(safeguard_results, captured), + safeguards=safeguards, + extra_headers={"anthropic-beta": "dangerous-tool-use-2026-09-03,interleaved-thinking-2025-05-14"}, + ) + + assert captured["body"]["safeguards"] == safeguards + assert captured["body"]["anthropic_beta"] == ["dangerous-tool-use-2026-09-03"] + assert response["safeguard_results"] == safeguard_results + + +@pytest.mark.asyncio +async def test_anthropic_messages_forwards_safeguards_and_dangerous_tool_use_beta_to_vertex( + local_beta_headers_config, +): + """Vertex rawPredict takes the beta as the `anthropic-beta` header and 400s on `safeguards` without it.""" + from litellm.llms.anthropic.experimental_pass_through.messages import handler + from litellm.llms.vertex_ai.vertex_llm_base import VertexBase + + safeguards, safeguard_results = _claude_code_auto_mode_request() + captured: dict[str, object] = {} + + with patch.object(VertexBase, "_ensure_access_token", return_value=("test-token", "test-project")): + response = await handler.anthropic_messages( + max_tokens=16, + messages=[{"role": "user", "content": "hi"}], + model="vertex_ai/claude-sonnet-5", + custom_llm_provider="vertex_ai", + vertex_project="test-project", + vertex_location="global", + vertex_credentials="{}", + client=_upstream_answering_with(safeguard_results, captured), + safeguards=safeguards, + extra_headers={"anthropic-beta": "dangerous-tool-use-2026-09-03,interleaved-thinking-2025-05-14"}, + ) + + assert captured["body"]["safeguards"] == safeguards + assert "anthropic_beta" not in captured["body"] + assert set(captured["anthropic-beta"].split(",")) == { + "dangerous-tool-use-2026-09-03", + "interleaved-thinking-2025-05-14", + } + assert response["safeguard_results"] == safeguard_results diff --git a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py index 7be005c0efe..e3c4a84cc60 100644 --- a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py @@ -1651,6 +1651,49 @@ def test_bedrock_messages_allowlist_filters_anthropic_only_fields(): assert set(result).issubset(cfg.BEDROCK_INVOKE_ALLOWED_TOP_LEVEL_FIELDS) +def test_bedrock_messages_forwards_safeguards_with_dangerous_tool_use_beta(local_beta_headers_config): + """ + Claude Code's server-side auto-mode classifier sends `safeguards` alongside the + dangerous-tool-use-2026-09-03 beta. Bedrock Invoke accepts the pair, answers + "safeguards: Extra inputs are not permitted" for the field alone, and returns + `safeguard_results: []` for the beta alone, so both must reach it unchanged. + """ + from litellm.types.router import GenericLiteLLMParams + + cfg = AmazonAnthropicClaudeMessagesConfig() + safeguards = [{"type": "dangerous_tool_use", "classifier_context": {"v": 1, "permission_mode": "auto"}}] + + result = cfg.transform_anthropic_messages_request( + model="us.anthropic.claude-sonnet-5", + messages=[{"role": "user", "content": [{"type": "text", "text": "Hello"}]}], + anthropic_messages_optional_request_params={"max_tokens": 64, "safeguards": safeguards}, + litellm_params=GenericLiteLLMParams(), + headers={"anthropic-beta": "dangerous-tool-use-2026-09-03,interleaved-thinking-2025-05-14"}, + ) + + assert result["safeguards"] == safeguards + assert result["anthropic_beta"] == ["dangerous-tool-use-2026-09-03"] + + +def test_bedrock_messages_stream_decoder_keeps_safeguard_results(): + """Bedrock streams the classifier verdicts on message_start and on the final message_delta, exactly as api.anthropic.com does.""" + decoder = AmazonAnthropicClaudeMessagesStreamDecoder(model="us.anthropic.claude-sonnet-5") + tool_verdicts = {"toolu_01": {"type": "evaluated", "outcome": "not_flagged"}} + safeguard_results = [{"type": "dangerous_tool_use", "status": {"type": "available", "tool_uses": tool_verdicts}}] + + message_delta = decoder._chunk_parser( + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None, "safeguard_results": safeguard_results}, + "usage": {"output_tokens": 1}, + "amazon-bedrock-invocationMetrics": {"inputTokenCount": 3, "outputTokenCount": 1}, + } + ) + + assert isinstance(message_delta, dict) + assert message_delta["delta"]["safeguard_results"] == safeguard_results + + def test_bedrock_messages_filters_user_provided_unsupported_beta_header(): """ In proxy deployments the client (e.g. Claude Code) doesn't know the backend diff --git a/tests/test_litellm/test_anthropic_beta_headers_filtering.py b/tests/test_litellm/test_anthropic_beta_headers_filtering.py index 3c967283abf..d600d2b734b 100644 --- a/tests/test_litellm/test_anthropic_beta_headers_filtering.py +++ b/tests/test_litellm/test_anthropic_beta_headers_filtering.py @@ -442,6 +442,20 @@ class TestAnthropicBetaHeadersFiltering: assert filtered == ["thinking-binding-controls-2026-08-01"] + @pytest.mark.parametrize("provider", ["anthropic", "bedrock", "vertex_ai"]) + def test_dangerous_tool_use_forwarded(self, provider): + """Claude Code's server-side auto-mode classifier sends `safeguards` together with + dangerous-tool-use-2026-09-03. Bedrock Invoke and Vertex rawPredict both answer + "safeguards: Extra inputs are not permitted" when the body field arrives without + the beta (probed 2026-09-21), so dropping the header turned every auto-mode turn + into a 400 on Vertex and silently disabled the classifier on Bedrock.""" + filtered = filter_and_transform_beta_headers( + beta_headers=["dangerous-tool-use-2026-09-03"], + provider=provider, + ) + + assert filtered == ["dangerous-tool-use-2026-09-03"] + def test_null_value_headers_filtered(self): """Test that headers with null values are always filtered out.""" for provider in [ From 10d343c3eee7472341c04e7ad5a1438f8eeb1a9d Mon Sep 17 00:00:00 2001 From: yuneng Date: Mon, 21 Sep 2026 19:06:13 +0000 Subject: [PATCH 06/12] fix(proxy): detach stored credential when model editor selects None (LIT-7597) Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../model_management_endpoints.py | 45 +- .../proxy/auth/test_model_checks.py | 86 ++++ .../test_model_management_endpoints.py | 399 ++++++++++++++++++ .../src/components/ModelInfoEditForm.tsx | 16 +- .../src/components/model_info_view.test.tsx | 113 ++++- .../src/components/model_info_view.tsx | 12 +- 6 files changed, 651 insertions(+), 20 deletions(-) diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index ea124776d0b..a8242904636 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -28,6 +28,7 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid from litellm.constants import LITELLM_PROXY_ADMIN_NAME +from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.ptu_pricing import ( CUSTOM_PRICING_FIELDS, PTU_EMPTIED_PRICING_FIELDS, @@ -145,7 +146,7 @@ if TYPE_CHECKING: from prisma import types as prisma_types router: Final = APIRouter() -CLEARABLE_LITELLM_PARAMS: Final = frozenset({"cache_control_injection_points"}) +CLEARABLE_LITELLM_PARAMS: Final = frozenset({"cache_control_injection_points", "litellm_credential_name"}) NULL_CLEARABLE_LITELLM_PARAMS: Final = frozenset((*SPECIAL_MODEL_INFO_PARAMS, *CLEARABLE_LITELLM_PARAMS)) @@ -332,6 +333,28 @@ def _raise_on_strategy_router_write_violation( ) +def _raise_on_invalid_credential_name(litellm_params: updateLiteLLMParams | None) -> None: + if litellm_params is None or "litellm_credential_name" not in litellm_params.model_fields_set: + return + credential_name: Final = litellm_params.litellm_credential_name + if credential_name is None: + return + if credential_name == "": + raise ProxyException( + message="litellm_credential_name cannot be an empty string. Send null to detach the stored credential or omit the field to leave it unchanged.", + type=ProxyErrorTypes.validation_error.value, + code=status.HTTP_400_BAD_REQUEST, + param="litellm_credential_name", + ) + if CredentialAccessor.find_credential(credential_name) is None: + raise ProxyException( + message=f"Credential '{credential_name}' not found. Create it via /credentials before attaching it to a model.", + type=ProxyErrorTypes.validation_error.value, + code=status.HTTP_400_BAD_REQUEST, + param="litellm_credential_name", + ) + + AUTO_ROUTER_CAPABILITY_SLOT_LOCK_KEY: Final = 5_872_301 _CAPABILITY_LOCK_SQL: Final = "SELECT 1 AS locked FROM pg_advisory_xact_lock($1)" _STORED_LITELLM_PARAMS_SQL: Final = ( @@ -1111,6 +1134,7 @@ async def patch_model( user_api_key_dict=user_api_key_dict, existing_litellm_params=db_model.litellm_params, ) + _raise_on_invalid_credential_name(patch_data.litellm_params) ModelManagementAuthChecks.can_user_set_aws_session_tags( litellm_params=patch_data.litellm_params, @@ -1921,21 +1945,28 @@ class ModelManagementAuthChecks: user_api_key_dict: UserAPIKeyAuth, existing_litellm_params: GenericLiteLLMParams | None = None, ) -> Literal[True]: - if litellm_params is None or litellm_params.litellm_credential_name is None: + if litellm_params is None: return True - if existing_litellm_params is not None and existing_litellm_params.litellm_credential_name is not None: - existing_credential_name: Final = decrypt_value_helper( + if "litellm_credential_name" not in litellm_params.model_fields_set: + return True + existing_credential_name: Final = ( + decrypt_value_helper( value=existing_litellm_params.litellm_credential_name, key="litellm_credential_name", exception_type="debug", return_original_value=True, ) - if litellm_params.litellm_credential_name == existing_credential_name: - return True + if existing_litellm_params is not None and existing_litellm_params.litellm_credential_name is not None + else None + ) + requested_credential_name: Final = litellm_params.litellm_credential_name + if requested_credential_name == existing_credential_name: + return True if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN: return True + action: Final = "detach" if requested_credential_name is None else "attach" raise ProxyException( - message=f"Only a proxy admin can attach a stored credential (litellm_credential_name) to a model. Your role={user_api_key_dict.user_role}.", + message=f"Only a proxy admin can {action} a stored credential (litellm_credential_name) on a model. Your role={user_api_key_dict.user_role}.", type=ProxyErrorTypes.auth_error.value, code=status.HTTP_403_FORBIDDEN, param="litellm_credential_name", diff --git a/tests/test_litellm/proxy/auth/test_model_checks.py b/tests/test_litellm/proxy/auth/test_model_checks.py index f10622e954b..13171a42cda 100644 --- a/tests/test_litellm/proxy/auth/test_model_checks.py +++ b/tests/test_litellm/proxy/auth/test_model_checks.py @@ -523,6 +523,92 @@ def test_wildcard_credential_hydration_preserves_missing_credential_name( } +def test_hydrate_credential_name_none_leaves_params_untouched(monkeypatch): + import litellm + from litellm.proxy.auth.model_checks import _hydrate_litellm_credential_name + from litellm.types.router import LiteLLM_Params + from litellm.types.utils import CredentialItem + + monkeypatch.setattr( + litellm, + "credential_list", + [ + CredentialItem( + credential_name="shared-credential", + credential_info={}, + credential_values={"api_key": "sk-shared"}, + ) + ], + ) + params = LiteLLM_Params(model="openai/gpt-4o", litellm_credential_name=None) + + result = _hydrate_litellm_credential_name(params) + + assert result is not None + assert result.api_key is None + assert result.litellm_credential_name is None + + +def test_hydrate_replaced_credential_uses_new_credential_values(monkeypatch): + import litellm + from litellm.proxy.auth.model_checks import _hydrate_litellm_credential_name + from litellm.types.router import LiteLLM_Params + from litellm.types.utils import CredentialItem + + monkeypatch.setattr( + litellm, + "credential_list", + [ + CredentialItem( + credential_name="shared-credential", + credential_info={}, + credential_values={"api_key": "sk-shared"}, + ), + CredentialItem( + credential_name="other-credential", + credential_info={}, + credential_values={"api_key": "sk-other"}, + ), + ], + ) + params = LiteLLM_Params(model="openai/gpt-4o", litellm_credential_name="other-credential") + + result = _hydrate_litellm_credential_name(params) + + assert result is not None + assert result.api_key == "sk-other" + assert result.litellm_credential_name is None + + +def test_hydrate_inline_api_key_wins_over_stored_credential(monkeypatch): + import litellm + from litellm.proxy.auth.model_checks import _hydrate_litellm_credential_name + from litellm.types.router import LiteLLM_Params + from litellm.types.utils import CredentialItem + + monkeypatch.setattr( + litellm, + "credential_list", + [ + CredentialItem( + credential_name="shared-credential", + credential_info={}, + credential_values={"api_key": "sk-shared"}, + ) + ], + ) + params = LiteLLM_Params( + model="openai/gpt-4o", + api_key="sk-inline", + litellm_credential_name="shared-credential", + ) + + result = _hydrate_litellm_credential_name(params) + + assert result is not None + assert result.api_key == "sk-inline" + + @pytest.mark.asyncio async def test_get_available_models_for_user_expands_query_team_wildcard( monkeypatch, diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py index 376309d8a7e..72db04dfbaa 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py @@ -308,6 +308,46 @@ class TestModelManagementAuthChecks: ) assert result is True + def test_can_user_attach_credential_non_admin_explicit_null_clear_fails(self): + from litellm.proxy._types import ProxyException + from litellm.types.router import updateLiteLLMParams as litellm_params + + with pytest.raises(ProxyException) as exc_info: + ModelManagementAuthChecks.can_user_attach_credential( + litellm_params=litellm_params(litellm_credential_name=None), + user_api_key_dict=self.team_admin_user, + existing_litellm_params=LiteLLM_Params( + model="test_model", litellm_credential_name="shared-credential" + ), + ) + + assert exc_info.value.code == "403" + assert exc_info.value.param == "litellm_credential_name" + + def test_can_user_attach_credential_admin_explicit_null_clear_succeeds(self): + from litellm.types.router import updateLiteLLMParams as litellm_params + + result = ModelManagementAuthChecks.can_user_attach_credential( + litellm_params=litellm_params(litellm_credential_name=None), + user_api_key_dict=self.admin_user, + existing_litellm_params=LiteLLM_Params( + model="test_model", litellm_credential_name="shared-credential" + ), + ) + + assert result is True + + def test_can_user_attach_credential_null_without_existing_allows_any_role(self): + from litellm.types.router import updateLiteLLMParams as litellm_params + + result = ModelManagementAuthChecks.can_user_attach_credential( + litellm_params=litellm_params(litellm_credential_name=None), + user_api_key_dict=self.team_admin_user, + existing_litellm_params=LiteLLM_Params(model="test_model"), + ) + + assert result is True + def test_can_user_attach_credential_unchanged_encrypted_existing_allows_any_role(self, monkeypatch): monkeypatch.setenv("LITELLM_SALT_KEY", "sk-1234") encrypted_name = encrypt_value_helper(value="shared-credential") @@ -4000,6 +4040,365 @@ class TestUpdateDBModelClearCacheControlInjectionPoints: assert params["tpm"] == 10 +class TestUpdateDBModelClearCredentialName: + def test_explicit_null_removes_stored_credential_name(self): + from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model + + db_model: Final = Deployment( + model_name="gpt-4", + litellm_params=LiteLLM_Params( + model="openai/gpt-4o", + api_base="https://api.openai.com/v1", + api_key="sk-real", + tpm=100, + litellm_credential_name="shared-credential", + ), + model_info=ModelInfo(id="dep-cred-1", team_id="team-keep", access_groups=["prod"]), + ) + update_patch: Final = updateDeployment( + litellm_params=updateLiteLLMParams(litellm_credential_name=None) + ) + + with patch("litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", side_effect=lambda value, **kwargs: value): + result: Final = update_db_model(db_model=db_model, updated_patch=update_patch) + + params: Final = json.loads(result["litellm_params"]) + info: Final = json.loads(result["model_info"]) + assert "litellm_credential_name" not in params + assert params["model"] == "openai/gpt-4o" + assert params["api_base"] == "https://api.openai.com/v1" + assert params["api_key"] == "sk-real" + assert params["tpm"] == 100 + assert info["team_id"] == "team-keep" + assert info["access_groups"] == ["prod"] + + def test_omitted_credential_name_keeps_stored_association(self): + from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model + + db_model: Final = Deployment( + model_name="gpt-4", + litellm_params=LiteLLM_Params( + model="openai/gpt-4o", + api_base="https://api.openai.com/v1", + api_key="sk-real", + tpm=100, + litellm_credential_name="shared-credential", + ), + model_info=ModelInfo(id="dep-cred-1", team_id="team-keep", access_groups=["prod"]), + ) + update_patch: Final = updateDeployment(litellm_params=updateLiteLLMParams(tpm=10)) + + with patch("litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", side_effect=lambda value, **kwargs: value): + result: Final = update_db_model(db_model=db_model, updated_patch=update_patch) + + params: Final = json.loads(result["litellm_params"]) + assert params["litellm_credential_name"] == "shared-credential" + assert params["tpm"] == 10 + + def test_null_clear_on_model_without_credential_is_noop(self): + from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model + + db_model: Final = Deployment( + model_name="gpt-4", + litellm_params=LiteLLM_Params(model="openai/gpt-4o", api_base="https://api.openai.com/v1"), + model_info=ModelInfo(id="dep-cred-1", team_id="team-keep", access_groups=["prod"]), + ) + update_patch: Final = updateDeployment( + litellm_params=updateLiteLLMParams(litellm_credential_name=None) + ) + + with patch("litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", side_effect=lambda value, **kwargs: value): + result: Final = update_db_model(db_model=db_model, updated_patch=update_patch) + + params: Final = json.loads(result["litellm_params"]) + assert "litellm_credential_name" not in params + assert params["model"] == "openai/gpt-4o" + assert params["api_base"] == "https://api.openai.com/v1" + + def test_null_credential_clear_alongside_pricing_clear(self): + from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model + + db_model: Final = Deployment( + model_name="gpt-4", + litellm_params=LiteLLM_Params( + model="openai/gpt-4o", + api_base="https://api.openai.com/v1", + input_cost_per_token=0.000001, + litellm_credential_name="shared-credential", + ), + model_info=ModelInfo(id="dep-cred-1", input_cost_per_token=0.000001), + ) + update_patch: Final = updateDeployment( + litellm_params=updateLiteLLMParams( + litellm_credential_name=None, + input_cost_per_token=None, + ) + ) + + with patch("litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", side_effect=lambda value, **kwargs: value): + result: Final = update_db_model(db_model=db_model, updated_patch=update_patch) + + params: Final = json.loads(result["litellm_params"]) + info: Final = json.loads(result["model_info"]) + assert "litellm_credential_name" not in params + assert "input_cost_per_token" not in params + assert "input_cost_per_token" not in info + + def test_replace_credential_name_keeps_other_params(self): + from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model + + db_model: Final = Deployment( + model_name="gpt-4", + litellm_params=LiteLLM_Params( + model="openai/gpt-4o", + api_base="https://api.openai.com/v1", + api_key="sk-real", + tpm=100, + litellm_credential_name="shared-credential", + ), + model_info=ModelInfo(id="dep-cred-1", team_id="team-keep", access_groups=["prod"]), + ) + update_patch: Final = updateDeployment( + litellm_params=updateLiteLLMParams(litellm_credential_name="other-credential") + ) + + with patch("litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", side_effect=lambda value, **kwargs: value): + result: Final = update_db_model(db_model=db_model, updated_patch=update_patch) + + params: Final = json.loads(result["litellm_params"]) + assert params["litellm_credential_name"] == "other-credential" + assert params["api_base"] == "https://api.openai.com/v1" + assert params["api_key"] == "sk-real" + assert params["tpm"] == 100 + + +class TestPatchModelCredentialName: + @staticmethod + async def _patch_model( + monkeypatch, + db_model: Deployment, + user_api_key_dict: UserAPIKeyAuth, + credential_name: str | None, + ) -> list[dict[str, object]]: + import litellm + from litellm.proxy.management_endpoints.model_management_endpoints import patch_model, update_db_model + from litellm.types.utils import CredentialItem + + monkeypatch.setattr( + litellm, + "credential_list", + [ + CredentialItem( + credential_name="shared-credential", + credential_info={}, + credential_values={"api_key": "sk-shared"}, + ), + CredentialItem( + credential_name="other-credential", + credential_info={}, + credential_values={"api_key": "sk-other"}, + ), + ], + ) + persisted: Final[list[dict[str, object]]] = [] + + async def persist_model(**kwargs): + row: Final = update_db_model(db_model=kwargs["db_model"], updated_patch=kwargs["patch_data"]) + persisted.append(row) + updated_row: Final = MagicMock() + updated_row.model_dump_json.return_value = "{}" + return updated_row + + with ( + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), + patch("litellm.proxy.proxy_server.llm_router", MagicMock()), + patch("litellm.proxy.proxy_server.store_model_in_db", True), + patch("litellm.proxy.proxy_server.premium_user", True), + patch( + "litellm.proxy.management_endpoints.model_management_endpoints.get_db_model", + new=AsyncMock(return_value=db_model), + ), + patch( + "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call", + new=AsyncMock(return_value=None), + ), + patch( + "litellm.proxy.management_endpoints.model_management_endpoints._update_team_model_in_db", + new=AsyncMock(side_effect=persist_model), + ), + patch( + "litellm.proxy.management_endpoints.model_management_endpoints.clear_cache", + new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)), + ), + patch( + "litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", + side_effect=lambda value, **kwargs: value, + ), + patch( + "litellm.proxy.management_endpoints.model_management_endpoints.raise_if_reload_degraded_serving" + ), + patch( + "litellm.proxy.management_endpoints.model_management_endpoints.create_object_audit_log", + new=AsyncMock(), + ), + patch( + "litellm.proxy.management_endpoints.model_management_endpoints.live_model_ids_snapshot", + return_value=frozenset(), + ), + ): + await patch_model( + model_id="dep-cred-1", + patch_data=updateDeployment( + litellm_params=updateLiteLLMParams(litellm_credential_name=credential_name) + ), + user_api_key_dict=user_api_key_dict, + ) + + return persisted + + @pytest.mark.asyncio + async def test_patch_model_rejects_empty_string_credential_name(self, monkeypatch): + from litellm.proxy._types import ProxyException + + db_model: Final = Deployment( + model_name="gpt-4", + litellm_params=LiteLLM_Params( + model="openai/gpt-4o", + api_base="https://api.openai.com/v1", + litellm_credential_name="shared-credential", + ), + model_info=ModelInfo(id="dep-cred-1"), + ) + + with pytest.raises(ProxyException) as exc_info: + await self._patch_model( + monkeypatch, + db_model, + self._admin_user(), + "", + ) + + assert exc_info.value.code == "400" + assert exc_info.value.param == "litellm_credential_name" + assert "empty" in exc_info.value.message.lower() + + @staticmethod + def _admin_user() -> UserAPIKeyAuth: + return UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + + @staticmethod + def _team_admin_user() -> UserAPIKeyAuth: + return UserAPIKeyAuth(user_id="team-admin", user_role=LitellmUserRoles.INTERNAL_USER, team_id="team-keep") + + @pytest.mark.asyncio + async def test_patch_model_rejects_unknown_credential_name(self, monkeypatch): + from litellm.proxy._types import ProxyException + + db_model: Final = Deployment( + model_name="gpt-4", + litellm_params=LiteLLM_Params( + model="openai/gpt-4o", + api_base="https://api.openai.com/v1", + litellm_credential_name="shared-credential", + ), + model_info=ModelInfo(id="dep-cred-1"), + ) + + with pytest.raises(ProxyException) as exc_info: + await self._patch_model( + monkeypatch, + db_model, + self._admin_user(), + "ghost-credential", + ) + + assert exc_info.value.code == "400" + assert "not found" in exc_info.value.message.lower() + + @pytest.mark.asyncio + async def test_patch_model_replaces_credential_name_and_preserves_other_params(self, monkeypatch): + db_model: Final = Deployment( + model_name="gpt-4", + litellm_params=LiteLLM_Params( + model="openai/gpt-4o", + api_base="https://api.openai.com/v1", + litellm_credential_name="shared-credential", + ), + model_info=ModelInfo(id="dep-cred-1"), + ) + + persisted: Final = await self._patch_model(monkeypatch, db_model, self._admin_user(), "other-credential") + params: Final = json.loads(persisted[0]["litellm_params"]) + assert params["litellm_credential_name"] == "other-credential" + assert params["api_base"] == "https://api.openai.com/v1" + + @pytest.mark.asyncio + async def test_patch_model_admin_null_clear_persists_without_credential(self, monkeypatch): + db_model: Final = Deployment( + model_name="gpt-4", + litellm_params=LiteLLM_Params( + model="openai/gpt-4o", + api_base="https://api.openai.com/v1", + litellm_credential_name="shared-credential", + ), + model_info=ModelInfo(id="dep-cred-1"), + ) + + persisted: Final = await self._patch_model(monkeypatch, db_model, self._admin_user(), None) + params: Final = json.loads(persisted[0]["litellm_params"]) + assert "litellm_credential_name" not in params + assert params["api_base"] == "https://api.openai.com/v1" + + @pytest.mark.asyncio + async def test_patch_model_rejects_non_admin_explicit_null_clear(self, monkeypatch): + from litellm.proxy._types import ProxyException + + db_model: Final = Deployment( + model_name="gpt-4", + litellm_params=LiteLLM_Params( + model="openai/gpt-4o", + api_base="https://api.openai.com/v1", + litellm_credential_name="shared-credential", + ), + model_info=ModelInfo(id="dep-cred-1"), + ) + + with pytest.raises(ProxyException) as exc_info: + await self._patch_model(monkeypatch, db_model, self._team_admin_user(), None) + + assert exc_info.value.code == "403" + assert exc_info.value.param == "litellm_credential_name" + + @pytest.mark.asyncio + async def test_patch_model_clear_then_reattach_round_trip(self, monkeypatch): + db_model: Final = Deployment( + model_name="gpt-4", + litellm_params=LiteLLM_Params( + model="openai/gpt-4o", + api_base="https://api.openai.com/v1", + litellm_credential_name="shared-credential", + ), + model_info=ModelInfo(id="dep-cred-1"), + ) + + cleared: Final = await self._patch_model(monkeypatch, db_model, self._admin_user(), None) + cleared_model: Final = Deployment.model_validate( + { + "model_name": db_model.model_name, + "litellm_params": json.loads(cleared[0]["litellm_params"]), + "model_info": json.loads(cleared[0]["model_info"]), + } + ) + reattached: Final = await self._patch_model( + monkeypatch, + cleared_model, + self._admin_user(), + "shared-credential", + ) + params: Final = json.loads(reattached[0]["litellm_params"]) + assert params["litellm_credential_name"] == "shared-credential" + + class TestGetModelInfoWithIdBlocked: """`ProxyConfig.get_model_info_with_id` must propagate the DB-level `blocked` column into the in-memory `model_info` dict so the router filter can read it.""" diff --git a/ui/litellm-dashboard/src/components/ModelInfoEditForm.tsx b/ui/litellm-dashboard/src/components/ModelInfoEditForm.tsx index d07e49e4712..b4fefe1d2c3 100644 --- a/ui/litellm-dashboard/src/components/ModelInfoEditForm.tsx +++ b/ui/litellm-dashboard/src/components/ModelInfoEditForm.tsx @@ -102,7 +102,7 @@ export interface ModelEditFormValues { vector_store_ids?: string[]; tags?: string[]; health_check_model?: string | null; - litellm_credential_name?: string; + litellm_credential_name?: string | null; litellm_extra_params?: string; model_info?: string; team_id?: string; @@ -139,7 +139,7 @@ const modelEditShape = { vector_store_ids: z.array(z.string()).optional(), tags: z.array(z.string()).optional(), health_check_model: z.string().nullish(), - litellm_credential_name: textish, + litellm_credential_name: z.string().nullish(), litellm_extra_params: textish, model_info: textish, team_id: textish, @@ -254,7 +254,7 @@ export const toModelEditFormValues = (localModelData: any, isWildcardModel: bool tags: Array.isArray(localModelData.litellm_params?.tags) ? localModelData.litellm_params.tags : [], // antd never mounted this field for a non-wildcard model, so the key must be absent, not null. ...(isWildcardModel ? { health_check_model: localModelData.model_info?.health_check_model } : {}), - litellm_credential_name: localModelData.litellm_params?.litellm_credential_name || "", + litellm_credential_name: localModelData.litellm_params?.litellm_credential_name ?? null, litellm_extra_params: JSON.stringify( Object.fromEntries( Object.entries(localModelData.litellm_params || {}).filter( @@ -635,8 +635,8 @@ const ModelInfoEditForm: React.FC = ({ {isEditing ? ( {({ id, value, onChange, onBlur }) => { - const items = [ - { value: "", label: "None" }, + const items: { value: string | null; label: string }[] = [ + { value: null, label: "None" }, ...credentialsList.map((credential) => ({ value: credential.credential_name, label: credential.credential_name, @@ -645,15 +645,15 @@ const ModelInfoEditForm: React.FC = ({ return (