From 484524b70b9cd28108547b201b2fa88e1fce27fb Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 16:04:13 -0700 Subject: [PATCH 01/15] fix(auth): fail closed when the team membership lookup hits a db outage --- litellm/proxy/auth/auth_checks.py | 28 ++++---- .../proxy/auth/test_auth_checks.py | 68 ++++++++++++++----- .../proxy/auth/test_resolvers_grants.py | 28 ++++++++ 3 files changed, 90 insertions(+), 34 deletions(-) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 61d2fa572a1..32ae9bc81df 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -2296,23 +2296,19 @@ async def _load_team_membership_on_cache_miss( parent_otel_span: Span | None, proxy_logging_obj: ProxyLogging | None, ) -> LiteLLM_TeamMembership | None: - try: - redis_cached: Final[object] = await user_api_key_cache.async_get_cache(key=cache_key) - redis_membership: Final = _membership_from_cached_payload(redis_cached) - if not isinstance(redis_membership, _TeamMembershipCacheMiss): - return redis_membership + redis_cached: Final[object] = await user_api_key_cache.async_get_cache(key=cache_key) + redis_membership: Final = _membership_from_cached_payload(redis_cached) + if not isinstance(redis_membership, _TeamMembershipCacheMiss): + return redis_membership - return await _fetch_team_membership_from_db( - user_id=user_id, - team_id=team_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - ) - except Exception: - verbose_proxy_logger.exception("Error getting team membership") - return None + return await _fetch_team_membership_from_db( + user_id=user_id, + team_id=team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) async def get_team_membership( diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 1ae986db23b..233660b634a 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -7198,7 +7198,7 @@ async def test_common_checks_skips_membership_load_when_no_check_reads_it(): @pytest.mark.asyncio -async def test_get_team_membership_db_error_returns_none_and_retries_next_call(): +async def test_get_team_membership_db_error_surfaces_and_retries_next_call(): from litellm.proxy.auth.auth_checks import get_team_membership from litellm.proxy.common_utils.user_api_key_cache import team_membership_reservation_cache_key @@ -7210,12 +7210,13 @@ async def test_get_team_membership_db_error_returns_none_and_retries_next_call() ) cache = UserApiKeyCache() - failed = await get_team_membership( - user_id="u-fail", - team_id="t-fail", - prisma_client=mock_prisma_client, - user_api_key_cache=cache, - ) + with pytest.raises(RuntimeError, match="db down"): + await get_team_membership( + user_id="u-fail", + team_id="t-fail", + prisma_client=mock_prisma_client, + user_api_key_cache=cache, + ) cached_after_failure = await cache.async_get_cache( key=team_membership_reservation_cache_key(user_id="u-fail", team_id="t-fail") ) @@ -7226,24 +7227,55 @@ async def test_get_team_membership_db_error_returns_none_and_retries_next_call() user_api_key_cache=cache, ) - assert failed is None assert cached_after_failure is None assert recovered is not None assert recovered.user_id == "u-fail" assert mock_prisma_client.db.litellm_teammembership.find_unique.await_count == 2 -@pytest.mark.asyncio -async def test_get_team_membership_string_prisma_client_returns_none(): - from litellm.proxy.auth.auth_checks import get_team_membership +class _UnreachableMembershipPrisma: + class db: + class litellm_teammembership: + @staticmethod + async def find_unique(where: dict[str, dict[str, str]], include: dict[str, bool]) -> None: + raise httpx.ConnectError("All connection attempts failed") - result = await get_team_membership( - user_id="u-str", - team_id="t-str", - prisma_client="hello-world", - user_api_key_cache=UserApiKeyCache(), - ) - assert result is None + +def _restricted_member_check_deps() -> dict[str, object]: + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.utils import ProxyLogging + + cache = UserApiKeyCache() + return { + "team_object": LiteLLM_TeamTable(team_id="team-outage", models=["claude-sonnet-5"]), + "valid_token": UserAPIKeyAuth(token="hashed-fake", user_id="bob", team_id="team-outage"), + "prisma_client": _UnreachableMembershipPrisma(), + "user_api_key_cache": cache, + "proxy_logging_obj": ProxyLogging(user_api_key_cache=cache), + } + + +@pytest.mark.asyncio +async def test_check_team_member_model_access_fails_closed_when_the_membership_read_hits_a_db_outage(): + """Regression: with the member's row uncached and the database unreachable, the loader used to swallow the + transport error and return None, which every check reads as "no per-member restriction", so a member + limited to other models got a 200. The outage must surface as the 503 the rest of auth answers with.""" + from litellm.proxy.auth.auth_checks import _check_team_member_model_access + from litellm.proxy.auth.auth_exception_handler import _as_proxy_exception + + with pytest.raises(httpx.ConnectError) as raised: + await _check_team_member_model_access( + model="claude-sonnet-5", llm_router=None, **_restricted_member_check_deps() + ) + + surfaced = _as_proxy_exception(raised.value) + assert (surfaced.code, surfaced.type) == ("503", ProxyErrorTypes.no_db_connection) + + +@pytest.mark.asyncio +async def test_check_team_member_budget_fails_closed_when_the_membership_read_hits_a_db_outage(): + with pytest.raises(httpx.ConnectError): + await _check_team_member_budget(user_object=None, **_restricted_member_check_deps()) @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/auth/test_resolvers_grants.py b/tests/test_litellm/proxy/auth/test_resolvers_grants.py index 3f6d943bf98..415d1191b5d 100644 --- a/tests/test_litellm/proxy/auth/test_resolvers_grants.py +++ b/tests/test_litellm/proxy/auth/test_resolvers_grants.py @@ -1,4 +1,5 @@ from fastapi import HTTPException +import httpx import pytest from litellm.proxy._types import ( @@ -8,6 +9,7 @@ from litellm.proxy._types import ( ProxyException, ) from litellm.proxy.auth.auth_checks import TeamNotFoundError, UserNotFoundError +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.auth.resolvers.grants import ( GrantResolver, LookupDegraded, @@ -172,6 +174,32 @@ async def test_resolve_identity_lets_loader_errors_surface(): await loaders.resolver().resolve_identity(UserLookup(user_id=USER_ID), team_id=None) +class _UnreachableMembershipPrisma: + class db: + class litellm_teammembership: + @staticmethod + async def find_unique(where: dict[str, dict[str, str]], include: dict[str, bool]) -> None: + raise httpx.ConnectError("All connection attempts failed") + + +async def test_resolve_marks_a_membership_read_that_hits_a_db_outage_as_degraded(): + """Regression: the real membership loader swallowed a database transport error into None, so this outcome + was ResolvedGrants with no membership, never LookupDegraded, and a member's own model or budget limits + silently dropped for the request.""" + loaders = _Loaders(user=_user(), team=_team()) + resolver = GrantResolver( + _UnreachableMembershipPrisma(), + UserApiKeyCache(), + load_user=loaders.load_user, + load_team=loaders.load_team, + ) + + outcome = await resolver.resolve(UserLookup(user_id=USER_ID), team_id=TEAM_ID) + + assert isinstance(outcome, LookupDegraded) + assert isinstance(outcome.error, httpx.ConnectError) + + def test_raise_public_maps_a_deleted_user_to_401(): with pytest.raises(ProxyException) as exc_info: raise_public(UserGone(user_id=USER_ID)) From b3b280d46381552b0624261b57e6911d08469fb9 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 19 Sep 2026 16:32:22 -0700 Subject: [PATCH 02/15] test(auth): model the membership row read in the fakes the loader now reaches --- tests/proxy_unit_tests/test_user_api_key_auth.py | 10 +++++++++- .../mcp_server/test_discoverable_endpoints.py | 4 +++- tests/test_litellm/proxy/auth/test_auth_checks.py | 3 --- tests/test_litellm/proxy/auth/test_resolvers_grants.py | 3 --- 4 files changed, 12 insertions(+), 8 deletions(-) diff --git a/tests/proxy_unit_tests/test_user_api_key_auth.py b/tests/proxy_unit_tests/test_user_api_key_auth.py index a8fce58c60b..f5e8d861d79 100644 --- a/tests/proxy_unit_tests/test_user_api_key_auth.py +++ b/tests/proxy_unit_tests/test_user_api_key_auth.py @@ -201,6 +201,14 @@ async def test_returned_user_api_key_auth(user_role, expected_role): assert new_obj.user_role == expected_role +class _NoMembershipRowPrisma: + class db: + class litellm_teammembership: + @staticmethod + async def find_unique(where: dict[str, dict[str, str]], include: dict[str, bool]) -> None: + return None + + @pytest.mark.parametrize("key_ownership", ["user_key", "team_key"]) @pytest.mark.asyncio async def test_aaauser_personal_budgets(key_ownership): @@ -253,7 +261,7 @@ async def test_aaauser_personal_budgets(key_ownership): setattr(litellm.proxy.proxy_server, "user_api_key_cache", user_api_key_cache) setattr(litellm.proxy.proxy_server, "master_key", "sk-1234") - setattr(litellm.proxy.proxy_server, "prisma_client", "hello-world") + setattr(litellm.proxy.proxy_server, "prisma_client", _NoMembershipRowPrisma()) request = Request(scope={"type": "http"}) request._url = URL(url="/chat/completions") diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index d7666f5e694..bed12892665 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -11515,7 +11515,9 @@ def jwt_oauth_identity(monkeypatch: pytest.MonkeyPatch) -> tuple["JWTHandler", " monkeypatch.setattr(proxy_server, "general_settings", {"enable_jwt_auth": True}) monkeypatch.setattr(proxy_server, "premium_user", True) monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) - monkeypatch.setattr(proxy_server, "prisma_client", MagicMock()) + prisma: Final = MagicMock() + prisma.db.litellm_teammembership.find_unique = AsyncMock(return_value=None) + monkeypatch.setattr(proxy_server, "prisma_client", prisma) return handler, signing_key diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 233660b634a..4f3e31fd70d 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -7257,9 +7257,6 @@ def _restricted_member_check_deps() -> dict[str, object]: @pytest.mark.asyncio async def test_check_team_member_model_access_fails_closed_when_the_membership_read_hits_a_db_outage(): - """Regression: with the member's row uncached and the database unreachable, the loader used to swallow the - transport error and return None, which every check reads as "no per-member restriction", so a member - limited to other models got a 200. The outage must surface as the 503 the rest of auth answers with.""" from litellm.proxy.auth.auth_checks import _check_team_member_model_access from litellm.proxy.auth.auth_exception_handler import _as_proxy_exception diff --git a/tests/test_litellm/proxy/auth/test_resolvers_grants.py b/tests/test_litellm/proxy/auth/test_resolvers_grants.py index 415d1191b5d..e61269ec1c7 100644 --- a/tests/test_litellm/proxy/auth/test_resolvers_grants.py +++ b/tests/test_litellm/proxy/auth/test_resolvers_grants.py @@ -183,9 +183,6 @@ class _UnreachableMembershipPrisma: async def test_resolve_marks_a_membership_read_that_hits_a_db_outage_as_degraded(): - """Regression: the real membership loader swallowed a database transport error into None, so this outcome - was ResolvedGrants with no membership, never LookupDegraded, and a member's own model or budget limits - silently dropped for the request.""" loaders = _Loaders(user=_user(), team=_team()) resolver = GrantResolver( _UnreachableMembershipPrisma(), From 45d22dc5e133dbfb6ee76edc79572f0c82f3ec40 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 21 Sep 2026 12:50:58 -0700 Subject: [PATCH 03/15] test(migrations): cover the release-to-release upgrade path The migration e2e harness only ever used one image: it seeded the database with the candidate build and then applied synthetic migrations on top. That proves the migration machinery (locking, crash recovery, legacy baselining, pooling) but never executes the real schema of release N against the real migrations of release N+1, which is the path operators actually run. Adds a baseline image alongside the candidate, so a test can seed with a published release and upgrade with the build under test. Suites: - test_upgrade.py: the candidate applies the pending release migrations, keys minted by the baseline release survive, and concurrent replicas upgrade a baseline database exactly once. - test_rolling_upgrade.py: a baseline replica keeps serving virtual-key auth while the candidate migrates underneath it, and both releases serve and resolve each other's keys during the overlap. This is the reported failure: a new column on LiteLLM_VerificationToken invalidates prepared plans on pods still running the old release, which the proxy reads whole-row, and auth starts failing until those pods leave service. - test_shaped_database.py: the upgrade completes and preserves rows on a populated spend log, rather than on the empty database every other migration test starts from. Every upgrade assertion is gated on the candidate having actually applied migrations the baseline had not, so a stale pin fails loudly instead of passing on an empty delta. CI adds two jobs to the migration_startup workflow. The baseline defaults to a committed release pin and is overridable per pipeline, matching how migration_candidate_image already works; only the upgrade jobs pull it. Verified against a real v1.101.0 -> v1.102.0 upgrade: 6 passed, with the baseline seeding 165 migrations and the candidate applying the 6 that landed between the two releases. --- .circleci/config.yml | 29 ++++- .circleci/scripts/run_migration_tests.py | 3 + tests/e2e/migrations/conftest.py | 29 +++++ tests/e2e/migrations/containers.py | 5 +- tests/e2e/migrations/test_rolling_upgrade.py | 55 +++++++++ tests/e2e/migrations/test_shaped_database.py | 43 +++++++ tests/e2e/migrations/test_upgrade.py | 47 ++++++++ tests/e2e/migrations/upgrade.py | 117 +++++++++++++++++++ 8 files changed, 326 insertions(+), 2 deletions(-) create mode 100644 tests/e2e/migrations/test_rolling_upgrade.py create mode 100644 tests/e2e/migrations/test_shaped_database.py create mode 100644 tests/e2e/migrations/test_upgrade.py create mode 100644 tests/e2e/migrations/upgrade.py diff --git a/.circleci/config.yml b/.circleci/config.yml index cc9aa7fe1c4..83acf0ac1c8 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -6,6 +6,9 @@ parameters: migration_candidate_image: type: string default: "" + migration_baseline_image: + type: string + default: "ghcr.io/berriai/litellm-database:v1.102.0" migration_source_sha: type: string default: "" @@ -2946,7 +2949,10 @@ jobs: parameters: suite: type: enum - enum: [startup, recovery, legacy] + enum: [startup, recovery, legacy, upgrade, shaped] + baseline: + type: boolean + default: false machine: image: ubuntu-2204:2024.04.1 resource_class: large @@ -2954,6 +2960,7 @@ jobs: environment: LITELLM_MIGRATION_TESTS: "1" LITELLM_MIGRATION_TEST_IMAGE: litellm-docker-database:ci + LITELLM_MIGRATION_BASELINE_IMAGE: << pipeline.parameters.migration_baseline_image >> MIGRATION_TEST_ADMIN_URL: postgresql://postgres:postgres@127.0.0.1:5432/postgres MIGRATION_TEST_CONTAINER_ADMIN_URL: postgresql://postgres:postgres@host.docker.internal:5432/postgres MIGRATION_TEST_OUTPUT: /tmp/migration-results @@ -2981,6 +2988,16 @@ jobs: - wait_for_service: url: tcp://localhost:5432 timeout: "60" + - when: + condition: << parameters.baseline >> + steps: + - run: + name: Pull the baseline release the upgrade starts from + environment: + BASELINE_IMAGE: << pipeline.parameters.migration_baseline_image >> + command: | + [[ "$BASELINE_IMAGE" =~ ^ghcr.io/berriai/[a-z0-9._/-]+(@sha256:[0-9a-f]{64}|:v[0-9][0-9a-z.-]*)$ ]] || exit 1 + docker pull "$BASELINE_IMAGE" - run: name: Run migration startup regressions environment: @@ -3188,6 +3205,16 @@ workflows: name: migration-legacy-and-pooling suite: legacy requires: [build_docker_database_image] + - migration_startup_tests: + name: migration-upgrade + suite: upgrade + baseline: true + requires: [build_docker_database_image] + - migration_startup_tests: + name: migration-upgrade-shaped + suite: shaped + baseline: true + requires: [build_docker_database_image] migration_startup_scheduled: triggers: - schedule: diff --git a/.circleci/scripts/run_migration_tests.py b/.circleci/scripts/run_migration_tests.py index 56029c406fb..5a73c54e3f5 100644 --- a/.circleci/scripts/run_migration_tests.py +++ b/.circleci/scripts/run_migration_tests.py @@ -13,6 +13,8 @@ SUITES: Final = { "startup": (("test_startup.py",), 12), "recovery": (("test_recovery.py",), 15), "legacy": (("test_legacy.py", "test_pooling.py"), 11), + "upgrade": (("test_upgrade.py", "test_rolling_upgrade.py"), 5), + "shaped": (("test_shaped_database.py",), 1), } @@ -93,6 +95,7 @@ def main() -> int: { **metadata, "suite": suite, + "baseline_image": os.environ.get("LITELLM_MIGRATION_BASELINE_IMAGE", ""), "expected_cases": expected, "passed": passed, "pytest_exit_code": result.returncode, diff --git a/tests/e2e/migrations/conftest.py b/tests/e2e/migrations/conftest.py index 735adeedbdb..b7604a4fdda 100644 --- a/tests/e2e/migrations/conftest.py +++ b/tests/e2e/migrations/conftest.py @@ -60,3 +60,32 @@ def containers(migration_image: str, tmp_path: Path, request: SubRequest) -> Con output: Final = Path(configured) / request.node.name if configured else tmp_path output.mkdir(parents=True, exist_ok=True) return Containers(migration_image, output) + + +@pytest.fixture(scope="session") +def baseline_image(tmp_path_factory: pytest.TempPathFactory) -> str: + configured: Final = os.environ.get("LITELLM_MIGRATION_BASELINE_IMAGE") + assert configured, "LITELLM_MIGRATION_BASELINE_IMAGE must name the released image the upgrade starts from" + image: Final = docker("image", "inspect", configured, "--format", "{{.Id}}") + assert image.startswith("sha256:"), "Unable to identify the baseline image" + output: Final = Path(os.environ.get("MIGRATION_TEST_OUTPUT", str(tmp_path_factory.getbasetemp()))) + output.mkdir(parents=True, exist_ok=True) + (output / "baseline-image.json").write_text(json.dumps({"requested": configured, "image_id": image})) + return image + + +@pytest.fixture(scope="session") +def baseline_template( + databases: Databases, baseline_image: str, tmp_path_factory: pytest.TempPathFactory +) -> Iterator[Database]: + output: Final = Path(os.environ.get("MIGRATION_TEST_OUTPUT", str(tmp_path_factory.getbasetemp()))) / "baseline-seed" + with databases.create() as database: + with Containers(baseline_image, output).start(database) as replica: + ready((replica,), database) + yield database + + +@pytest.fixture +def baseline_database(databases: Databases, baseline_template: Database) -> Iterator[Database]: + with databases.create(baseline_template) as database: + yield database diff --git a/tests/e2e/migrations/containers.py b/tests/e2e/migrations/containers.py index 0f5793b81dd..dd126b994d3 100644 --- a/tests/e2e/migrations/containers.py +++ b/tests/e2e/migrations/containers.py @@ -5,7 +5,7 @@ import subprocess import time from collections.abc import Callable, Generator, Mapping from contextlib import contextmanager -from dataclasses import dataclass +from dataclasses import dataclass, replace from pathlib import Path from typing import Final from uuid import uuid4 @@ -123,6 +123,9 @@ class Containers: image: str output: Path + def using(self, image: str) -> "Containers": + return replace(self, image=image) + @contextmanager def start( self, diff --git a/tests/e2e/migrations/test_rolling_upgrade.py b/tests/e2e/migrations/test_rolling_upgrade.py new file mode 100644 index 00000000000..4d60d20df4a --- /dev/null +++ b/tests/e2e/migrations/test_rolling_upgrade.py @@ -0,0 +1,55 @@ +from typing import Final + +import pytest + +from .containers import Containers, ready +from .database import Database +from .upgrade import ( + CACHED_PLAN, + assert_history_clean, + assert_upgraded, + auth_traffic, + confirm, + keep_serving, + migration_names, + provision, +) + +pytestmark: Final = [pytest.mark.e2e, pytest.mark.migration_startup] + + +class TestRollingUpgrade: + def test_baseline_replica_keeps_serving_while_the_candidate_migrates( + self, containers: Containers, baseline_image: str, baseline_database: Database + ) -> None: + with containers.using(baseline_image).start(baseline_database) as old: + ready((old,), baseline_database) + key, _ = provision(old) + before: Final = migration_names(baseline_database) + with auth_traffic(old, key) as traffic: + keep_serving(traffic, "the baseline replica authenticating before the upgrade") + with containers.start(baseline_database) as new: + ready((new,), baseline_database) + assert_upgraded(before, migration_names(baseline_database)) + keep_serving(traffic, "the baseline replica authenticating after the schema moved") + assert_history_clean(baseline_database) + assert CACHED_PLAN not in old.logs(), "The baseline replica hit a stale prepared statement" + assert old.state().Running, "The baseline replica died during the upgrade" + + def test_both_releases_serve_and_share_keys_during_the_overlap( + self, containers: Containers, baseline_image: str, baseline_database: Database + ) -> None: + with containers.using(baseline_image).start(baseline_database) as old: + ready((old,), baseline_database) + old_key, old_alias = provision(old) + before: Final = migration_names(baseline_database) + with containers.start(baseline_database) as new: + ready((new,), baseline_database) + assert_upgraded(before, migration_names(baseline_database)) + new_key, new_alias = provision(new) + with auth_traffic(old, old_key) as old_traffic, auth_traffic(new, new_key) as new_traffic: + keep_serving(old_traffic, "the baseline replica serving through the overlap") + keep_serving(new_traffic, "the candidate replica serving through the overlap") + confirm(old, new_key, new_alias) + confirm(new, old_key, old_alias) + assert CACHED_PLAN not in old.logs(), "The baseline replica hit a stale prepared statement" diff --git a/tests/e2e/migrations/test_shaped_database.py b/tests/e2e/migrations/test_shaped_database.py new file mode 100644 index 00000000000..20c4368ae33 --- /dev/null +++ b/tests/e2e/migrations/test_shaped_database.py @@ -0,0 +1,43 @@ +from typing import Final + +import pytest + +from .containers import Containers, ready +from .database import Database +from .upgrade import assert_history_clean, assert_upgraded, confirm, migration_names, provision + +SPEND_ROWS: Final = 20_000 + +pytestmark: Final = [pytest.mark.e2e, pytest.mark.migration_startup] + + +def seed_spend_logs(database: Database, rows: int) -> None: + database.execute( + 'INSERT INTO "LiteLLM_SpendLogs" (request_id, call_type, "startTime", "endTime") ' + "SELECT 'upgrade-shape-' || g, 'acompletion', now() - (g || ' seconds')::interval, " + "now() - (g || ' seconds')::interval FROM generate_series(1, %s) AS g", + (rows,), + ) + assert database.query('SELECT count(*) FROM "LiteLLM_SpendLogs"') == ((rows,),) + + +class TestPopulatedDatabaseUpgrade: + def test_upgrade_completes_and_preserves_a_populated_spend_log( + self, containers: Containers, baseline_image: str, baseline_database: Database + ) -> None: + with containers.using(baseline_image).start(baseline_database) as old: + ready((old,), baseline_database) + key, alias = provision(old) + seed_spend_logs(baseline_database, SPEND_ROWS) + before: Final = migration_names(baseline_database) + with containers.start(baseline_database) as new: + ready((new,), baseline_database) + assert_upgraded(before, migration_names(baseline_database)) + confirm(new, key, alias) + assert_history_clean(baseline_database) + assert baseline_database.query('SELECT count(*) FROM "LiteLLM_SpendLogs"') == ((SPEND_ROWS,),), ( + "The upgrade lost spend rows" + ) + assert baseline_database.query( + 'SELECT count(*) FROM "LiteLLM_SpendLogs" WHERE "startTime" IS NULL OR "endTime" IS NULL' + ) == ((0,),), "The upgrade nulled timestamps on existing spend rows" diff --git a/tests/e2e/migrations/test_upgrade.py b/tests/e2e/migrations/test_upgrade.py new file mode 100644 index 00000000000..23f0bbe9124 --- /dev/null +++ b/tests/e2e/migrations/test_upgrade.py @@ -0,0 +1,47 @@ +from contextlib import ExitStack +from typing import Final + +import pytest + +from .checks import start_replicas +from .containers import Containers, ready +from .database import Database +from .upgrade import assert_history_clean, assert_upgraded, confirm, migration_names, provision + +pytestmark: Final = [pytest.mark.e2e, pytest.mark.migration_startup] + + +class TestReleaseUpgrade: + def test_candidate_applies_the_pending_release_migrations( + self, containers: Containers, baseline_database: Database + ) -> None: + before: Final = migration_names(baseline_database) + with containers.start(baseline_database) as replica: + ready((replica,), baseline_database) + assert_upgraded(before, migration_names(baseline_database)) + assert_history_clean(baseline_database) + + def test_upgrade_preserves_keys_minted_by_the_baseline_release( + self, containers: Containers, baseline_image: str, baseline_database: Database + ) -> None: + with containers.using(baseline_image).start(baseline_database) as old: + ready((old,), baseline_database) + key, alias = provision(old) + confirm(old, key, alias) + before: Final = migration_names(baseline_database) + with containers.start(baseline_database) as new: + ready((new,), baseline_database) + assert_upgraded(before, migration_names(baseline_database)) + confirm(new, key, alias) + + def test_concurrent_replicas_upgrade_a_baseline_database_once( + self, containers: Containers, baseline_database: Database + ) -> None: + before: Final = migration_names(baseline_database) + with ExitStack() as stack: + ready(start_replicas(stack, containers, baseline_database), baseline_database) + assert_upgraded(before, migration_names(baseline_database)) + assert_history_clean(baseline_database) + assert baseline_database.query("SELECT count(*) FROM _prisma_migrations WHERE applied_steps_count > 1") == ( + (0,), + ), "A migration was executed more than once across the upgrading replicas" diff --git a/tests/e2e/migrations/upgrade.py b/tests/e2e/migrations/upgrade.py new file mode 100644 index 00000000000..2758d0fb974 --- /dev/null +++ b/tests/e2e/migrations/upgrade.py @@ -0,0 +1,117 @@ +from __future__ import annotations + +import threading +from collections.abc import Generator +from contextlib import contextmanager +from dataclasses import dataclass, field +from typing import Final +from uuid import uuid4 + +from e2e_http import Result, Success, unwrap +from models import ( + KeyGenerateBody, + KeyGenerateResponse, + KeyInfoParams, + KeyInfoResponse, + ModelsListParams, + ModelsListResponse, +) +from pydantic import BaseModel + +from .containers import Replica, until +from .database import Database + +CACHED_PLAN: Final = "cached plan must not change result type" + + +def provision(replica: Replica) -> tuple[str, str]: + alias: Final = f"upgrade-{uuid4().hex}" + key: Final = unwrap( + replica.transport.post( + "/key/generate", + headers=replica.transport.master, + json=KeyGenerateBody(key_alias=alias), + response_type=KeyGenerateResponse, + ) + ).key + return key, alias + + +def confirm(replica: Replica, key: str, alias: str) -> None: + info: Final = unwrap( + replica.transport.get( + "/key/info", + headers=replica.transport.master, + params=KeyInfoParams(key=key), + response_type=KeyInfoResponse, + ) + ) + assert info.info.key_alias == alias, "Key minted on one release did not resolve on the other" + + +@dataclass(slots=True) +class Outcomes: + served: int = 0 + failures: list[str] = field(default_factory=list) + + def record(self, result: Result[BaseModel]) -> None: + match result: + case Success(): + self.served += 1 + case _: + self.failures.append(result.model_dump_json()) + + +@contextmanager +def auth_traffic(replica: Replica, key: str, interval: float = 0.05) -> Generator[Outcomes]: + outcomes: Final = Outcomes() + stop: Final = threading.Event() + + def drive() -> None: + while not stop.is_set(): + outcomes.record( + replica.transport.get( + "/v1/models", + headers=replica.transport.bearer(key), + params=ModelsListParams(), + response_type=ModelsListResponse, + timeout=10, + ) + ) + stop.wait(interval) + + thread: Final = threading.Thread(target=drive, name="upgrade-auth-traffic", daemon=True) + thread.start() + try: + yield outcomes + finally: + stop.set() + thread.join(30) + assert not thread.is_alive(), "Auth traffic thread did not stop" + + +def keep_serving(outcomes: Outcomes, description: str, calls: int = 20) -> int: + target: Final = outcomes.served + calls + until(description, lambda: outcomes.served >= target or bool(outcomes.failures)) + assert not outcomes.failures, f"Virtual-key auth failed during {description}: {outcomes.failures[:5]}" + return outcomes.served + + +def migration_names(database: Database) -> frozenset[str]: + return frozenset(str(row[0]) for row in database.query("SELECT migration_name FROM _prisma_migrations")) + + +def assert_history_clean(database: Database) -> None: + assert database.query( + "SELECT count(*) FROM _prisma_migrations WHERE finished_at IS NULL OR rolled_back_at IS NOT NULL" + ) == ((0,),), "The upgrade left an unfinished or rolled-back migration behind" + + +def assert_upgraded(before: frozenset[str], after: frozenset[str]) -> frozenset[str]: + applied: Final = after - before + assert applied, ( + "The candidate applied no migrations the baseline release had not: the pinned " + "LITELLM_MIGRATION_BASELINE_IMAGE is at or ahead of the candidate, so this suite proves nothing" + ) + assert not before - after, "The upgrade removed migration history the baseline release had already applied" + return applied From dc85812971ef1ba5a143f9ee87a8ef5f395852dd Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Mon, 21 Sep 2026 13:15:46 -0700 Subject: [PATCH 04/15] test(migrations): close the gaps the upgrade assertions left open Three holes in the new suite, all of which let a test pass without proving what its name claims: - A migration recorded twice, once per replica, each with applied_steps_count = 1, slipped past both the step-count check and migration_names(), which collapses the history into a set. Reject duplicate migration_name rows outright. - auth_traffic only asserted the failures it had seen by the time keep_serving hit its target. A request failing after that, or on the other replica while the test waited on one stream, was recorded and never read. Assert the recorded failures once the thread has joined. - The rolling test warmed the baseline replica's virtual-key cache before the upgrade, and that cache holds for 60 seconds by default (UserAPIKeyCacheTTLEnum.in_memory_cache_ttl). The candidate migrates well inside that window, so the post-upgrade requests could be served from cache without ever repeating the whole-row token lookup that the stale prepared statement breaks. Drive the baseline replica with a key minted after the schema moved, which it has never seen and must resolve from the database. Re-ran against v1.101.0 -> v1.102.0: 6 passed. --- tests/e2e/migrations/test_rolling_upgrade.py | 2 ++ tests/e2e/migrations/upgrade.py | 7 +++++++ 2 files changed, 9 insertions(+) diff --git a/tests/e2e/migrations/test_rolling_upgrade.py b/tests/e2e/migrations/test_rolling_upgrade.py index 4d60d20df4a..5ad74e0ba8c 100644 --- a/tests/e2e/migrations/test_rolling_upgrade.py +++ b/tests/e2e/migrations/test_rolling_upgrade.py @@ -32,6 +32,8 @@ class TestRollingUpgrade: ready((new,), baseline_database) assert_upgraded(before, migration_names(baseline_database)) keep_serving(traffic, "the baseline replica authenticating after the schema moved") + with auth_traffic(old, provision(new)[0]) as uncached: + keep_serving(uncached, "the baseline replica resolving a key minted after the schema moved") assert_history_clean(baseline_database) assert CACHED_PLAN not in old.logs(), "The baseline replica hit a stale prepared statement" assert old.state().Running, "The baseline replica died during the upgrade" diff --git a/tests/e2e/migrations/upgrade.py b/tests/e2e/migrations/upgrade.py index 2758d0fb974..2123f86450a 100644 --- a/tests/e2e/migrations/upgrade.py +++ b/tests/e2e/migrations/upgrade.py @@ -88,6 +88,9 @@ def auth_traffic(replica: Replica, key: str, interval: float = 0.05) -> Generato stop.set() thread.join(30) assert not thread.is_alive(), "Auth traffic thread did not stop" + assert not outcomes.failures, ( + f"Virtual-key auth failed on {replica.name} after the traffic window closed: {outcomes.failures[:5]}" + ) def keep_serving(outcomes: Outcomes, description: str, calls: int = 20) -> int: @@ -105,6 +108,10 @@ def assert_history_clean(database: Database) -> None: assert database.query( "SELECT count(*) FROM _prisma_migrations WHERE finished_at IS NULL OR rolled_back_at IS NOT NULL" ) == ((0,),), "The upgrade left an unfinished or rolled-back migration behind" + assert database.query( + "SELECT count(*) FROM (SELECT migration_name FROM _prisma_migrations GROUP BY migration_name " + "HAVING count(*) > 1) duplicated" + ) == ((0,),), "A migration was recorded more than once, so it ran on more than one replica" def assert_upgraded(before: frozenset[str], after: frozenset[str]) -> frozenset[str]: From ae69a8c79a82d54d50436c1eb06ab9cca3625574 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:30:48 +0000 Subject: [PATCH 05/15] feat(rust): add Azure Key Vault secret manager backend Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/workflows/test-rust.yml | 2 +- litellm-rust/Cargo.lock | 21 ++ litellm-rust/Cargo.toml | 1 + litellm-rust/crates/secrets-azure/Cargo.toml | 24 ++ .../crates/secrets-azure/src/error.rs | 25 ++ .../crates/secrets-azure/src/key_vault.rs | 122 ++++++++++ litellm-rust/crates/secrets-azure/src/lib.rs | 7 + .../tests/fixtures/key_vault_parity.json | 8 + .../crates/secrets-azure/tests/key_vault.rs | 220 ++++++++++++++++++ .../crates/secrets-azure/tests/live.rs | 30 +++ litellm-rust/crates/secrets/Cargo.toml | 2 + litellm-rust/crates/secrets/src/error.rs | 3 + litellm-rust/crates/secrets/src/handler.rs | 9 + litellm-rust/crates/secrets/src/lib.rs | 2 + litellm-rust/crates/secrets/tests/handler.rs | 61 +++++ .../test_secret_manager_handler.py | 105 +++++++++ 16 files changed, 641 insertions(+), 1 deletion(-) create mode 100644 litellm-rust/crates/secrets-azure/Cargo.toml create mode 100644 litellm-rust/crates/secrets-azure/src/error.rs create mode 100644 litellm-rust/crates/secrets-azure/src/key_vault.rs create mode 100644 litellm-rust/crates/secrets-azure/src/lib.rs create mode 100644 litellm-rust/crates/secrets-azure/tests/fixtures/key_vault_parity.json create mode 100644 litellm-rust/crates/secrets-azure/tests/key_vault.rs create mode 100644 litellm-rust/crates/secrets-azure/tests/live.rs create mode 100644 tests/test_litellm/secret_managers/test_secret_manager_handler.py diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml index 278fa7c425f..e7cb9984676 100644 --- a/.github/workflows/test-rust.yml +++ b/.github/workflows/test-rust.yml @@ -130,7 +130,7 @@ jobs: - name: Test secret manager feature combinations run: | cargo test -p litellm-auth-gcp --locked --no-default-features - for features in '' aws google aws,google; do + for features in '' aws google azure aws,google aws,azure google,azure aws,google,azure; do cargo test -p litellm-secrets --locked --no-default-features --features "$features" done diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index ed4ae4e3353..f7b1e371c0b 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2697,6 +2697,7 @@ dependencies = [ "jsonwebtoken", "litellm-core-utils", "litellm-secrets-aws", + "litellm-secrets-azure", "litellm-secrets-google", "litellm-secrets-types", "moka", @@ -2731,6 +2732,26 @@ dependencies = [ "wiremock", ] +[[package]] +name = "litellm-secrets-azure" +version = "0.1.0" +dependencies = [ + "litellm-auth-azure", + "litellm-auth-types", + "litellm-core-utils", + "litellm-secrets-types", + "percent-encoding", + "reqwest 0.12.28", + "rstest", + "serde", + "serde_json", + "sha2 0.10.9", + "thiserror 2.0.19", + "tokio", + "veil", + "wiremock", +] + [[package]] name = "litellm-secrets-google" version = "0.1.0" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 570d0dd3568..802a29898d3 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -22,6 +22,7 @@ litellm-secrets = { path = "crates/secrets" } litellm-secrets-types = { path = "crates/secrets-types" } litellm-secrets-aws = { path = "crates/secrets-aws" } litellm-secrets-google = { path = "crates/secrets-google" } +litellm-secrets-azure = { path = "crates/secrets-azure" } litellm-http = { path = "crates/http" } litellm-llms = { path = "crates/llms" } litellm-types = { path = "crates/types" } diff --git a/litellm-rust/crates/secrets-azure/Cargo.toml b/litellm-rust/crates/secrets-azure/Cargo.toml new file mode 100644 index 00000000000..d559592cb34 --- /dev/null +++ b/litellm-rust/crates/secrets-azure/Cargo.toml @@ -0,0 +1,24 @@ +[package] +name = "litellm-secrets-azure" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +litellm-auth-azure.workspace = true +litellm-auth-types.workspace = true +litellm-secrets-types.workspace = true +litellm-core-utils.workspace = true +reqwest.workspace = true +serde.workspace = true +serde_json.workspace = true +thiserror.workspace = true +veil.workspace = true +percent-encoding = "2.3" + +[dev-dependencies] +tokio.workspace = true +wiremock = "0.6.5" +rstest.workspace = true +sha2.workspace = true diff --git a/litellm-rust/crates/secrets-azure/src/error.rs b/litellm-rust/crates/secrets-azure/src/error.rs new file mode 100644 index 00000000000..9b20efe4f7c --- /dev/null +++ b/litellm-rust/crates/secrets-azure/src/error.rs @@ -0,0 +1,25 @@ +#[derive(thiserror::Error, veil::Redact)] +pub enum Error { + #[error("{0} environment variable is missing")] + MissingEnvironment(&'static str), + #[error("AZURE_KEY_VAULT_URI is not a valid https vault URL")] + VaultUri, + #[error("Azure Key Vault credentials are not configured")] + MissingCredentials, + #[error(transparent)] + Auth( + #[from] + #[redact] + litellm_auth_types::Error, + ), + #[error("Azure Key Vault request failed")] + Http( + #[source] + #[redact] + reqwest::Error, + ), + #[error("Azure Key Vault returned HTTP {0}")] + Status(u16), + #[error("Azure Key Vault response is missing the secret value")] + MissingValue, +} diff --git a/litellm-rust/crates/secrets-azure/src/key_vault.rs b/litellm-rust/crates/secrets-azure/src/key_vault.rs new file mode 100644 index 00000000000..262b19386df --- /dev/null +++ b/litellm-rust/crates/secrets-azure/src/key_vault.rs @@ -0,0 +1,122 @@ +use std::sync::Arc; + +use litellm_auth_azure::{AzureAuthInputs, AzureAuthService}; +use litellm_core_utils::settings::Lookup; +use litellm_secrets_types::{Secret, SecretValue}; +use serde::Deserialize; + +use crate::Error; + +const AZURE_KEY_VAULT_URI: &str = "AZURE_KEY_VAULT_URI"; +const API_VERSION: &str = "7.4"; + +#[derive(Clone)] +pub struct AzureKeyVault { + client: reqwest::Client, + vault: reqwest::Url, + auth: Arc, + inputs: AzureAuthInputs, + environment: Arc, +} + +#[derive(Deserialize)] +struct SecretResponse { + value: Option, +} + +impl AzureKeyVault { + pub fn with_client( + client: reqwest::Client, + vault: reqwest::Url, + environment: Arc, + ) -> Result { + if vault.host_str().is_none() { + return Err(Error::VaultUri); + } + let scope = scope_for(&vault); + let inputs = AzureAuthInputs::from_sourced_optional_params( + serde_json::json!({ + "azure_scope": scope, + "enable_azure_ad_token_refresh": true, + }) + .as_object() + .expect("static Azure auth inputs object"), + &std::collections::BTreeMap::new(), + )?; + Ok(Self { + client, + vault, + auth: Arc::new(AzureAuthService::default()), + inputs, + environment, + }) + } + + pub fn new(environment: Arc) -> Result { + let value = environment + .get(AZURE_KEY_VAULT_URI) + .ok_or(Error::MissingEnvironment(AZURE_KEY_VAULT_URI))?; + let vault = reqwest::Url::parse(&value).map_err(|_| Error::VaultUri)?; + if vault.scheme() != "https" || vault.host_str().is_none() { + return Err(Error::VaultUri); + } + Self::with_client(reqwest::Client::new(), vault, environment) + } + + pub fn scope(&self) -> &str { + self.inputs + .azure_scope + .as_value() + .map(|value| value.value().as_str()) + .unwrap_or_default() + } + + pub async fn get_secret_from_azure_key_vault( + &self, + name: &str, + ) -> Result, Error> { + let token = self + .auth + .get_azure_ad_token(&self.inputs, &|key| self.environment.get(key)) + .await? + .ok_or(Error::MissingCredentials)?; + let encoded_name = encode_name(name); + let url = self + .vault + .join(&format!("secrets/{encoded_name}?api-version={API_VERSION}")) + .map_err(|_| Error::VaultUri)?; + let response = self + .client + .get(url) + .bearer_auth(token.value().secret().expose()) + .send() + .await + .map_err(Error::Http)?; + if response.status() == reqwest::StatusCode::NOT_FOUND { + return Ok(None); + } + if response.status() != reqwest::StatusCode::OK { + return Err(Error::Status(response.status().as_u16())); + } + let payload: SecretResponse = response.json().await.map_err(Error::Http)?; + let value = payload.value.ok_or(Error::MissingValue)?; + Ok(Some(Secret::String(SecretValue::new(value)))) + } +} + +fn scope_for(vault: &reqwest::Url) -> String { + let host = vault.host_str().unwrap_or_default(); + let resource = host + .split_once('.') + .map_or(host, |(_, remainder)| remainder); + format!("https://{resource}/.default") +} + +fn encode_name(name: &str) -> String { + percent_encoding::utf8_percent_encode(name, percent_encoding::NON_ALPHANUMERIC) + .to_string() + .replace("%2D", "-") + .replace("%2E", ".") + .replace("%5F", "_") + .replace("%7E", "~") +} diff --git a/litellm-rust/crates/secrets-azure/src/lib.rs b/litellm-rust/crates/secrets-azure/src/lib.rs new file mode 100644 index 00000000000..c0094fc033b --- /dev/null +++ b/litellm-rust/crates/secrets-azure/src/lib.rs @@ -0,0 +1,7 @@ +#![forbid(unsafe_code)] + +mod error; +mod key_vault; + +pub use error::Error; +pub use key_vault::AzureKeyVault; diff --git a/litellm-rust/crates/secrets-azure/tests/fixtures/key_vault_parity.json b/litellm-rust/crates/secrets-azure/tests/fixtures/key_vault_parity.json new file mode 100644 index 00000000000..c4a83cd150a --- /dev/null +++ b/litellm-rust/crates/secrets-azure/tests/fixtures/key_vault_parity.json @@ -0,0 +1,8 @@ +{ + "cases": [ + {"name": "plain_value", "secret_name": "OPENAI-API-KEY", "response": {"status": 200, "body": {"value": "sk-parity-1", "id": "https://example.vault.azure.net/secrets/OPENAI-API-KEY/abc"}}, "expected": {"value": "sk-parity-1"}}, + {"name": "json_value_is_kept_as_string", "secret_name": "JSON-SECRET", "response": {"status": 200, "body": {"value": "{\"api_key\": \"nested\"}", "id": "https://example.vault.azure.net/secrets/JSON-SECRET/abc"}}, "expected": {"value": "{\"api_key\": \"nested\"}"}}, + {"name": "missing_secret", "secret_name": "MISSING", "response": {"status": 404, "body": {"error": {"code": "SecretNotFound", "message": "not found"}}}, "expected": {"missing": true}}, + {"name": "forbidden", "secret_name": "FORBIDDEN", "response": {"status": 403, "body": {"error": {"code": "Forbidden", "message": "denied"}}}, "expected": {"error": true}} + ] +} diff --git a/litellm-rust/crates/secrets-azure/tests/key_vault.rs b/litellm-rust/crates/secrets-azure/tests/key_vault.rs new file mode 100644 index 00000000000..aae09369243 --- /dev/null +++ b/litellm-rust/crates/secrets-azure/tests/key_vault.rs @@ -0,0 +1,220 @@ +use std::sync::Arc; + +use litellm_secrets_azure::{AzureKeyVault, Error}; +use litellm_secrets_types::{Secret, SecretValue}; +use serde::Deserialize; +use wiremock::{ + Mock, MockServer, ResponseTemplate, + matchers::{header, path, query_param}, +}; + +fn manager(server: &MockServer) -> AzureKeyVault { + AzureKeyVault::with_client( + reqwest::Client::new(), + server.uri().parse().unwrap(), + Arc::new(|name: &str| (name == "AZURE_AD_TOKEN").then(|| "fake".to_owned())), + ) + .unwrap() +} + +#[tokio::test] +async fn reads_secret_with_bearer_token_and_api_version() { + let server = MockServer::start().await; + Mock::given(path("/secrets/OPENAI-API-KEY")) + .and(query_param("api-version", "7.4")) + .and(header("authorization", "Bearer fake")) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(serde_json::json!({"value": "s3cret", "id": "secret-id"})), + ) + .expect(1) + .mount(&server) + .await; + + let secret = manager(&server) + .get_secret_from_azure_key_vault("OPENAI-API-KEY") + .await + .unwrap() + .unwrap(); + + assert_eq!(secret, Secret::String(SecretValue::new("s3cret"))); +} + +#[tokio::test] +async fn percent_encodes_secret_name_path_segment() { + let server = MockServer::start().await; + Mock::given(path("/secrets/name%2Fwith%20spaces")) + .and(query_param("api-version", "7.4")) + .respond_with( + ResponseTemplate::new(200).set_body_json(serde_json::json!({"value": "value"})), + ) + .expect(1) + .mount(&server) + .await; + + let secret = manager(&server) + .get_secret_from_azure_key_vault("name/with spaces") + .await + .unwrap() + .unwrap(); + + assert_eq!(secret.as_str(), Some("value")); +} + +#[rstest::rstest] +#[case::not_found(404, None)] +#[case::forbidden(403, Some(403))] +#[tokio::test] +async fn handles_statuses(#[case] status: u16, #[case] expected_status: Option) { + let server = MockServer::start().await; + Mock::given(path("/secrets/NAME")) + .respond_with(ResponseTemplate::new(status)) + .expect(1) + .mount(&server) + .await; + + let result = manager(&server) + .get_secret_from_azure_key_vault("NAME") + .await; + + match expected_status { + None => assert_eq!(result.unwrap(), None), + Some(status) => assert!(matches!(result, Err(Error::Status(actual)) if actual == status)), + } +} + +#[tokio::test] +async fn missing_value_is_an_error() { + let server = MockServer::start().await; + Mock::given(path("/secrets/NAME")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({}))) + .expect(1) + .mount(&server) + .await; + + assert!(matches!( + manager(&server) + .get_secret_from_azure_key_vault("NAME") + .await, + Err(Error::MissingValue) + )); +} + +#[test] +fn new_validates_vault_environment() { + assert!(matches!( + AzureKeyVault::new(Arc::new(|_: &str| None)), + Err(Error::MissingEnvironment("AZURE_KEY_VAULT_URI")) + )); + assert!(matches!( + AzureKeyVault::new(Arc::new(|name: &str| { + (name == "AZURE_KEY_VAULT_URI").then(|| "http://vault.example".to_owned()) + })), + Err(Error::VaultUri) + )); + assert!(matches!( + AzureKeyVault::new(Arc::new(|name: &str| { + (name == "AZURE_KEY_VAULT_URI").then(|| "vault.example".to_owned()) + })), + Err(Error::VaultUri) + )); +} + +#[rstest::rstest] +#[case("https://myvault.vault.azure.net", "https://vault.azure.net/.default")] +#[case( + "https://v.vault.usgovcloudapi.net/", + "https://vault.usgovcloudapi.net/.default" +)] +#[case("http://localhost:8080", "https://localhost/.default")] +#[test] +fn derives_scope_from_vault_host(#[case] uri: &str, #[case] expected: &str) { + let manager = AzureKeyVault::with_client( + reqwest::Client::new(), + uri.parse().unwrap(), + Arc::new(|_: &str| None), + ) + .unwrap(); + + assert_eq!(manager.scope(), expected); +} + +#[tokio::test] +async fn missing_credentials_do_not_request_vault() { + let server = MockServer::start().await; + Mock::given(path("/secrets/NAME")) + .respond_with(ResponseTemplate::new(200)) + .expect(0) + .mount(&server) + .await; + + assert!( + manager_without_credentials(&server) + .get_secret_from_azure_key_vault("NAME") + .await + .is_err() + ); +} + +fn manager_without_credentials(server: &MockServer) -> AzureKeyVault { + AzureKeyVault::with_client( + reqwest::Client::new(), + server.uri().parse().unwrap(), + Arc::new(|_: &str| None), + ) + .unwrap() +} + +#[derive(Deserialize)] +struct Fixture { + cases: Vec, +} + +#[derive(Deserialize)] +struct FixtureCase { + secret_name: String, + response: FixtureResponse, + expected: FixtureExpected, +} + +#[derive(Deserialize)] +struct FixtureResponse { + status: u16, + body: serde_json::Value, +} + +#[derive(Deserialize)] +struct FixtureExpected { + value: Option, + missing: Option, + error: Option, +} + +#[tokio::test] +async fn parity_fixture_matches_python_backend_contract() { + let fixture: Fixture = + serde_json::from_str(include_str!("fixtures/key_vault_parity.json")).unwrap(); + for case in fixture.cases { + let server = MockServer::start().await; + Mock::given(path(format!("/secrets/{}", case.secret_name))) + .respond_with( + ResponseTemplate::new(case.response.status).set_body_json(case.response.body), + ) + .expect(1) + .mount(&server) + .await; + let result = manager(&server) + .get_secret_from_azure_key_vault(&case.secret_name) + .await; + if case.expected.missing == Some(true) { + assert_eq!(result.unwrap(), None); + } else if case.expected.error == Some(true) { + assert!(result.is_err()); + } else { + assert_eq!( + result.unwrap().unwrap().as_str(), + case.expected.value.as_deref() + ); + } + } +} diff --git a/litellm-rust/crates/secrets-azure/tests/live.rs b/litellm-rust/crates/secrets-azure/tests/live.rs new file mode 100644 index 00000000000..a062ba95070 --- /dev/null +++ b/litellm-rust/crates/secrets-azure/tests/live.rs @@ -0,0 +1,30 @@ +use std::sync::Arc; + +use litellm_core_utils::settings::ProcessEnvironment; +use litellm_secrets_azure::AzureKeyVault; +use litellm_secrets_types::Secret; + +#[tokio::test] +#[ignore] +async fn reads_a_real_secret() { + let environment = Arc::new(ProcessEnvironment); + let manager = AzureKeyVault::new(environment).unwrap(); + let name = std::env::var("AZURE_KEY_VAULT_LIVE_SECRET_NAME").unwrap(); + let secret = manager + .get_secret_from_azure_key_vault(&name) + .await + .unwrap() + .unwrap(); + assert!(matches!(&secret, Secret::String(_))); + let host = std::env::var("AZURE_KEY_VAULT_URI") + .unwrap() + .parse::() + .unwrap() + .host_str() + .unwrap() + .to_owned(); + let value_len = secret.as_str().unwrap().len(); + println!( + "native provider=litellm-secrets-azure vault_host={host} secret={name} value_len={value_len}" + ); +} diff --git a/litellm-rust/crates/secrets/Cargo.toml b/litellm-rust/crates/secrets/Cargo.toml index a7e7ec80636..f56ad100337 100644 --- a/litellm-rust/crates/secrets/Cargo.toml +++ b/litellm-rust/crates/secrets/Cargo.toml @@ -9,11 +9,13 @@ repository.workspace = true default = [] aws = ["dep:litellm-secrets-aws"] google = ["dep:litellm-secrets-google"] +azure = ["dep:litellm-secrets-azure"] [dependencies] litellm-secrets-types.workspace = true litellm-secrets-aws = { workspace = true, optional = true } litellm-secrets-google = { workspace = true, optional = true } +litellm-secrets-azure = { workspace = true, optional = true } litellm-core-utils.workspace = true base64.workspace = true serde.workspace = true diff --git a/litellm-rust/crates/secrets/src/error.rs b/litellm-rust/crates/secrets/src/error.rs index 0c6e681b8aa..3729b5c2f2c 100644 --- a/litellm-rust/crates/secrets/src/error.rs +++ b/litellm-rust/crates/secrets/src/error.rs @@ -30,4 +30,7 @@ pub enum Error { #[cfg(feature = "google")] #[error(transparent)] Google(#[from] litellm_secrets_google::Error), + #[cfg(feature = "azure")] + #[error(transparent)] + Azure(#[from] litellm_secrets_azure::Error), } diff --git a/litellm-rust/crates/secrets/src/handler.rs b/litellm-rust/crates/secrets/src/handler.rs index 943ffdf6158..84360b39f26 100644 --- a/litellm-rust/crates/secrets/src/handler.rs +++ b/litellm-rust/crates/secrets/src/handler.rs @@ -13,6 +13,8 @@ pub enum SecretManager { GoogleKms(crate::google::GoogleKms), #[cfg(feature = "google")] GoogleSecretManager(crate::google::GoogleSecretManager), + #[cfg(feature = "azure")] + AzureKeyVault(crate::azure::AzureKeyVault), } impl SecretManager { @@ -27,6 +29,8 @@ impl SecretManager { Self::GoogleKms(_) => KeyManagementSystem::GoogleKms, #[cfg(feature = "google")] Self::GoogleSecretManager(_) => KeyManagementSystem::GoogleSecretManager, + #[cfg(feature = "azure")] + Self::AzureKeyVault(_) => KeyManagementSystem::AzureKeyVault, } } } @@ -78,6 +82,11 @@ pub async fn get_secret_from_manager( .get_secret_from_google_secret_manager(secret_name) .await .map_err(Error::from), + #[cfg(feature = "azure")] + SecretManager::AzureKeyVault(client) => client + .get_secret_from_azure_key_vault(secret_name) + .await + .map_err(Error::from), } } diff --git a/litellm-rust/crates/secrets/src/lib.rs b/litellm-rust/crates/secrets/src/lib.rs index ff2e95f7b2f..52251bb593a 100644 --- a/litellm-rust/crates/secrets/src/lib.rs +++ b/litellm-rust/crates/secrets/src/lib.rs @@ -17,5 +17,7 @@ pub use state::{SecretManagerState, secret_manager_would_be_consulted}; #[cfg(feature = "aws")] pub use litellm_secrets_aws as aws; +#[cfg(feature = "azure")] +pub use litellm_secrets_azure as azure; #[cfg(feature = "google")] pub use litellm_secrets_google as google; diff --git a/litellm-rust/crates/secrets/tests/handler.rs b/litellm-rust/crates/secrets/tests/handler.rs index a2cbbd843e1..b5c7b9e0cfb 100644 --- a/litellm-rust/crates/secrets/tests/handler.rs +++ b/litellm-rust/crates/secrets/tests/handler.rs @@ -105,3 +105,64 @@ async fn google_handler_requires_canonical_base64_and_preserves_plaintext_whites Err(Error::MissingCiphertext) )); } + +#[cfg(feature = "azure")] +#[tokio::test] +async fn azure_handler_reads_missing_and_failed_secrets() { + use litellm_secrets::{ + Error, KeyManagementSettings, KeyManagementSystem, SecretManager, azure::AzureKeyVault, + get_secret_from_manager, + }; + use wiremock::{ + Mock, MockServer, ResponseTemplate, + matchers::{path, query_param}, + }; + + let server = MockServer::start().await; + Mock::given(path("/secrets/KEY")) + .and(query_param("api-version", "7.4")) + .respond_with( + ResponseTemplate::new(200).set_body_json(serde_json::json!({"value": "value"})), + ) + .expect(1) + .mount(&server) + .await; + let manager = SecretManager::AzureKeyVault( + AzureKeyVault::with_client( + reqwest::Client::new(), + server.uri().parse().unwrap(), + std::sync::Arc::new(|name: &str| (name == "AZURE_AD_TOKEN").then(|| "fake".to_owned())), + ) + .unwrap(), + ); + assert_eq!(manager.system(), KeyManagementSystem::AzureKeyVault); + let settings = KeyManagementSettings::default(); + let value = get_secret_from_manager(&manager, "KEY", &settings, &|_: &str| None) + .await + .unwrap() + .unwrap(); + assert_eq!(value.as_str(), Some("value")); + + let not_found = Mock::given(path("/secrets/MISSING")) + .respond_with(ResponseTemplate::new(404)) + .expect(1) + .mount_as_scoped(&server) + .await; + assert_eq!( + get_secret_from_manager(&manager, "MISSING", &settings, &|_: &str| None) + .await + .unwrap(), + None + ); + drop(not_found); + + Mock::given(path("/secrets/FAILED")) + .respond_with(ResponseTemplate::new(500)) + .expect(1) + .mount(&server) + .await; + assert!(matches!( + get_secret_from_manager(&manager, "FAILED", &settings, &|_: &str| None).await, + Err(Error::Azure(_)) + )); +} diff --git a/tests/test_litellm/secret_managers/test_secret_manager_handler.py b/tests/test_litellm/secret_managers/test_secret_manager_handler.py new file mode 100644 index 00000000000..0925198992b --- /dev/null +++ b/tests/test_litellm/secret_managers/test_secret_manager_handler.py @@ -0,0 +1,105 @@ +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +import pytest +from pydantic import BaseModel, ConfigDict + +from litellm.secret_managers.secret_manager_handler import get_secret_from_manager +from litellm.types.secret_managers.main import KeyManagementSystem + + +def _azure_exception_types() -> tuple[type[Exception], type[Exception]]: + try: + from azure.core.exceptions import ( + HttpResponseError, + ResourceNotFoundError, + ) + except ImportError: + return Exception, Exception + return HttpResponseError, ResourceNotFoundError + + +_AZURE_EXCEPTION_TYPES: Final[tuple[type[Exception], type[Exception]]] = _azure_exception_types() +AzureHttpResponseError: Final[type[Exception]] = _AZURE_EXCEPTION_TYPES[0] +AzureResourceNotFoundError: Final[type[Exception]] = _AZURE_EXCEPTION_TYPES[1] + + +class FixtureResponse(BaseModel): + model_config = ConfigDict(frozen=True) + + status: int + body: dict[str, object] + + +class FixtureExpected(BaseModel): + model_config = ConfigDict(frozen=True) + + value: str | None = None + missing: bool = False + error: bool = False + + +class FixtureCase(BaseModel): + model_config = ConfigDict(frozen=True) + + name: str + secret_name: str + response: FixtureResponse + expected: FixtureExpected + + +class Fixture(BaseModel): + model_config = ConfigDict(frozen=True) + + cases: tuple[FixtureCase, ...] + + +@dataclass(frozen=True, slots=True) +class FakeSecret: + value: str | None + + +@dataclass(frozen=True, slots=True) +class FakeAzureKeyVaultClient: + status: int + value: str | None + + def get_secret(self, name: str) -> FakeSecret: + if self.status == 404: + raise AzureResourceNotFoundError() + if self.status != 200: + raise AzureHttpResponseError() + return FakeSecret(value=self.value) + + +FIXTURE_PATH: Path = ( + Path(__file__).parents[3] + / "litellm-rust/crates/secrets-azure/tests/fixtures/key_vault_parity.json" +) + + +def test_azure_key_vault_matches_rust_parity_fixture() -> None: + fixture: Fixture = Fixture.model_validate_json(FIXTURE_PATH.read_text()) + for case in fixture.cases: + value: object = case.response.body.get("value") + secret: str | None = value if isinstance(value, str) else None + client: FakeAzureKeyVaultClient = FakeAzureKeyVaultClient( + status=case.response.status, + value=secret, + ) + if case.expected.missing or case.expected.error: + with pytest.raises(Exception): + get_secret_from_manager( + secret_name=case.secret_name, + key_manager=KeyManagementSystem.AZURE_KEY_VAULT.value, + client=client, + ) + continue + + result: str | None = get_secret_from_manager( + secret_name=case.secret_name, + key_manager=KeyManagementSystem.AZURE_KEY_VAULT.value, + client=client, + ) + assert result == case.expected.value From bfb4a8a2b33ed60fd083d59806f33aea4e5b5f59 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:36:10 +0000 Subject: [PATCH 06/15] refactor(rust): build Azure Key Vault auth inputs directly and keep credential tests offline Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/crates/auth-azure/src/lib.rs | 2 +- litellm-rust/crates/secrets-azure/Cargo.toml | 2 +- .../crates/secrets-azure/src/key_vault.rs | 42 +++++++++---------- .../crates/secrets-azure/tests/key_vault.rs | 4 +- 4 files changed, 24 insertions(+), 26 deletions(-) diff --git a/litellm-rust/crates/auth-azure/src/lib.rs b/litellm-rust/crates/auth-azure/src/lib.rs index e76227d6aa2..5c7c654b69d 100644 --- a/litellm-rust/crates/auth-azure/src/lib.rs +++ b/litellm-rust/crates/auth-azure/src/lib.rs @@ -4,4 +4,4 @@ mod resolve; mod types; pub use resolve::AzureAuthService; -pub use types::AzureAuthInputs; +pub use types::{AzureAuthInputs, ConfigValue}; diff --git a/litellm-rust/crates/secrets-azure/Cargo.toml b/litellm-rust/crates/secrets-azure/Cargo.toml index d559592cb34..96db7f235ef 100644 --- a/litellm-rust/crates/secrets-azure/Cargo.toml +++ b/litellm-rust/crates/secrets-azure/Cargo.toml @@ -12,7 +12,6 @@ litellm-secrets-types.workspace = true litellm-core-utils.workspace = true reqwest.workspace = true serde.workspace = true -serde_json.workspace = true thiserror.workspace = true veil.workspace = true percent-encoding = "2.3" @@ -21,4 +20,5 @@ percent-encoding = "2.3" tokio.workspace = true wiremock = "0.6.5" rstest.workspace = true +serde_json.workspace = true sha2.workspace = true diff --git a/litellm-rust/crates/secrets-azure/src/key_vault.rs b/litellm-rust/crates/secrets-azure/src/key_vault.rs index 262b19386df..e12289b83f5 100644 --- a/litellm-rust/crates/secrets-azure/src/key_vault.rs +++ b/litellm-rust/crates/secrets-azure/src/key_vault.rs @@ -1,21 +1,28 @@ use std::sync::Arc; -use litellm_auth_azure::{AzureAuthInputs, AzureAuthService}; +use litellm_auth_azure::{AzureAuthInputs, AzureAuthService, ConfigValue}; +use litellm_auth_types::{InputSource, Sourced}; use litellm_core_utils::settings::Lookup; use litellm_secrets_types::{Secret, SecretValue}; +use percent_encoding::{AsciiSet, NON_ALPHANUMERIC}; use serde::Deserialize; use crate::Error; const AZURE_KEY_VAULT_URI: &str = "AZURE_KEY_VAULT_URI"; const API_VERSION: &str = "7.4"; +const PATH_SEGMENT: &AsciiSet = &NON_ALPHANUMERIC + .remove(b'-') + .remove(b'.') + .remove(b'_') + .remove(b'~'); #[derive(Clone)] pub struct AzureKeyVault { client: reqwest::Client, vault: reqwest::Url, auth: Arc, - inputs: AzureAuthInputs, + inputs: Arc, environment: Arc, } @@ -33,21 +40,19 @@ impl AzureKeyVault { if vault.host_str().is_none() { return Err(Error::VaultUri); } - let scope = scope_for(&vault); - let inputs = AzureAuthInputs::from_sourced_optional_params( - serde_json::json!({ - "azure_scope": scope, - "enable_azure_ad_token_refresh": true, - }) - .as_object() - .expect("static Azure auth inputs object"), - &std::collections::BTreeMap::new(), - )?; + let inputs = AzureAuthInputs { + azure_scope: ConfigValue::Value(Sourced::new( + scope_for(&vault), + InputSource::Deployment, + )), + enable_azure_ad_token_refresh: Sourced::new(true, InputSource::Deployment), + ..AzureAuthInputs::default() + }; Ok(Self { client, vault, auth: Arc::new(AzureAuthService::default()), - inputs, + inputs: Arc::new(inputs), environment, }) } @@ -80,7 +85,7 @@ impl AzureKeyVault { .get_azure_ad_token(&self.inputs, &|key| self.environment.get(key)) .await? .ok_or(Error::MissingCredentials)?; - let encoded_name = encode_name(name); + let encoded_name = percent_encoding::utf8_percent_encode(name, PATH_SEGMENT); let url = self .vault .join(&format!("secrets/{encoded_name}?api-version={API_VERSION}")) @@ -111,12 +116,3 @@ fn scope_for(vault: &reqwest::Url) -> String { .map_or(host, |(_, remainder)| remainder); format!("https://{resource}/.default") } - -fn encode_name(name: &str) -> String { - percent_encoding::utf8_percent_encode(name, percent_encoding::NON_ALPHANUMERIC) - .to_string() - .replace("%2D", "-") - .replace("%2E", ".") - .replace("%5F", "_") - .replace("%7E", "~") -} diff --git a/litellm-rust/crates/secrets-azure/tests/key_vault.rs b/litellm-rust/crates/secrets-azure/tests/key_vault.rs index aae09369243..cf9102d0b45 100644 --- a/litellm-rust/crates/secrets-azure/tests/key_vault.rs +++ b/litellm-rust/crates/secrets-azure/tests/key_vault.rs @@ -160,7 +160,9 @@ fn manager_without_credentials(server: &MockServer) -> AzureKeyVault { AzureKeyVault::with_client( reqwest::Client::new(), server.uri().parse().unwrap(), - Arc::new(|_: &str| None), + Arc::new(|name: &str| { + (name == "AZURE_CREDENTIAL").then(|| "ClientSecretCredential".to_owned()) + }), ) .unwrap() } From df8966591939e6aa4f747429b762ec33a6721543 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:37:57 +0000 Subject: [PATCH 07/15] test: narrow Azure parity exception assertion Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../secret_managers/test_secret_manager_handler.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/test_litellm/secret_managers/test_secret_manager_handler.py b/tests/test_litellm/secret_managers/test_secret_manager_handler.py index 0925198992b..b4838912d27 100644 --- a/tests/test_litellm/secret_managers/test_secret_manager_handler.py +++ b/tests/test_litellm/secret_managers/test_secret_manager_handler.py @@ -89,7 +89,9 @@ def test_azure_key_vault_matches_rust_parity_fixture() -> None: value=secret, ) if case.expected.missing or case.expected.error: - with pytest.raises(Exception): + with pytest.raises( + AzureResourceNotFoundError if case.expected.missing else AzureHttpResponseError + ): get_secret_from_manager( secret_name=case.secret_name, key_manager=KeyManagementSystem.AZURE_KEY_VAULT.value, From b41edb57c1ff349c0fd537a20a84571b25053b28 Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 21 Sep 2026 20:42:03 +0000 Subject: [PATCH 08/15] fix(bedrock): keep batch S3 credentials out of chat requests and debug logs Register s3_access_key_id, s3_secret_access_key and s3_encryption_key_id as LiteLLM-owned batch params so they are no longer forwarded to Bedrock as additionalModelRequestFields (which 400s ordinary chat on a batch-configured deployment), keep them on CredentialLiteLLMParams so the batch/file paths still receive them, and redact the S3 credential key names in debug logs. Resolves LIT-8290 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../litellm_core_utils/secret_redaction.py | 2 +- litellm/types/router.py | 2 + litellm/types/utils.py | 3 ++ .../coverage_registry/llm_conversational.yaml | 1 + tests/e2e/coverage_registry/schema.py | 1 + .../test_bedrock_provider_matrix_e2e.py | 44 +++++++++++++++++++ tests/e2e/models.py | 1 + tests/test_litellm/test_secret_redaction.py | 19 ++++++++ tests/test_litellm/test_utils.py | 27 ++++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 8 ++++ 10 files changed, 107 insertions(+), 1 deletion(-) diff --git a/litellm/litellm_core_utils/secret_redaction.py b/litellm/litellm_core_utils/secret_redaction.py index 390abf41955..70b2cd08b4c 100644 --- a/litellm/litellm_core_utils/secret_redaction.py +++ b/litellm/litellm_core_utils/secret_redaction.py @@ -65,7 +65,7 @@ def _build_secret_patterns() -> "re.Pattern[str]": # private_key with PEM-aware value capture r"""private_key['\"]?\s*[:=]\s*['\"]?(?:-----BEGIN[A-Z \-]*PRIVATE KEY-----[\s\S]*?-----END[A-Z \-]*PRIVATE KEY-----|[^\s,'\"})\]{}>]+)""", r"(?:master_key|xai_key|database_url|db_url|connection_string|" - r"aws_secret_access_key|aws_session_token|aws_access_key_id|" + r"aws_secret_access_key|aws_session_token|aws_access_key_id|s3_secret_access_key|s3_access_key_id|" r"signing_key|encryption_key|" r"auth_token|access_token|refresh_token|" r"slack_webhook_url|webhook_url|" diff --git a/litellm/types/router.py b/litellm/types/router.py index fd426835d65..8ad01206811 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -305,6 +305,8 @@ class CredentialLiteLLMParams(BaseModel): s3_bucket_name: str | None = None s3_endpoint_url: str | None = None s3_region_name: str | None = None + s3_access_key_id: str | None = None + s3_secret_access_key: str | None = None s3_encryption_key_id: str | None = None s3_bucket_owner: str | None = None aws_batch_role_arn: str | None = None diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 82cd0250857..3695be2641a 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3840,6 +3840,9 @@ bedrock_batch_litellm_params: Final = ( "s3_endpoint_url", "s3_output_bucket_name", "s3_bucket_owner", + "s3_access_key_id", + "s3_secret_access_key", + "s3_encryption_key_id", "bedrock_tags", ) diff --git a/tests/e2e/coverage_registry/llm_conversational.yaml b/tests/e2e/coverage_registry/llm_conversational.yaml index 36fbd39154d..49d4d92ff0b 100644 --- a/tests/e2e/coverage_registry/llm_conversational.yaml +++ b/tests/e2e/coverage_registry/llm_conversational.yaml @@ -31,6 +31,7 @@ - {id: llm.chat_completions.bedrock_converse.thinking.nonstream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: bedrock_converse, capability: thinking, streaming: nonstream, assertions: [works], source: "model_prices json", rationale: "Anthropic thinking on Bedrock"} - {id: llm.chat_completions.bedrock_converse.response_headers.nonstream.works, module: llm, tier: P0, subject_endpoint: chat_completions, route: bedrock_converse, capability: response_headers, streaming: nonstream, assertions: [works], source: "llms/bedrock/chat/converse_handler.py:248", rationale: "Bedrock request ids must surface as llm_provider-* response headers on /chat/completions so callers can correlate calls with AWS-side logs (#37003)", fail_before_fix: proven} - {id: llm.chat_completions.bedrock_converse.response_headers.stream.works, module: llm, tier: P0, subject_endpoint: chat_completions, route: bedrock_converse, capability: response_headers, streaming: stream, assertions: [works], source: "llms/bedrock/chat/converse_handler.py:154", rationale: "The llm_provider-* headers must also surface on streaming /chat/completions, where CustomStreamWrapper carries them instead of the nonstream setter"} +- {id: llm.chat_completions.bedrock_converse.batch_deployment.nonstream.works, module: llm, tier: P0, subject_endpoint: chat_completions, route: bedrock_converse, capability: batch_deployment, streaming: nonstream, assertions: [works], source: "types/utils.py bedrock_batch_litellm_params", rationale: "A deployment carrying the documented batch-only S3 keys (s3_access_key_id, s3_secret_access_key, s3_encryption_key_id) must still serve ordinary chat; unregistered keys fall into optional_params and are forwarded as additionalModelRequestFields, which Bedrock 400s and which puts the S3 secret in the request body and debug log (LIT-8290)", fail_before_fix: proven} - {id: llm.chat_completions.bedrock_invoke.basic.nonstream.works, module: llm, tier: P0, subject_endpoint: chat_completions, route: bedrock_invoke, capability: basic, streaming: nonstream, assertions: [works], source: "proxy_server.py:8455", rationale: "Regional inference-profile ids (us.anthropic.*) over the invoke route, the deployment shape behind a customer timeout report on v1.90.0"} - {id: llm.chat_completions.bedrock_invoke.basic.stream.works, module: llm, tier: P0, subject_endpoint: chat_completions, route: bedrock_invoke, capability: basic, streaming: stream, assertions: [works], source: "proxy_server.py:8455", rationale: "Streaming with regional inference-profile ids over the invoke route"} - {id: llm.chat_completions.vertex.basic.nonstream.works, module: llm, tier: P0, subject_endpoint: chat_completions, route: vertex, capability: basic, streaming: nonstream, assertions: [works], source: "proxy_server.py:8455", rationale: "P0 route; Vertex AI"} diff --git a/tests/e2e/coverage_registry/schema.py b/tests/e2e/coverage_registry/schema.py index fa6dad90126..8b0d38a083c 100644 --- a/tests/e2e/coverage_registry/schema.py +++ b/tests/e2e/coverage_registry/schema.py @@ -63,6 +63,7 @@ LlmRoute = Literal[ LlmCapability = Literal[ "assume_role", "basic", + "batch_deployment", "count_tokens", "govcloud_partition", "input_validation", diff --git a/tests/e2e/llm_translation/test_bedrock_provider_matrix_e2e.py b/tests/e2e/llm_translation/test_bedrock_provider_matrix_e2e.py index 3c6aaa75ab3..5f0a931109c 100644 --- a/tests/e2e/llm_translation/test_bedrock_provider_matrix_e2e.py +++ b/tests/e2e/llm_translation/test_bedrock_provider_matrix_e2e.py @@ -131,6 +131,50 @@ class TestBedrockResponseHeaders: _assert_request_id_header(result) +def _register_bedrock_batch_deployment(client: PassthroughClient, resources: ResourceManager) -> str: + model = f"e2e-bedrock-batch-chat-{unique_marker()}" + model_id = client.proxy.create_model( + model, + LiteLLMParamsBody( + model=CONVERSE_REGIONAL_BACKEND, + aws_access_key_id="os.environ/AWS_ACCESS_KEY_ID", + aws_secret_access_key="os.environ/AWS_SECRET_ACCESS_KEY", + aws_region_name="os.environ/AWS_REGION", + s3_bucket_name="os.environ/AWS_BATCH_S3_BUCKET", + s3_access_key_id="os.environ/AWS_ACCESS_KEY_ID", + s3_secret_access_key="os.environ/AWS_SECRET_ACCESS_KEY", + s3_encryption_key_id=f"alias/e2e-unused-{unique_marker()}", + aws_batch_role_arn="os.environ/AWS_BATCH_ROLE_ARN", + ), + ) + resources.defer(lambda: client.proxy.delete_model(model_id)) + return model + + +class TestBedrockBatchDeploymentServesChat: + @pytest.mark.covers( + "llm.chat_completions.bedrock_converse.batch_deployment.nonstream.works", + exercised_on=[], + ) + def test_batch_s3_keys_do_not_break_chat( + self, client: PassthroughClient, resources: ResourceManager + ) -> None: + model = _register_bedrock_batch_deployment(client, resources) + key = resources.key() + + result = client.proxy.transport.send( + "/chat/completions", + headers=client.proxy.transport.bearer(key), + json=ChatBody(model=model, messages=_prompt(), max_tokens=64), + ) + + assert result.ok, ( + f"chat on a batch-configured deployment failed: {result.status_code} {result.body[:300]}; " + "batch-only S3 keys were forwarded to Bedrock as additionalModelRequestFields" + ) + _assert_completion(ChatResponse.model_validate_json(result.body)) + + class TestBedrockInvokeRegionalModelIds: @pytest.mark.covers("llm.chat_completions.bedrock_invoke.basic.nonstream.works", exercised_on=[]) def test_invoke_regional_id_completes( diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 47ef672ebec..441110574eb 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -991,6 +991,7 @@ class LiteLLMParamsBody(BaseModel): s3_region_name: str | None = None s3_access_key_id: str | None = None s3_secret_access_key: str | None = None + s3_encryption_key_id: str | None = None aws_batch_role_arn: str | None = None aws_role_name: str | None = None aws_session_name: str | None = None diff --git a/tests/test_litellm/test_secret_redaction.py b/tests/test_litellm/test_secret_redaction.py index 85933fbf9e8..f58ade11d1c 100644 --- a/tests/test_litellm/test_secret_redaction.py +++ b/tests/test_litellm/test_secret_redaction.py @@ -629,6 +629,25 @@ def test_aws_credential_redaction_catches_quoted_values(): assert redact_string(safe) == safe +def test_bedrock_batch_s3_credential_redaction_in_deployment_dump(): + """The router logs each deployment's litellm_params at DEBUG. A Bedrock batch + deployment carries s3_secret_access_key there, which the aws_* key-name rule + did not cover, so the S3 secret was printed verbatim (LIT-8290).""" + cases = ( + "{'s3_secret_access_key': 'wJalrXUtnFEMIK7MDENGbPxRfiCYEXAMPLEKEY'}", + "s3_secret_access_key=wJalrXUtnFEMIK7MDENGbPxRfiCYEXAMPLEKEY", + "{'s3_access_key_id': 'not-an-akia-shaped-value'}", + ) + for secret_line in cases: + result = redact_string(secret_line) + assert "REDACTED" in result, f"S3 credential redaction missed: {secret_line!r}" + assert "wJalrXUtnFEMIK7MDENGbPxRfiCYEXAMPLEKEY" not in result + assert "not-an-akia-shaped-value" not in result + + safe = "'s3_bucket_name': 'my-batch-bucket'" + assert redact_string(safe) == safe + + @pytest.mark.parametrize( "extra", ( diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 8e419b15044..6cc8b7f7565 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -4653,6 +4653,33 @@ def test_bedrock_batch_params_never_reach_the_provider(): ) +def test_documented_batch_s3_credentials_never_reach_the_provider(): + """The Bedrock batch docs tell users to put s3_access_key_id, s3_secret_access_key + and s3_encryption_key_id on the deployment. Left unregistered they are swept into + additionalModelRequestFields, Bedrock 400s ordinary chat on that deployment with + `s3_secret_access_key: Extra inputs are not permitted`, and the S3 secret is sent + to the provider and printed in the debug log (LIT-8290). + """ + configured = { + "s3_access_key_id": "configured-access-key-id", + "s3_secret_access_key": "configured-secret-access-key", + "s3_encryption_key_id": "arn:aws:kms:us-east-1:000000000000:key/configured", + } + kwargs = {"a_real_provider_specific_param": 1, **configured} + + non_default = get_non_default_completion_params(dict(kwargs)) + + assert non_default == {"a_real_provider_specific_param": 1}, ( + "documented batch S3 credentials leaked into the provider params: " + f"{sorted(set(non_default) - {'a_real_provider_specific_param'})}" + ) + + batch_params = dict(GenericLiteLLMParams(**kwargs)) + assert {field: batch_params.get(field) for field in configured} == configured, ( + "registering these must not strip them from the batch path" + ) + + def test_client_side_timeout_marker_never_reaches_the_provider(): """The proxy stamps kwargs["client_side_timeout"] = True whenever a request carries a caller-supplied timeout (body timeout / request_timeout / stream_timeout or the diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 440958fa412..eb09105817a 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -31211,6 +31211,8 @@ export interface components { regional_processing_uplift_multiplier_us?: number | null; /** Rpm */ rpm?: number | null; + /** S3 Access Key Id */ + s3_access_key_id?: string | null; /** S3 Bucket Name */ s3_bucket_name?: string | null; /** S3 Bucket Owner */ @@ -31223,6 +31225,8 @@ export interface components { s3_output_bucket_name?: string | null; /** S3 Region Name */ s3_region_name?: string | null; + /** S3 Secret Access Key */ + s3_secret_access_key?: string | null; /** Search Context Cost Per Query */ search_context_cost_per_query?: { [key: string]: unknown; @@ -42011,6 +42015,8 @@ export interface components { regional_processing_uplift_multiplier_us?: number | null; /** Rpm */ rpm?: number | null; + /** S3 Access Key Id */ + s3_access_key_id?: string | null; /** S3 Bucket Name */ s3_bucket_name?: string | null; /** S3 Bucket Owner */ @@ -42023,6 +42029,8 @@ export interface components { s3_output_bucket_name?: string | null; /** S3 Region Name */ s3_region_name?: string | null; + /** S3 Secret Access Key */ + s3_secret_access_key?: string | null; /** Search Context Cost Per Query */ search_context_cost_per_query?: { [key: string]: unknown; From e5a9b5c113b9d8107d4d60bef61cba1354aa77b2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 21 Sep 2026 13:51:03 -0700 Subject: [PATCH 09/15] fix(alerting): deliver every distinct alert queued in one flush window --- .../SlackAlerting/batching_handler.py | 40 ++++---- .../SlackAlerting/slack_alerting.py | 12 +-- litellm/types/integrations/slack_alerting.py | 15 ++- .../SlackAlerting/test_batching_handler.py | 98 +++++++++++++++++++ 4 files changed, 136 insertions(+), 29 deletions(-) create mode 100644 tests/test_litellm/integrations/SlackAlerting/test_batching_handler.py diff --git a/litellm/integrations/SlackAlerting/batching_handler.py b/litellm/integrations/SlackAlerting/batching_handler.py index 1c35a15d5a1..d152985a2c5 100644 --- a/litellm/integrations/SlackAlerting/batching_handler.py +++ b/litellm/integrations/SlackAlerting/batching_handler.py @@ -1,14 +1,18 @@ """ Handles Batching + sending Httpx Post requests to slack -Slack alerts are sent every 10s or when events are greater than X events +Slack alerts are sent every DEFAULT_FLUSH_INTERVAL_SECONDS or when events are greater than X events see custom_batch_logger.py for more details / defaults """ +from collections import Counter +from collections.abc import Sequence +from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Final from litellm._logging import verbose_proxy_logger +from litellm.types.integrations.slack_alerting import AlertQueueItem, AlertType from .ms_teams import MS_TEAMS_ALERTING_DESTINATION, build_ms_teams_payload @@ -20,26 +24,20 @@ else: SlackAlertingType = Any -def squash_payloads(queue): - squashed: Final = {} - if len(queue) == 0: - return squashed - if len(queue) == 1: - return {"key": {"item": queue[0], "count": 1}} +@dataclass(frozen=True, slots=True) +class SquashedAlert: + item: AlertQueueItem + count: int - for item in queue: - url = item["url"] - alert_type = item["alert_type"] - _key = (url, alert_type) - if _key in squashed: - squashed[_key]["count"] += 1 - # Merge the payloads +def _squash_key(item: AlertQueueItem) -> tuple[str, AlertType | str, str]: + return (item["url"], item["alert_type"], item["payload"]["text"]) - else: - squashed[_key] = {"item": item, "count": 1} - return squashed +def squash_payloads(queue: Sequence[AlertQueueItem]) -> tuple[SquashedAlert, ...]: + counts: Final = Counter(_squash_key(item) for item in queue) + first_item_by_key: Final = {_squash_key(item): item for item in reversed(queue)} + return tuple(SquashedAlert(item=first_item_by_key[key], count=count) for key, count in counts.items()) def _print_alerting_payload_warning(payload: dict, slackAlertingInstance: SlackAlertingType): @@ -53,17 +51,15 @@ def _print_alerting_payload_warning(payload: dict, slackAlertingInstance: SlackA verbose_proxy_logger.warning(payload) -async def send_to_webhook(slackAlertingInstance: SlackAlertingType, item, count): +async def send_to_webhook(slackAlertingInstance: SlackAlertingType, item: AlertQueueItem, count: int) -> None: """ Send a single slack alert to the webhook """ import json - payload: Final = item.get("payload", {}) + text: Final = item["payload"]["text"] + payload: Final = {"text": text if count == 1 else f"[Num Alerts: {count}]\n\n{text}"} try: - if count > 1: - payload["text"] = f"[Num Alerts: {count}]\n\n{payload['text']}" - request_body: Final = ( build_ms_teams_payload(payload["text"]) if item.get("format") == MS_TEAMS_ALERTING_DESTINATION else payload ) diff --git a/litellm/integrations/SlackAlerting/slack_alerting.py b/litellm/integrations/SlackAlerting/slack_alerting.py index 66e2754d5ad..7b579546403 100644 --- a/litellm/integrations/SlackAlerting/slack_alerting.py +++ b/litellm/integrations/SlackAlerting/slack_alerting.py @@ -1583,12 +1583,12 @@ Model Info: if not self.log_queue: return - squashed_queue: Final = squash_payloads(self.log_queue) - tasks: Final = [ - send_to_webhook(slackAlertingInstance=self, item=item["item"], count=item["count"]) - for item in squashed_queue.values() - ] - await asyncio.gather(*tasks) + await asyncio.gather( + *( + send_to_webhook(slackAlertingInstance=self, item=squashed.item, count=squashed.count) + for squashed in squash_payloads(self.log_queue) + ) + ) self.log_queue.clear() async def _flush_digest_buckets(self): diff --git a/litellm/types/integrations/slack_alerting.py b/litellm/types/integrations/slack_alerting.py index 64c0c530e9b..33bb446364e 100644 --- a/litellm/types/integrations/slack_alerting.py +++ b/litellm/types/integrations/slack_alerting.py @@ -1,11 +1,12 @@ import os import time +from collections.abc import Mapping from datetime import datetime as dt from enum import Enum from typing import Any, Final, Literal, Optional, Union from pydantic import BaseModel, Field -from typing_extensions import TypedDict +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm.types.utils import LiteLLMPydanticObjectBase @@ -235,6 +236,18 @@ DEFAULT_ALERT_TYPES: Final[list[AlertType]] = [ ] +class AlertText(TypedDict): + text: ReadOnly[str] + + +class AlertQueueItem(TypedDict): + url: ReadOnly[str] + headers: ReadOnly[Mapping[str, str]] + payload: ReadOnly[AlertText] + alert_type: ReadOnly[AlertType | str] + format: NotRequired[ReadOnly[str]] + + class HangingRequestData(BaseModel): request_id: str model: str diff --git a/tests/test_litellm/integrations/SlackAlerting/test_batching_handler.py b/tests/test_litellm/integrations/SlackAlerting/test_batching_handler.py new file mode 100644 index 00000000000..9052cb8bb5d --- /dev/null +++ b/tests/test_litellm/integrations/SlackAlerting/test_batching_handler.py @@ -0,0 +1,98 @@ +import json +from typing import Final +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from litellm.integrations.SlackAlerting.ms_teams import MS_TEAMS_WEBHOOK_URL_ENV +from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting +from litellm.types.integrations.slack_alerting import AlertQueueItem, AlertType + +SLACK_WEBHOOK_URL: Final = "https://hooks.slack.com/services/test" +THRESHOLD_ALERT: Final = "User Budget: 15% or less of budget remaining\n\n*user_id:* `user-a`" +CROSSED_ALERT: Final = "User Budget: Budget Crossed\n\n*user_id:* `user-b`" + + +def _slack_alerting_recording_posts(alerting: list[str]) -> SlackAlerting: + slack_alerting: Final = SlackAlerting(alerting=alerting) + slack_alerting.periodic_started = True + response: Final = MagicMock() + response.status_code = 200 + slack_alerting.async_http_handler = MagicMock() + slack_alerting.async_http_handler.post = AsyncMock(return_value=response) + return slack_alerting + + +def _queued_slack_alert(text: str) -> AlertQueueItem: + return { + "url": SLACK_WEBHOOK_URL, + "headers": {"Content-type": "application/json"}, + "payload": {"text": text}, + "alert_type": AlertType.budget_alerts, + } + + +def _posted_bodies(slack_alerting: SlackAlerting) -> tuple[dict, ...]: + return tuple(json.loads(call.kwargs["data"]) for call in slack_alerting.async_http_handler.post.call_args_list) + + +async def _send_budget_alert(slack_alerting: SlackAlerting, message: str) -> None: + await slack_alerting.send_alert( + message=message, + level="High", + alert_type=AlertType.budget_alerts, + alerting_metadata={}, + ) + + +@pytest.mark.asyncio +async def test_async_send_batch_delivers_every_distinct_alert_queued_in_one_flush(monkeypatch): + monkeypatch.setenv("SLACK_WEBHOOK_URL", SLACK_WEBHOOK_URL) + slack_alerting: Final = _slack_alerting_recording_posts(["slack"]) + await _send_budget_alert(slack_alerting, THRESHOLD_ALERT) + await _send_budget_alert(slack_alerting, CROSSED_ALERT) + + await slack_alerting.async_send_batch() + + posted_texts: Final = tuple(body["text"] for body in _posted_bodies(slack_alerting)) + assert len(posted_texts) == 2 + assert THRESHOLD_ALERT in posted_texts[0] + assert CROSSED_ALERT in posted_texts[1] + assert not any(text.startswith("[Num Alerts") for text in posted_texts) + assert slack_alerting.log_queue == [] + + +@pytest.mark.asyncio +async def test_async_send_batch_collapses_only_identical_alerts(): + slack_alerting: Final = _slack_alerting_recording_posts(["slack"]) + slack_alerting.log_queue.extend( + ( + _queued_slack_alert(THRESHOLD_ALERT), + _queued_slack_alert(CROSSED_ALERT), + _queued_slack_alert(THRESHOLD_ALERT), + ) + ) + + await slack_alerting.async_send_batch() + + assert _posted_bodies(slack_alerting) == ( + {"text": f"[Num Alerts: 2]\n\n{THRESHOLD_ALERT}"}, + {"text": CROSSED_ALERT}, + ) + + +@pytest.mark.asyncio +async def test_async_send_batch_delivers_every_distinct_ms_teams_alert(monkeypatch): + monkeypatch.setenv(MS_TEAMS_WEBHOOK_URL_ENV, "https://teams.example/webhook") + slack_alerting: Final = _slack_alerting_recording_posts(["ms_teams"]) + await _send_budget_alert(slack_alerting, THRESHOLD_ALERT) + await _send_budget_alert(slack_alerting, CROSSED_ALERT) + + await slack_alerting.async_send_batch() + + card_texts: Final = tuple( + body["attachments"][0]["content"]["body"][0]["text"] for body in _posted_bodies(slack_alerting) + ) + assert len(card_texts) == 2 + assert THRESHOLD_ALERT in card_texts[0] + assert CROSSED_ALERT in card_texts[1] From d8ce49de064a10ed785544d9c809541661def4e0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 21 Sep 2026 13:51:26 -0700 Subject: [PATCH 10/15] fix(proxy): evict the cached user row when SCIM or /user/delete removes a user --- .../internal_user_endpoints.py | 1 + .../management_endpoints/scim/scim_v2.py | 4 ++ .../scim/test_scim_key_deactivation.py | 39 +++++++++++++++++++ .../test_internal_user_endpoints.py | 39 +++++++++++++++++++ 4 files changed, 83 insertions(+) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 1c986305c21..c6b89096ca9 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -2518,6 +2518,7 @@ async def delete_user( ## DELETE USERS deleted_users: Final = await _user_table(prisma_client).delete_many(where={"user_id": {"in": data.user_ids}}) + await evict_and_broadcast(cache_keys=tuple(data.user_ids), user_api_key_cache=user_api_key_cache) return deleted_users diff --git a/litellm/proxy/management_endpoints/scim/scim_v2.py b/litellm/proxy/management_endpoints/scim/scim_v2.py index b676c0ddb82..3292a0141d1 100644 --- a/litellm/proxy/management_endpoints/scim/scim_v2.py +++ b/litellm/proxy/management_endpoints/scim/scim_v2.py @@ -1871,6 +1871,10 @@ async def delete_user( # Delete user await _table(UserRepository(prisma_client)).delete(where={"user_id": user_id}) + from litellm.proxy.proxy_server import user_api_key_cache + + await evict_and_broadcast(cache_keys=(user_id,), user_api_key_cache=user_api_key_cache) + return Response(status_code=204) except Exception as e: raise handle_exception_on_proxy(e) diff --git a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_key_deactivation.py b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_key_deactivation.py index 2cfbfbbb3cb..d35a676a28b 100644 --- a/tests/test_litellm/proxy/management_endpoints/scim/test_scim_key_deactivation.py +++ b/tests/test_litellm/proxy/management_endpoints/scim/test_scim_key_deactivation.py @@ -601,3 +601,42 @@ async def test_scim_status_write_refreshes_user_cache( else: assert cached is None broadcast.assert_awaited_once_with(cache_key=user_id) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure", [None, "delete"]) +async def test_scim_delete_user_evicts_cached_user_row(failure: str | None) -> None: + from typing import Final + + from litellm.proxy._types import ProxyException + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + + user_id: Final = "scim-deleted-user" + saved: Final = LiteLLM_UserTable(user_id=user_id, user_email="x@example.com", teams=[], metadata={}) + client, db = _build_prisma_with_keys([], mock_user=saved.model_copy(deep=True)) + if failure == "delete": + db.litellm_usertable.delete.side_effect = RuntimeError("user delete failed") + cache: Final = UserApiKeyCache() + await cache.async_set_cache(key=user_id, value=saved, model_type=LiteLLM_UserTable) + with ( + patch("litellm.proxy.proxy_server.prisma_client", client), # test-quality-ok: substitute the database dependency + patch("litellm.proxy.proxy_server.user_api_key_cache", cache), # test-quality-ok: exercise a real isolated cache + patch("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()), # test-quality-ok: isolate the logging dependency + patch( # test-quality-ok: observe the Redis publication boundary + "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation", + new_callable=AsyncMock, + ) as broadcast, + ): + if failure == "delete": + with pytest.raises(ProxyException, match="user delete failed"): + await delete_user(user_id=user_id) + else: + response: Final = await delete_user(user_id=user_id) + assert response.status_code == 204 + cached: Final = await cache.async_get_cache(key=user_id, model_type=LiteLLM_UserTable) + if failure == "delete": + assert cached == saved + broadcast.assert_not_awaited() + else: + assert cached is None + broadcast.assert_awaited_once_with(cache_key=user_id) diff --git a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py index b8f1aa0330b..ede0ecb2790 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py @@ -4691,3 +4691,42 @@ async def test_user_update_hashes_and_persists_strong_password(_admin_prisma, mo written_data = mock_prisma_client.update_data.call_args.kwargs["data"] assert written_data.get("password") is not None assert written_data["password"] != strong_password + + +@pytest.mark.asyncio +async def test_delete_user_evicts_cached_user_rows(mocker: MockerFixture) -> None: + from litellm.proxy._types import DeleteUserRequest, LiteLLM_UserTable + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache + from litellm.proxy.management_endpoints.internal_user_endpoints import delete_user + + deleted: Final = LiteLLM_UserTable(user_id="user-gone", user_email="gone@example.test", teams=[]) + survivor: Final = LiteLLM_UserTable(user_id="user-stays", user_email="stays@example.test", teams=[]) + prisma_client: Final = mocker.MagicMock() + prisma_client.db.litellm_usertable.find_unique = mocker.AsyncMock(return_value=deleted) + prisma_client.db.litellm_teamtable.find_many = mocker.AsyncMock(return_value=[]) + prisma_client.db.litellm_jwtkeymapping.find_many = mocker.AsyncMock(return_value=[]) + prisma_client.db.litellm_verificationtoken.find_many = mocker.AsyncMock(return_value=[]) + prisma_client.db.litellm_verificationtoken.delete_many = mocker.AsyncMock(return_value=0) + prisma_client.db.litellm_invitationlink.delete_many = mocker.AsyncMock(return_value=0) + prisma_client.db.litellm_organizationmembership.delete_many = mocker.AsyncMock(return_value=0) + prisma_client.db.litellm_teammembership.delete_many = mocker.AsyncMock(return_value=0) + prisma_client.db.litellm_usertable.delete_many = mocker.AsyncMock(return_value=1) + mocker.patch("litellm.proxy.proxy_server.prisma_client", prisma_client) # test-quality-ok: substitute the database dependency + cache: Final = UserApiKeyCache() + for row in (deleted, survivor): + await cache.async_set_cache(key=row.user_id, value=row, model_type=LiteLLM_UserTable) + mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", cache) # test-quality-ok: exercise a real isolated cache + mocker.patch("litellm.proxy.proxy_server.proxy_logging_obj", None) # test-quality-ok: delete_user reads it off proxy_server at call time + broadcast: Final = mocker.patch( # test-quality-ok: observe the Redis publication boundary + "litellm.proxy.common_utils.auth_cache_invalidation_pubsub.publish_auth_cache_invalidation", + new_callable=mocker.AsyncMock, + ) + + await delete_user( + data=DeleteUserRequest(user_ids=[deleted.user_id]), + user_api_key_dict=UserAPIKeyAuth(user_id="proxy-admin", user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + assert await cache.async_get_cache(key=deleted.user_id, model_type=LiteLLM_UserTable) is None + assert await cache.async_get_cache(key=survivor.user_id, model_type=LiteLLM_UserTable) == survivor + broadcast.assert_awaited_once_with(cache_key=deleted.user_id) From 5acfcfa4febefce2f9df7d6876915278bf91d229 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 21:00:48 +0000 Subject: [PATCH 11/15] feat(rust): native Azure Blob response cache backend Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/Cargo.lock | 82 +++ litellm-rust/Cargo.toml | 1 + litellm-rust/crates/auth-azure/src/types.rs | 15 + .../crates/cache-azure-blob/Cargo.toml | 22 + .../crates/cache-azure-blob/src/cache.rs | 246 +++++++ .../cache-azure-blob/src/cache/tests.rs | 692 ++++++++++++++++++ .../crates/cache-azure-blob/src/credential.rs | 84 +++ .../crates/cache-azure-blob/src/lib.rs | 5 + .../crates/cache-azure-blob/src/tests.rs | 0 litellm-rust/crates/python-bridge/Cargo.toml | 1 + .../crates/python-bridge/src/cache/config.rs | 49 +- .../crates/python-bridge/src/cache/facade.rs | 88 ++- .../crates/python-bridge/src/cache/handle.rs | 17 +- .../crates/python-bridge/src/cache/native.rs | 44 +- tests/test_litellm_rust/test_cache.py | 117 +++ 15 files changed, 1439 insertions(+), 24 deletions(-) create mode 100644 litellm-rust/crates/cache-azure-blob/Cargo.toml create mode 100644 litellm-rust/crates/cache-azure-blob/src/cache.rs create mode 100644 litellm-rust/crates/cache-azure-blob/src/cache/tests.rs create mode 100644 litellm-rust/crates/cache-azure-blob/src/credential.rs create mode 100644 litellm-rust/crates/cache-azure-blob/src/lib.rs create mode 100644 litellm-rust/crates/cache-azure-blob/src/tests.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index ed4ae4e3353..5b112519add 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -115,6 +115,28 @@ dependencies = [ "pin-project-lite", ] +[[package]] +name = "async-stream" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b5a71a6f37880a80d1d7f19efd781e4b5de42c88f0722cc13bcb6cc2cfe8476" +dependencies = [ + "async-stream-impl", + "futures-core", + "pin-project-lite", +] + +[[package]] +name = "async-stream-impl" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c7c24de15d275a1ecfd47a380fb4d5ec9bfe0933f309ed5e705b775596a3574d" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + [[package]] name = "async-trait" version = "0.1.91" @@ -599,6 +621,37 @@ dependencies = [ "url", ] +[[package]] +name = "azure_storage_blob" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "17b10207ecf7d666df6940b50051f433b3cd5d2b9b1dd190613208d7a84e7eed" +dependencies = [ + "async-stream", + "async-trait", + "azure_core", + "azure_storage_common", + "bytes", + "futures", + "percent-encoding", + "pin-project", + "serde", + "serde_json", + "time", + "tokio", +] + +[[package]] +name = "azure_storage_common" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0af2e6aeb8d76b17fc998f453c320913f73787b944e3cc29509d19411fa0321d" +dependencies = [ + "azure_core", + "serde", + "time", +] + [[package]] name = "base64" version = "0.13.1" @@ -2464,6 +2517,23 @@ dependencies = [ "tokio", ] +[[package]] +name = "litellm-cache-azure-blob" +version = "0.1.0" +dependencies = [ + "async-trait", + "azure_core", + "azure_storage_blob", + "futures-util", + "litellm-auth-azure", + "litellm-auth-types", + "litellm-cache", + "litellm-cache-response", + "serde_json", + "tokio", + "url", +] + [[package]] name = "litellm-cache-memory" version = "0.1.0" @@ -2666,6 +2736,7 @@ dependencies = [ "litellm-auth", "litellm-auth-gcp", "litellm-cache", + "litellm-cache-azure-blob", "litellm-cache-memory", "litellm-cache-redis", "litellm-cache-response", @@ -3487,6 +3558,16 @@ version = "1.2.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a1d01941d82fa2ab50be1e79e6714289dd7cde78eba4c074bc5a4374f650dfe0" +[[package]] +name = "quick-xml" +version = "0.41.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e660451e55124f798a69a5af3f49ccfbefbd41910eefd25caf2393e1f3473ec1" +dependencies = [ + "memchr", + "serde", +] + [[package]] name = "quinn" version = "0.11.11" @@ -5030,6 +5111,7 @@ dependencies = [ "base64 0.22.1", "bytes", "futures", + "quick-xml", "serde", "serde_json", "url", diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 570d0dd3568..58fc1ba1d05 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -27,6 +27,7 @@ litellm-llms = { path = "crates/llms" } litellm-types = { path = "crates/types" } litellm-core-utils = { path = "crates/core-utils" } litellm-cache = { path = "crates/cache" } +litellm-cache-azure-blob = { path = "crates/cache-azure-blob" } litellm-cache-memory = { path = "crates/cache-memory" } litellm-cache-redis = { path = "crates/cache-redis" } litellm-cache-response = { path = "crates/cache-response" } diff --git a/litellm-rust/crates/auth-azure/src/types.rs b/litellm-rust/crates/auth-azure/src/types.rs index d5a00f09751..a3a898f000f 100644 --- a/litellm-rust/crates/auth-azure/src/types.rs +++ b/litellm-rust/crates/auth-azure/src/types.rs @@ -51,6 +51,21 @@ pub struct AzureAuthInputs { } impl AzureAuthInputs { + pub fn default_credential_for_scope(scope: &str) -> Self { + Self { + azure_scope: ConfigValue::Value(Sourced::new( + scope.to_string(), + InputSource::Deployment, + )), + azure_credential: ConfigValue::Value(Sourced::new( + "DefaultAzureCredential".to_string(), + InputSource::Deployment, + )), + enable_azure_ad_token_refresh: Sourced::new(true, InputSource::Deployment), + ..Self::default() + } + } + pub fn or_configured_token_refresh(self, enabled: bool) -> Self { if *self.enable_azure_ad_token_refresh.value() || !enabled { return self; diff --git a/litellm-rust/crates/cache-azure-blob/Cargo.toml b/litellm-rust/crates/cache-azure-blob/Cargo.toml new file mode 100644 index 00000000000..55abaff1975 --- /dev/null +++ b/litellm-rust/crates/cache-azure-blob/Cargo.toml @@ -0,0 +1,22 @@ +[package] +name = "litellm-cache-azure-blob" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +litellm-auth-azure.workspace = true +litellm-auth-types.workspace = true +litellm-cache.workspace = true + +async-trait = "0.1" +azure_core = "1.1.0" +azure_storage_blob = "1.1.0" +futures-util.workspace = true +tokio.workspace = true +url.workspace = true + +[dev-dependencies] +litellm-cache-response.workspace = true +serde_json.workspace = true diff --git a/litellm-rust/crates/cache-azure-blob/src/cache.rs b/litellm-rust/crates/cache-azure-blob/src/cache.rs new file mode 100644 index 00000000000..f28b8c9a641 --- /dev/null +++ b/litellm-rust/crates/cache-azure-blob/src/cache.rs @@ -0,0 +1,246 @@ +use std::{sync::Arc, time::Duration}; + +use azure_core::{ + credentials::TokenCredential, + error::ErrorKind, + http::{ClientOptions, RequestContent}, +}; +use azure_storage_blob::{ + BlobContainerClient, BlobContainerClientOptions, + models::{BlobClientUploadOptions, StorageErrorCode}, +}; +use futures_util::{TryStreamExt, future::try_join_all}; +use litellm_cache::{ + BaseCache, BatchCache, CacheCodec, CacheConnectionResult, CacheConnectionStatus, Error, + ExactCacheContext, FlushCache, +}; +use tokio::runtime::Handle; +use url::Url; + +use crate::credential::AzureBlobCredential; + +/// Synchronous methods block on `runtime` and therefore must run outside of it +pub struct AzureBlobCache { + container: BlobContainerClient, + codec: C, + runtime: Handle, + account_url: String, + container_name: String, +} + +impl AzureBlobCache { + pub async fn connect( + account_url: &str, + container: &str, + codec: C, + runtime: Handle, + ) -> Result { + Self::connect_with_options( + account_url, + container, + Some(Arc::new(AzureBlobCredential::default())), + ClientOptions::default(), + codec, + runtime, + ) + .await + } + + pub async fn connect_with_options( + account_url: &str, + container: &str, + credential: Option>, + client_options: ClientOptions, + codec: C, + runtime: Handle, + ) -> Result { + let mut url = Url::parse(account_url).map_err(|_| Error::Unavailable)?; + let account_url = url.as_str().trim_end_matches('/').to_string(); + url.path_segments_mut() + .map_err(|()| Error::Unavailable)? + .pop_if_empty() + .push(container); + let client = BlobContainerClient::new( + url, + credential, + Some(BlobContainerClientOptions { + client_options, + ..BlobContainerClientOptions::default() + }), + ) + .map_err(|_| Error::Unavailable)?; + let cache = Self { + container: client, + codec, + runtime, + account_url, + container_name: container.to_string(), + }; + cache.create_container().await?; + Ok(cache) + } + + pub fn account_url(&self) -> &str { + &self.account_url + } + + pub fn container_name(&self) -> &str { + &self.container_name + } + + async fn create_container(&self) -> Result<(), Error> { + match self.container.create(None).await { + Ok(_) => Ok(()), + Err(error) if is_storage_error(&error, StorageErrorCode::ContainerAlreadyExists) => { + Ok(()) + } + Err(_) => Err(Error::Unavailable), + } + } + + async fn upload(&self, key: &str, value: &C::Value, overwrite: bool) -> Result<(), Error> { + let payload = self.codec.encode(value)?; + let options = (!overwrite).then(|| BlobClientUploadOptions::default().if_not_exists()); + match self + .container + .blob_client(key) + .upload(RequestContent::from(payload), options) + .await + { + Ok(_) => Ok(()), + Err(error) if is_storage_error(&error, StorageErrorCode::BlobAlreadyExists) => Ok(()), + Err(_) => Err(Error::Unavailable), + } + } + + async fn download(&self, key: &str) -> Result, Error> { + let response = match self.container.blob_client(key).download(None).await { + Ok(response) => response, + Err(error) if is_storage_error(&error, StorageErrorCode::BlobNotFound) => { + return Ok(None); + } + Err(_) => return Err(Error::Unavailable), + }; + let bytes = response + .body + .collect() + .await + .map_err(|_| Error::Unavailable)?; + self.codec.decode(&bytes).map(Some) + } + + async fn delete_all_blobs(&self) -> Result<(), Error> { + let mut pages = self + .container + .list_blobs(None) + .map_err(|_| Error::Unavailable)? + .into_pages(); + while let Some(page) = pages.try_next().await.map_err(|_| Error::Unavailable)? { + let page = page.into_model().map_err(|_| Error::Unavailable)?; + for name in page.blob_items.into_iter().filter_map(|item| item.name) { + self.container + .blob_client(&name) + .delete(None) + .await + .map_err(|_| Error::Unavailable)?; + } + } + Ok(()) + } + + fn block_on(&self, future: impl Future) -> T { + self.runtime.block_on(future) + } +} + +fn is_storage_error(error: &azure_core::Error, code: StorageErrorCode) -> bool { + matches!( + error.kind(), + ErrorKind::HttpResponse { + error_code: Some(error_code), + .. + } if error_code == code.as_ref() + ) +} + +impl BaseCache for AzureBlobCache { + type Value = C::Value; + type Context = ExactCacheContext; + + fn get_ttl(&self, _: &ExactCacheContext) -> Option { + None + } + + fn set_cache(&self, key: &str, value: C::Value, _: &ExactCacheContext) -> Result<(), Error> { + self.block_on(self.upload(key, &value, false)) + } + + fn get_cache(&self, key: &str, _: &ExactCacheContext) -> Result, Error> { + self.block_on(self.download(key)) + } + + async fn async_set_cache( + &self, + key: &str, + value: C::Value, + _: ExactCacheContext, + ) -> Result<(), Error> { + self.upload(key, &value, true).await + } + + async fn async_get_cache( + &self, + key: &str, + _: &ExactCacheContext, + ) -> Result, Error> { + self.download(key).await + } + + async fn async_set_cache_pipeline( + &self, + entries: Vec<(String, C::Value)>, + _: ExactCacheContext, + ) -> Result<(), Error> { + try_join_all( + entries + .iter() + .map(|(key, value)| self.upload(key, value, true)), + ) + .await + .map(drop) + } + + async fn disconnect(&self) -> Result<(), Error> { + Ok(()) + } + + async fn test_connection(&self) -> Result { + Ok(match self.container.get_properties(None).await { + Ok(_) => CacheConnectionResult { + status: CacheConnectionStatus::Success, + message: "Azure Blob cache connection test successful".into(), + error: None, + }, + Err(error) => CacheConnectionResult { + status: CacheConnectionStatus::Failed, + message: format!("Azure Blob connection failed: {error}"), + error: Some(error.to_string()), + }, + }) + } +} + +impl BatchCache for AzureBlobCache {} + +impl FlushCache for AzureBlobCache { + fn flush_cache(&self) -> Result<(), Error> { + self.block_on(self.delete_all_blobs()) + } + + async fn async_flush_cache(&self) -> Result<(), Error> { + self.delete_all_blobs().await + } +} + +#[cfg(test)] +mod tests; diff --git a/litellm-rust/crates/cache-azure-blob/src/cache/tests.rs b/litellm-rust/crates/cache-azure-blob/src/cache/tests.rs new file mode 100644 index 00000000000..fd116a0e28f --- /dev/null +++ b/litellm-rust/crates/cache-azure-blob/src/cache/tests.rs @@ -0,0 +1,692 @@ +use std::{ + collections::BTreeMap, + sync::{Arc, Mutex}, + time::Duration, +}; + +use azure_core::http::{ + AsyncRawResponse, Body, ClientOptions, HttpClient, Method, Request, StatusCode, Transport, + headers::{HeaderName, Headers}, +}; +use litellm_cache::{ + BaseCache, BatchCache, BatchEntry, CacheConnectionStatus, Error, ExactCacheContext, FlushCache, +}; +use litellm_cache_response::{ + CacheEntry, CacheKeyField, CacheKeyInput, ResponseCache, ResponseCacheCodec, + ResponseCacheRequest, cache_key, +}; +use serde_json::json; +use tokio::runtime::Runtime; + +use super::AzureBlobCache; + +const ACCOUNT_URL: &str = "https://example.blob.core.windows.net"; +const CONTAINER: &str = "litellm-cache"; +const IF_NONE_MATCH: HeaderName = HeaderName::from_static("if-none-match"); +const ERROR_CODE: HeaderName = HeaderName::from_static("x-ms-error-code"); + +#[derive(Clone, Debug, PartialEq, Eq)] +struct RecordedRequest { + method: Method, + path: String, + query: String, + if_none_match: Option, +} + +#[derive(Default)] +struct FakeState { + container_exists: bool, + blobs: BTreeMap>, + requests: Vec, + failing: bool, +} + +#[derive(Clone, Default)] +struct FakeBlobService { + state: Arc>, +} + +impl std::fmt::Debug for FakeBlobService { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str("FakeBlobService") + } +} + +impl FakeBlobService { + fn with_existing_container() -> Self { + let service = Self::default(); + service.state.lock().unwrap().container_exists = true; + service + } + + fn blob(&self, name: &str) -> Option> { + self.state.lock().unwrap().blobs.get(name).cloned() + } + + fn blob_names(&self) -> Vec { + self.state.lock().unwrap().blobs.keys().cloned().collect() + } + + fn seed_blob(&self, name: &str, bytes: &[u8]) { + self.state + .lock() + .unwrap() + .blobs + .insert(name.to_string(), bytes.to_vec()); + } + + fn set_failing(&self, failing: bool) { + self.state.lock().unwrap().failing = failing; + } + + fn requests(&self) -> Vec { + self.state.lock().unwrap().requests.clone() + } + + fn container_exists(&self) -> bool { + self.state.lock().unwrap().container_exists + } + + fn respond(status: StatusCode, error_code: Option<&str>, body: Vec) -> AsyncRawResponse { + let mut headers = Headers::new(); + if let Some(code) = error_code { + headers.insert(ERROR_CODE, code.to_string()); + } + AsyncRawResponse::from_bytes(status, headers, body) + } + + fn list_body(state: &FakeState) -> Vec { + let mut xml = String::from( + r#""#, + ); + for name in state.blobs.keys() { + xml.push_str(&format!( + "{name}BlockBlob" + )); + } + xml.push_str(""); + xml.into_bytes() + } +} + +#[async_trait::async_trait] +impl HttpClient for FakeBlobService { + async fn execute_request(&self, request: &Request) -> azure_core::Result { + let mut state = self.state.lock().unwrap(); + let path = request.url().path().to_string(); + let query = request.url().query().unwrap_or_default().to_string(); + let if_none_match = request + .headers() + .get_optional_str(&IF_NONE_MATCH) + .map(str::to_owned); + state.requests.push(RecordedRequest { + method: request.method(), + path: path.clone(), + query: query.clone(), + if_none_match: if_none_match.clone(), + }); + if state.failing { + return Ok(Self::respond( + StatusCode::Forbidden, + Some("AuthorizationFailure"), + Vec::new(), + )); + } + let container_path = format!("/{CONTAINER}"); + let blob_name = path + .strip_prefix(&format!("{container_path}/")) + .map(str::to_owned); + let is_container = path == container_path && query.contains("restype=container"); + let response = match (request.method(), is_container, blob_name) { + (Method::Put, true, None) if state.container_exists => Self::respond( + StatusCode::Conflict, + Some("ContainerAlreadyExists"), + Vec::new(), + ), + (Method::Put, true, None) => { + state.container_exists = true; + Self::respond(StatusCode::Created, None, Vec::new()) + } + (Method::Get, true, None) if query.contains("comp=list") => { + Self::respond(StatusCode::Ok, None, Self::list_body(&state)) + } + (Method::Get, true, None) if state.container_exists => { + Self::respond(StatusCode::Ok, None, Vec::new()) + } + (Method::Get, true, None) => { + Self::respond(StatusCode::NotFound, Some("ContainerNotFound"), Vec::new()) + } + (Method::Put, false, Some(name)) => { + if if_none_match.as_deref() == Some("*") && state.blobs.contains_key(&name) { + Self::respond(StatusCode::Conflict, Some("BlobAlreadyExists"), Vec::new()) + } else { + let bytes = match request.body() { + Body::Bytes(bytes) => bytes.to_vec(), + Body::SeekableStream(_) => panic!("unexpected streaming upload"), + }; + state.blobs.insert(name, bytes); + Self::respond(StatusCode::Created, None, Vec::new()) + } + } + (Method::Get, false, Some(name)) => match state.blobs.get(&name) { + Some(bytes) => Self::respond(StatusCode::Ok, None, bytes.clone()), + None => Self::respond(StatusCode::NotFound, Some("BlobNotFound"), Vec::new()), + }, + (Method::Delete, false, Some(name)) => match state.blobs.remove(&name) { + Some(_) => Self::respond(StatusCode::Accepted, None, Vec::new()), + None => Self::respond(StatusCode::NotFound, Some("BlobNotFound"), Vec::new()), + }, + (method, _, _) => panic!("unexpected request {method:?} {path}?{query}"), + }; + Ok(response) + } +} + +struct Fixture { + runtime: Runtime, + service: FakeBlobService, + cache: Arc>, +} + +impl Fixture { + fn new(service: FakeBlobService) -> Self { + let runtime = Runtime::new().unwrap(); + let cache = runtime + .block_on(Self::connect(&service, runtime.handle().clone())) + .unwrap(); + Self { + runtime, + service, + cache: Arc::new(cache), + } + } + + async fn connect( + service: &FakeBlobService, + handle: tokio::runtime::Handle, + ) -> Result, Error> { + AzureBlobCache::connect_with_options( + ACCOUNT_URL, + CONTAINER, + None, + ClientOptions { + transport: Some(Transport::new(Arc::new(service.clone()))), + ..ClientOptions::default() + }, + ResponseCacheCodec, + handle, + ) + .await + } + + fn response_cache(&self) -> ResponseCache> { + ResponseCache::new(self.cache.clone()) + } + + fn stored_json(&self, key: &str) -> serde_json::Value { + serde_json::from_slice(&self.service.blob(key).expect("blob should exist")).unwrap() + } +} + +fn request(model: &str) -> ResponseCacheRequest { + ResponseCacheRequest::new(CacheKeyInput { + fields: vec![CacheKeyField { + name: "model".into(), + value: Some(model.into()), + api_parameter: true, + internal_parameter: false, + }], + preset: None, + namespace: None, + include_provider_parameters: false, + }) +} + +fn now() -> Duration { + Duration::from_secs(1_700_000_000) +} + +fn entry(value: serde_json::Value) -> CacheEntry { + CacheEntry { + timestamp: Some(1_700_000_000.5), + response: value, + } +} + +fn no_ttl() -> ExactCacheContext { + ExactCacheContext::default() +} + +fn with_ttl(seconds: u64) -> ExactCacheContext { + ExactCacheContext { + ttl: Some(Duration::from_secs(seconds)), + } +} + +#[test] +fn connect_creates_the_container_once() { + let fixture = Fixture::new(FakeBlobService::default()); + assert!(fixture.service.container_exists()); + assert_eq!( + fixture.service.requests(), + vec![RecordedRequest { + method: Method::Put, + path: format!("/{CONTAINER}"), + query: "restype=container".into(), + if_none_match: None, + }] + ); + assert_eq!(fixture.cache.account_url(), ACCOUNT_URL); + assert_eq!(fixture.cache.container_name(), CONTAINER); +} + +#[test] +fn connect_accepts_an_existing_container() { + let fixture = Fixture::new(FakeBlobService::with_existing_container()); + assert!(fixture.service.container_exists()); + assert_eq!(fixture.service.requests().len(), 1); +} + +#[test] +fn connect_accepts_account_urls_with_trailing_slash() { + let runtime = Runtime::new().unwrap(); + let service = FakeBlobService::default(); + let cache = runtime + .block_on(AzureBlobCache::connect_with_options( + "https://example.blob.core.windows.net/", + CONTAINER, + None, + ClientOptions { + transport: Some(Transport::new(Arc::new(service.clone()))), + ..ClientOptions::default() + }, + ResponseCacheCodec, + runtime.handle().clone(), + )) + .unwrap(); + assert_eq!(service.requests()[0].path, format!("/{CONTAINER}")); + assert_eq!(cache.account_url(), "https://example.blob.core.windows.net"); +} + +#[test] +fn connect_surfaces_service_failures() { + let runtime = Runtime::new().unwrap(); + let service = FakeBlobService::default(); + service.set_failing(true); + let result = runtime.block_on(Fixture::connect(&service, runtime.handle().clone())); + assert!(matches!(result, Err(Error::Unavailable))); +} + +#[test] +fn sync_set_and_get_round_trip_python_json_shape() { + let fixture = Fixture::new(FakeBlobService::default()); + let value = entry(json!({"choices": [{"message": {"content": "héllo 🌍"}}]})); + fixture + .cache + .set_cache("key-1", value.clone(), &no_ttl()) + .unwrap(); + + assert_eq!( + fixture.stored_json("key-1"), + json!({ + "timestamp": 1_700_000_000.5, + "response": {"choices": [{"message": {"content": "héllo 🌍"}}]} + }) + ); + assert_eq!( + fixture.cache.get_cache("key-1", &no_ttl()).unwrap(), + Some(value) + ); +} + +#[test] +fn sync_set_does_not_overwrite_an_existing_blob() { + let fixture = Fixture::new(FakeBlobService::default()); + fixture + .cache + .set_cache("key", entry(json!({"v": "first"})), &no_ttl()) + .unwrap(); + fixture + .cache + .set_cache("key", entry(json!({"v": "second"})), &no_ttl()) + .unwrap(); + + assert_eq!( + fixture.stored_json("key")["response"], + json!({"v": "first"}) + ); + let uploads: Vec<_> = fixture + .service + .requests() + .into_iter() + .filter(|request| request.method == Method::Put && request.path.ends_with("/key")) + .collect(); + assert_eq!(uploads.len(), 2); + assert!( + uploads + .iter() + .all(|request| request.if_none_match.as_deref() == Some("*")) + ); +} + +#[test] +fn async_set_overwrites_an_existing_blob() { + let fixture = Fixture::new(FakeBlobService::default()); + fixture.runtime.block_on(async { + fixture + .cache + .async_set_cache("key", entry(json!({"v": "first"})), no_ttl()) + .await + .unwrap(); + fixture + .cache + .async_set_cache("key", entry(json!({"v": "second"})), no_ttl()) + .await + .unwrap(); + assert_eq!( + fixture + .cache + .async_get_cache("key", &no_ttl()) + .await + .unwrap(), + Some(entry(json!({"v": "second"}))) + ); + }); + assert_eq!( + fixture.stored_json("key")["response"], + json!({"v": "second"}) + ); + assert!( + fixture + .service + .requests() + .iter() + .filter(|request| request.method == Method::Put && request.path.ends_with("/key")) + .all(|request| request.if_none_match.is_none()) + ); +} + +#[test] +fn missing_blobs_are_misses() { + let fixture = Fixture::new(FakeBlobService::default()); + assert_eq!(fixture.cache.get_cache("absent", &no_ttl()).unwrap(), None); + assert_eq!( + fixture + .runtime + .block_on(fixture.cache.async_get_cache("absent", &no_ttl())) + .unwrap(), + None + ); +} + +#[test] +fn ttl_is_ignored_and_entries_never_expire() { + let fixture = Fixture::new(FakeBlobService::default()); + assert_eq!(fixture.cache.get_ttl(&with_ttl(1)), None); + assert_eq!(fixture.cache.get_ttl(&no_ttl()), None); + + fixture + .cache + .set_cache("key", entry(json!("value")), &with_ttl(1)) + .unwrap(); + std::thread::sleep(Duration::from_millis(1100)); + assert_eq!( + fixture.cache.get_cache("key", &with_ttl(1)).unwrap(), + Some(entry(json!("value"))) + ); + assert!( + fixture + .service + .requests() + .iter() + .all(|request| !request.query.contains("expiry")) + ); +} + +#[test] +fn malformed_blobs_are_invalid_entries_and_response_cache_misses() { + let fixture = Fixture::new(FakeBlobService::default()); + fixture.service.seed_blob("broken-json", b"{not json"); + fixture + .service + .seed_blob("broken-utf8", &[0xff, 0xfe, 0x22]); + fixture + .service + .seed_blob("wrong-shape", br#"{"timestamp": "yesterday"}"#); + + for key in ["broken-json", "broken-utf8", "wrong-shape"] { + assert!(matches!( + fixture.cache.get_cache(key, &no_ttl()), + Err(Error::InvalidEntry) + )); + } + + let response_cache = fixture.response_cache(); + let broken = request("broken"); + fixture + .service + .seed_blob(&cache_key(&broken.key), b"{not json"); + assert_eq!(response_cache.lookup(&broken, now()).unwrap(), None); + assert_eq!( + fixture + .runtime + .block_on(response_cache.async_lookup(&broken, now())) + .unwrap(), + None + ); +} + +#[test] +fn batch_get_preserves_order_and_marks_misses_and_invalid_entries() { + let fixture = Fixture::new(FakeBlobService::default()); + fixture + .cache + .set_cache("a", entry(json!("A")), &no_ttl()) + .unwrap(); + fixture + .cache + .set_cache("c", entry(json!("C")), &no_ttl()) + .unwrap(); + fixture.service.seed_blob("bad", b"nope"); + let keys = ["c", "missing", "a", "bad"].map(String::from); + + let sync = fixture.cache.batch_get_cache(&keys, &no_ttl()).unwrap(); + assert_eq!( + sync, + vec![ + BatchEntry::Hit(entry(json!("C"))), + BatchEntry::Miss, + BatchEntry::Hit(entry(json!("A"))), + BatchEntry::Invalid, + ] + ); + + let asynchronous = fixture + .runtime + .block_on(fixture.cache.async_batch_get_cache(keys.to_vec(), no_ttl())) + .unwrap(); + assert_eq!(asynchronous, sync); + + let response_cache = fixture.response_cache(); + let requests = [request("hit"), request("missing"), request("bad")]; + response_cache + .store(&requests[0], json!("HIT"), now()) + .unwrap(); + fixture + .service + .seed_blob(&cache_key(&requests[2].key), b"nope"); + let hits = response_cache.lookup_batch(&requests, now()).unwrap(); + assert_eq!(hits.values, vec![Some(json!("HIT")), None, None]); + assert_eq!(hits.missing_indices, vec![1, 2]); + let async_hits = fixture + .runtime + .block_on(response_cache.async_lookup_batch(&requests, now())) + .unwrap(); + assert_eq!(async_hits.values, hits.values); +} + +#[test] +fn async_pipeline_writes_every_entry_with_overwrite() { + let fixture = Fixture::new(FakeBlobService::default()); + fixture.service.seed_blob("k2", b"stale"); + fixture + .runtime + .block_on(fixture.cache.async_set_cache_pipeline( + vec![ + ("k1".into(), entry(json!({"n": 1}))), + ("k2".into(), entry(json!({"n": 2}))), + ("k3".into(), entry(json!({"n": 3}))), + ], + with_ttl(30), + )) + .unwrap(); + assert_eq!(fixture.service.blob_names(), ["k1", "k2", "k3"]); + assert_eq!(fixture.stored_json("k2")["response"], json!({"n": 2})); +} + +#[test] +fn flush_deletes_every_blob_in_the_container() { + let fixture = Fixture::new(FakeBlobService::default()); + for key in ["x", "y", "z"] { + fixture + .cache + .set_cache(key, entry(json!(key)), &no_ttl()) + .unwrap(); + } + fixture.cache.flush_cache().unwrap(); + assert!(fixture.service.blob_names().is_empty()); + assert!(fixture.service.container_exists()); + + fixture + .cache + .set_cache("again", entry(json!(1)), &no_ttl()) + .unwrap(); + fixture + .runtime + .block_on(fixture.cache.async_flush_cache()) + .unwrap(); + assert!(fixture.service.blob_names().is_empty()); +} + +#[test] +fn service_failures_map_to_unavailable() { + let fixture = Fixture::new(FakeBlobService::default()); + fixture.service.set_failing(true); + assert!(matches!( + fixture.cache.get_cache("key", &no_ttl()), + Err(Error::Unavailable) + )); + assert!(matches!( + fixture.cache.set_cache("key", entry(json!(1)), &no_ttl()), + Err(Error::Unavailable) + )); + assert!(matches!( + fixture.cache.flush_cache(), + Err(Error::Unavailable) + )); + assert!(matches!( + fixture.runtime.block_on( + fixture + .cache + .async_set_cache_pipeline(vec![("k".into(), entry(json!(1)))], no_ttl()) + ), + Err(Error::Unavailable) + )); +} + +#[test] +fn test_connection_reports_container_reachability() { + let fixture = Fixture::new(FakeBlobService::default()); + let ok = fixture + .runtime + .block_on(fixture.cache.test_connection()) + .unwrap(); + assert_eq!(ok.status, CacheConnectionStatus::Success); + assert!(ok.error.is_none()); + + fixture.service.set_failing(true); + let failed = fixture + .runtime + .block_on(fixture.cache.test_connection()) + .unwrap(); + assert_eq!(failed.status, CacheConnectionStatus::Failed); + assert!(failed.error.is_some()); +} + +#[test] +fn disconnect_is_idempotent_and_keeps_data() { + let fixture = Fixture::new(FakeBlobService::default()); + fixture + .cache + .set_cache("key", entry(json!(1)), &no_ttl()) + .unwrap(); + fixture.runtime.block_on(async { + fixture.cache.disconnect().await.unwrap(); + fixture.cache.disconnect().await.unwrap(); + }); + assert_eq!( + fixture.cache.get_cache("key", &no_ttl()).unwrap(), + Some(entry(json!(1))) + ); +} + +#[test] +fn response_cache_stores_and_reads_through_the_backend() { + let fixture = Fixture::new(FakeBlobService::default()); + let response_cache = fixture.response_cache(); + let mut request = request("gpt"); + request.context = with_ttl(60); + let response = json!({"id": "chatcmpl-1"}); + response_cache + .store(&request, response.clone(), now()) + .unwrap(); + assert_eq!( + fixture.stored_json(&cache_key(&request.key)), + json!({"timestamp": 1_700_000_000.0, "response": {"id": "chatcmpl-1"}}) + ); + assert_eq!( + response_cache + .lookup(&request, now() + Duration::from_secs(3600)) + .unwrap(), + Some(response.clone()) + ); + assert_eq!( + fixture + .runtime + .block_on(response_cache.async_lookup(&request, now() + Duration::from_secs(3600))) + .unwrap(), + Some(response.clone()) + ); + fixture.runtime.block_on(async { + response_cache + .async_store(&request, json!("replaced"), now()) + .await + .unwrap(); + assert_eq!( + response_cache.async_lookup(&request, now()).await.unwrap(), + Some(json!("replaced")) + ); + response_cache.async_flush().await.unwrap(); + assert_eq!( + response_cache.async_lookup(&request, now()).await.unwrap(), + None + ); + }); +} + +#[test] +fn non_object_responses_are_written_serialized_like_python() { + let fixture = Fixture::new(FakeBlobService::default()); + fixture + .cache + .set_cache("s", entry(json!("plain")), &no_ttl()) + .unwrap(); + assert_eq!( + fixture.stored_json("s"), + json!({"timestamp": 1_700_000_000.5, "response": "\"plain\""}) + ); + assert_eq!( + fixture.cache.get_cache("s", &no_ttl()).unwrap(), + Some(entry(json!("plain"))) + ); +} diff --git a/litellm-rust/crates/cache-azure-blob/src/credential.rs b/litellm-rust/crates/cache-azure-blob/src/credential.rs new file mode 100644 index 00000000000..d1a3d0e44ec --- /dev/null +++ b/litellm-rust/crates/cache-azure-blob/src/credential.rs @@ -0,0 +1,84 @@ +use std::{ + fmt, + sync::Arc, + time::{Duration, SystemTime}, +}; + +use azure_core::{ + credentials::{AccessToken, TokenCredential, TokenRequestOptions}, + error::ErrorKind, + time::OffsetDateTime, +}; +use litellm_auth_azure::{AzureAuthInputs, AzureAuthService}; +use litellm_auth_types::ResolvedCredential; + +const STATIC_TOKEN_LIFETIME: Duration = Duration::from_secs(300); +const LLM_TOKEN_ENV: &str = "AZURE_AD_TOKEN"; + +type EnvLookup = Arc Option + Send + Sync>; + +pub struct AzureBlobCredential { + service: AzureAuthService, + env_lookup: EnvLookup, +} + +impl fmt::Debug for AzureBlobCredential { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str("AzureBlobCredential") + } +} + +impl Default for AzureBlobCredential { + fn default() -> Self { + Self::new( + AzureAuthService::default(), + Arc::new(|name| std::env::var(name).ok()), + ) + } +} + +impl AzureBlobCredential { + pub fn new(service: AzureAuthService, env_lookup: EnvLookup) -> Self { + Self { + service, + env_lookup, + } + } +} + +#[async_trait::async_trait] +impl TokenCredential for AzureBlobCredential { + async fn get_token( + &self, + scopes: &[&str], + _options: Option>, + ) -> azure_core::Result { + let env_lookup = &self.env_lookup; + let lookup = move |name: &str| (name != LLM_TOKEN_ENV).then(|| env_lookup(name)).flatten(); + let credential = self + .service + .get_azure_ad_token( + &AzureAuthInputs::default_credential_for_scope(&scopes.join(" ")), + &lookup, + ) + .await + .map_err(|error| { + azure_core::Error::with_message(ErrorKind::Credential, error.to_string()) + })? + .ok_or_else(|| { + azure_core::Error::with_message( + ErrorKind::Credential, + "no Azure credential is available for blob storage", + ) + })?; + let (token, expires_on) = match credential.into_value() { + ResolvedCredential::AccessToken { token, expires_on } => (token, expires_on), + ResolvedCredential::Static(token) => (token, None), + }; + let expires_on = expires_on.unwrap_or_else(|| SystemTime::now() + STATIC_TOKEN_LIFETIME); + Ok(AccessToken::new( + token.expose().to_string(), + OffsetDateTime::from(expires_on), + )) + } +} diff --git a/litellm-rust/crates/cache-azure-blob/src/lib.rs b/litellm-rust/crates/cache-azure-blob/src/lib.rs new file mode 100644 index 00000000000..5ae752c111d --- /dev/null +++ b/litellm-rust/crates/cache-azure-blob/src/lib.rs @@ -0,0 +1,5 @@ +mod cache; +mod credential; + +pub use cache::AzureBlobCache; +pub use credential::AzureBlobCredential; diff --git a/litellm-rust/crates/cache-azure-blob/src/tests.rs b/litellm-rust/crates/cache-azure-blob/src/tests.rs new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 1eb2ec28036..7e48874a30b 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -21,6 +21,7 @@ tiktoken = ["litellm-token-counter/tiktoken"] [dependencies] bytes.workspace = true litellm-cache.workspace = true +litellm-cache-azure-blob.workspace = true litellm-cache-memory.workspace = true litellm-cache-redis.workspace = true litellm-cache-response.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/cache/config.rs b/litellm-rust/crates/python-bridge/src/cache/config.rs index 0e7d6aee11d..91f56157bdc 100644 --- a/litellm-rust/crates/python-bridge/src/cache/config.rs +++ b/litellm-rust/crates/python-bridge/src/cache/config.rs @@ -73,9 +73,15 @@ pub(super) struct RedisCacheConfig { pub(super) connection: RedisConnectionConfig, } +pub(super) struct AzureBlobCacheConfig { + pub(super) account_url: String, + pub(super) container: String, +} + pub(super) enum CacheBackendConfig { Memory(MemoryCacheConfig), Redis(Box), + AzureBlob(AzureBlobCacheConfig), } #[allow(dead_code, reason = "consumed by the cache activation follow-up")] @@ -142,13 +148,18 @@ impl NativeCacheConfig { }))), Err(reason) => Ok(CacheConfigProjection::Unsupported(reason)), }, + Some(CacheType::AzureBlob) => project_azure_blob(&backend).map(|backend| { + CacheConfigProjection::Native(Box::new(Self { + policy, + backend: CacheBackendConfig::AzureBlob(backend), + })) + }), Some( CacheType::RedisSemantic | CacheType::ValkeySemantic | CacheType::S3 | CacheType::Disk | CacheType::QdrantSemantic - | CacheType::AzureBlob | CacheType::Gcs, ) | None => Ok(CacheConfigProjection::Unsupported( @@ -158,12 +169,12 @@ impl NativeCacheConfig { } pub(super) fn service_mismatch(&self, service: &NativeResponseCache) -> Option<&'static str> { - if service.default_ttl() - != Some(match &self.backend { - CacheBackendConfig::Memory(config) => config.default_ttl, - CacheBackendConfig::Redis(config) => config.default_ttl, - }) - { + let default_ttl = match &self.backend { + CacheBackendConfig::Memory(config) => Some(config.default_ttl), + CacheBackendConfig::Redis(config) => Some(config.default_ttl), + CacheBackendConfig::AzureBlob(_) => None, + }; + if service.default_ttl() != default_ttl { return Some("facade and native backend default TTLs must match"); } match &self.backend { @@ -185,10 +196,34 @@ impl NativeCacheConfig { CacheBackendConfig::Redis(config) => (service.namespace() != config.namespace.as_deref()) .then_some("facade and native backend namespaces must match"), + CacheBackendConfig::AzureBlob(config) => match service.azure_blob_identity() { + None => Some("facade and native backend types must match"), + Some((account_url, container)) + if account_url != config.account_url || container != config.container => + { + Some("facade and native backend containers must match") + } + Some(_) => None, + }, } } } +#[inline(never)] +fn project_azure_blob(backend: &Bound<'_, PyAny>) -> PyResult { + let client = backend.getattr("container_client")?; + let container = client.getattr("container_name")?.extract::()?; + let url = client.getattr("url")?.extract::()?; + let account_url = url + .strip_suffix(container.as_str()) + .and_then(|url| url.strip_suffix('/')) + .ok_or_else(|| PyValueError::new_err("Azure Blob container URL is malformed"))?; + Ok(AzureBlobCacheConfig { + account_url: account_url.to_string(), + container, + }) +} + #[inline(never)] fn project_memory(backend: &Bound<'_, PyAny>) -> PyResult { let max_size_kib = backend.getattr("max_size_per_item")?.extract::()?; diff --git a/litellm-rust/crates/python-bridge/src/cache/facade.rs b/litellm-rust/crates/python-bridge/src/cache/facade.rs index f2f86c14b37..1df6312503f 100644 --- a/litellm-rust/crates/python-bridge/src/cache/facade.rs +++ b/litellm-rust/crates/python-bridge/src/cache/facade.rs @@ -32,10 +32,23 @@ struct RedisPoolGuard { max_connections: usize, } +struct AzureBlobClientGuard { + sync_client: Py, + async_client: Py, + url: String, + container_name: String, +} + +enum ConnectionGuard { + None, + RedisPool(RedisPoolGuard), + AzureBlob(AzureBlobClientGuard), +} + pub(super) struct FacadeGuard { outer: ObjectGuard, backend: ObjectGuard, - redis_pool: Option, + connection: ConnectionGuard, } impl ObjectGuard { @@ -176,6 +189,60 @@ impl RedisPoolGuard { } } +impl AzureBlobClientGuard { + fn capture(backend: &Bound<'_, PyAny>) -> PyResult { + let sync_client = backend.getattr("container_client")?; + Ok(Self { + url: sync_client.getattr("url")?.extract::()?, + container_name: sync_client.getattr("container_name")?.extract::()?, + sync_client: sync_client.unbind(), + async_client: backend.getattr("async_container_client")?.unbind(), + }) + } + + fn matches(&self, py: Python<'_>, backend: &Bound<'_, PyAny>) -> PyResult { + let sync_client = backend.getattr("container_client")?; + Ok(self.sync_client.bind(py).is(&sync_client) + && self + .async_client + .bind(py) + .is(&backend.getattr("async_container_client")?) + && self.url == sync_client.getattr("url")?.extract::()? + && self.container_name == sync_client.getattr("container_name")?.extract::()?) + } + + fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + visit.call(&self.sync_client)?; + visit.call(&self.async_client) + } +} + +impl ConnectionGuard { + fn capture(kind: &str, backend: &Bound<'_, PyAny>) -> PyResult { + Ok(match kind { + "redis" => Self::RedisPool(RedisPoolGuard::capture(backend)?), + "azure-blob" => Self::AzureBlob(AzureBlobClientGuard::capture(backend)?), + _ => Self::None, + }) + } + + fn matches(&self, py: Python<'_>, backend: &Bound<'_, PyAny>) -> PyResult { + match self { + Self::None => Ok(true), + Self::RedisPool(guard) => guard.matches(py, backend), + Self::AzureBlob(guard) => guard.matches(py, backend), + } + } + + fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + match self { + Self::None => Ok(()), + Self::RedisPool(guard) => guard.traverse(visit), + Self::AzureBlob(guard) => guard.traverse(visit), + } + } +} + impl FacadeGuard { pub(super) fn capture( py: Python<'_>, @@ -192,6 +259,11 @@ impl FacadeGuard { let (module, name, cache_kind) = match kind { "memory" => ("litellm.caching.in_memory_cache", "InMemoryCache", "local"), "redis" => ("litellm.caching.redis_cache", "RedisCache", "redis"), + "azure-blob" => ( + "litellm.caching.azure_blob_cache", + "AzureBlobCache", + "azure-blob", + ), _ => unreachable!(), }; let backend = facade.getattr("cache")?; @@ -237,9 +309,7 @@ impl FacadeGuard { "redis_flush_size", ], )?, - redis_pool: (kind == "redis") - .then(|| RedisPoolGuard::capture(&backend)) - .transpose()?, + connection: ConnectionGuard::capture(kind, &backend)?, }) } @@ -251,19 +321,13 @@ impl FacadeGuard { if !self.backend.matches(py, &backend)? { return Ok(false); } - match &self.redis_pool { - Some(guard) => guard.matches(py, &backend), - None => Ok(true), - } + self.connection.matches(py, &backend) } pub(super) fn traverse(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { self.outer.traverse(&visit)?; self.backend.traverse(&visit)?; - if let Some(guard) = &self.redis_pool { - guard.traverse(&visit)?; - } - Ok(()) + self.connection.traverse(&visit) } } diff --git a/litellm-rust/crates/python-bridge/src/cache/handle.rs b/litellm-rust/crates/python-bridge/src/cache/handle.rs index 8251b3df06c..69988980d64 100644 --- a/litellm-rust/crates/python-bridge/src/cache/handle.rs +++ b/litellm-rust/crates/python-bridge/src/cache/handle.rs @@ -1,4 +1,4 @@ -use litellm_host_python::release_gil; +use litellm_host_python::{release_gil, run_sync_value}; use pyo3::{PyTraverseError, PyVisit, exceptions::PyRuntimeError, prelude::*}; use super::{cache_error, facade::FacadeGuard, native::NativeResponseCache, request::duration}; @@ -51,6 +51,21 @@ impl CacheTestHandle { }) } + #[staticmethod] + #[pyo3(signature = (account_url, container))] + fn azure_blob(py: Python<'_>, account_url: String, container: String) -> PyResult { + let service = run_sync_value(py, async move { + NativeResponseCache::azure_blob(&account_url, &container) + .await + .map_err(cache_error) + })?; + Ok(Self { + service, + guard: None, + pid: std::process::id(), + }) + } + #[getter] fn backend(&self) -> &'static str { self.service.kind() diff --git a/litellm-rust/crates/python-bridge/src/cache/native.rs b/litellm-rust/crates/python-bridge/src/cache/native.rs index a9475429e45..08c65905e84 100644 --- a/litellm-rust/crates/python-bridge/src/cache/native.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native.rs @@ -1,6 +1,7 @@ use std::{sync::Arc, time::Duration}; use litellm_cache::{CacheCodec, CacheConnectionResult, Error}; +use litellm_cache_azure_blob::AzureBlobCache; use litellm_cache_memory::InMemoryCache; use litellm_cache_redis::RedisCache; use litellm_cache_response::{ @@ -15,6 +16,7 @@ pub(super) enum NativeResponseCache { cache: Arc>>, buffer: Option>, }, + AzureBlob(Arc>>), } impl NativeResponseCache { @@ -43,6 +45,29 @@ impl NativeResponseCache { buffer: None, }) } + + pub async fn azure_blob(account_url: &str, container: &str) -> Result { + let backend = AzureBlobCache::connect( + account_url, + container, + ResponseCacheCodec, + tokio::runtime::Handle::current(), + ) + .await?; + Ok(Self::AzureBlob(Arc::new(ResponseCache::new(Arc::new( + backend, + ))))) + } + + pub fn azure_blob_identity(&self) -> Option<(&str, &str)> { + match self { + Self::AzureBlob(cache) => Some(( + cache.backend().account_url(), + cache.backend().container_name(), + )), + Self::Memory(_) | Self::Redis { .. } => None, + } + } } impl NativeResponseCache { @@ -50,6 +75,7 @@ impl NativeResponseCache { match self { Self::Memory(_) => "memory", Self::Redis { .. } => "redis", + Self::AzureBlob(_) => "azure-blob", } } @@ -57,12 +83,13 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.default_ttl(), Self::Redis { cache, .. } => cache.default_ttl(), + Self::AzureBlob(cache) => cache.default_ttl(), } } pub fn namespace(&self) -> Option<&str> { match self { - Self::Memory(_) => None, + Self::Memory(_) | Self::AzureBlob(_) => None, Self::Redis { cache, .. } => cache.backend().namespace(), } } @@ -70,14 +97,14 @@ impl NativeResponseCache { pub fn capacity(&self) -> Option { match self { Self::Memory(cache) => Some(cache.backend().max_size_in_memory()), - Self::Redis { .. } => None, + Self::Redis { .. } | Self::AzureBlob(_) => None, } } pub fn max_entry_bytes(&self) -> Option { match self { Self::Memory(cache) => cache.backend().max_entry_bytes(), - Self::Redis { .. } => None, + Self::Redis { .. } | Self::AzureBlob(_) => None, } } @@ -87,7 +114,7 @@ impl NativeResponseCache { cache, buffer: flush_size.map(|flush_size| Arc::new(WriteBuffer::new(flush_size))), }, - memory => memory, + other => other, } } @@ -99,6 +126,7 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.lookup(request, now), Self::Redis { cache, .. } => cache.lookup(request, now), + Self::AzureBlob(cache) => cache.lookup(request, now), } } @@ -111,6 +139,7 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.store(request, response, now), Self::Redis { cache, .. } => cache.store(request, response, now), + Self::AzureBlob(cache) => cache.store(request, response, now), } } @@ -122,6 +151,7 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.lookup_batch(requests, now), Self::Redis { cache, .. } => cache.lookup_batch(requests, now), + Self::AzureBlob(cache) => cache.lookup_batch(requests, now), } } @@ -133,6 +163,7 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.async_lookup(request, now).await, Self::Redis { cache, .. } => cache.async_lookup(request, now).await, + Self::AzureBlob(cache) => cache.async_lookup(request, now).await, } } @@ -152,6 +183,7 @@ impl NativeResponseCache { cache, buffer: Some(buffer), } => buffer.async_store(cache, request, response, now).await, + Self::AzureBlob(cache) => cache.async_store(request, response, now).await, } } @@ -163,6 +195,7 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.async_lookup_batch(requests, now).await, Self::Redis { cache, .. } => cache.async_lookup_batch(requests, now).await, + Self::AzureBlob(cache) => cache.async_lookup_batch(requests, now).await, } } @@ -174,6 +207,7 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.async_store_batch(entries, now).await, Self::Redis { cache, .. } => cache.async_store_batch(entries, now).await, + Self::AzureBlob(cache) => cache.async_store_batch(entries, now).await, } } @@ -186,6 +220,7 @@ impl NativeResponseCache { } cache.async_flush().await } + Self::AzureBlob(cache) => cache.async_flush().await, } } @@ -193,6 +228,7 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.test_connection().await, Self::Redis { cache, .. } => cache.test_connection().await, + Self::AzureBlob(cache) => cache.test_connection().await, } } } diff --git a/tests/test_litellm_rust/test_cache.py b/tests/test_litellm_rust/test_cache.py index c35cb1a20fb..99e371d8c2e 100644 --- a/tests/test_litellm_rust/test_cache.py +++ b/tests/test_litellm_rust/test_cache.py @@ -2,8 +2,10 @@ import asyncio import contextvars import gc import json +import os import threading import time +import uuid import weakref from collections.abc import Generator from types import SimpleNamespace @@ -13,8 +15,10 @@ from urllib.parse import urlparse import fakeredis import pytest import redis +from azure.storage.blob import ContainerClient import litellm +from litellm.caching.azure_blob_cache import AzureBlobCache from litellm.caching.caching import Cache, disable_cache, enable_cache, update_cache from litellm.caching.in_memory_cache import InMemoryCache from litellm.rust_bridge import _native @@ -45,6 +49,36 @@ def redis_url() -> Generator[str]: worker.join(timeout=5) +@pytest.fixture +def azure_blob_facade() -> Generator[Cache]: + account_url: Final = os.environ.get("AZURE_BLOB_CACHE_ACCOUNT_URL") + if account_url is None: + pytest.skip( + "live Azure Blob parity needs AZURE_BLOB_CACHE_ACCOUNT_URL plus DefaultAzureCredential inputs in the environment" + ) + facade: Final = Cache( + type=LiteLLMCacheType.AZURE_BLOB, + azure_account_url=account_url, + azure_blob_container=f"litellm-parity-{uuid.uuid4().hex[:12]}", + ) + backend: Final = facade.cache + assert isinstance(backend, AzureBlobCache) + try: + yield facade + finally: + backend.container_client.delete_container() + asyncio.run(backend.disconnect()) + + +def azure_blob_handle(facade: Cache) -> _native._CacheTestHandle: + backend: Final = facade.cache + assert isinstance(backend, AzureBlobCache) + return _native._CacheTestHandle.azure_blob( + backend.container_client.url.removesuffix(f"/{backend.container_client.container_name}"), + backend.container_client.container_name, + ) + + def test_existing_constructor_and_global_are_unchanged() -> None: facade: Final = Cache(type=LiteLLMCacheType.LOCAL) assert type(facade.cache) is InMemoryCache @@ -361,6 +395,89 @@ def test_facade_registration_rejects_mismatched_capacity() -> None: _native._CacheTestHandle.memory(capacity=7)._bind_facade(facade) +def test_azure_blob_facade_serves_natively_and_python_reads_the_same_blobs(azure_blob_facade: Cache) -> None: + backend: Final = azure_blob_facade.cache + assert isinstance(backend, AzureBlobCache) + handle: Final = azure_blob_handle(azure_blob_facade) + assert handle.backend == "azure-blob" + account_url: Final = backend.container_client.url.removesuffix(f"/{backend.container_client.container_name}") + with pytest.raises(TypeError, match="containers must match"): + _native._CacheTestHandle.azure_blob(account_url, f"{backend.container_client.container_name}-other")._bind_facade( + azure_blob_facade + ) + handle._bind_facade(azure_blob_facade) + resolver: Final = _native._CacheTestResolver(SimpleNamespace(cache=azure_blob_facade)) + native: Final = resolver.resolve() + assert native.kind == "native" + + response: Final = {"choices": [{"text": "caf\u00e9 \u2603"}], "usage": {"total_tokens": 3}, "flag": True, "empty": None} + native.store({**request("sync"), "ttl_seconds": 0.001}, response) + native.store(request("sync"), {"choices": [{"text": "second"}]}) + time.sleep(0.01) + stored: Final = json.loads(backend.container_client.download_blob("sync").readall()) + assert stored["response"] == response + assert isinstance(stored["timestamp"], float) + assert native.lookup(request("sync")) == response + assert cast(CacheLookup, azure_blob_facade).get_cache(cache_key="sync") == response + + backend.set_cache("python", {"timestamp": time.time(), "response": response}) + backend.set_cache("legacy", "bare legacy value") + backend.container_client.upload_blob("invalid", b"{not json", overwrite=True) + assert native.lookup(request("python")) == response + assert native.lookup(request("legacy")) == cast(CacheLookup, azure_blob_facade).get_cache(cache_key="legacy") + assert native.lookup_batch([request("python"), request("missing"), request("invalid"), request("sync")]) == { + "values": [response, None, None, response], + "missing_indices": [1, 2], + } + + with rebound(azure_blob_facade, "ttl", 12): + assert resolver.resolve().kind == "python_callback" + with rebound(backend, "container_client", ContainerClient.from_container_url(backend.container_client.url)): + assert resolver.resolve().kind == "python_callback" + + def custom_get(*_args: object, **_kwargs: object) -> None: + return None + + with rebound(backend, "get_cache", custom_get): + assert resolver.resolve().kind == "python_callback" + assert resolver.resolve().kind == "python_callback" + assert cast(CacheLookup, azure_blob_facade).get_cache(cache_key="sync") == response + + class CustomBlobCache(AzureBlobCache): + pass + + with rebound(azure_blob_facade, "cache", CustomBlobCache(account_url, backend.container_client.container_name)): + assert resolver.resolve().kind == "python_callback" + with pytest.raises(TypeError): + azure_blob_handle(azure_blob_facade)._bind_facade(azure_blob_facade) + + +async def test_azure_blob_native_async_writes_overwrite_batch_and_flush_like_python(azure_blob_facade: Cache) -> None: + backend: Final = azure_blob_facade.cache + assert isinstance(backend, AzureBlobCache) + azure_blob_handle(azure_blob_facade)._bind_facade(azure_blob_facade) + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=azure_blob_facade)).resolve() + assert binding.kind == "native" + ping: Final = cast(dict[str, object], await binding.ping()) + assert ping["status"] == "success", ping + + await binding.async_store(request("async"), {"value": 1}) + await binding.async_store({**request("async"), "ttl_seconds": 0.001}, {"value": 2}) + time.sleep(0.01) + assert await binding.async_lookup(request("async")) == {"value": 2} + assert await backend.async_get_cache("async") == json.loads(backend.container_client.download_blob("async").readall()) + assert cast(CacheLookup, azure_blob_facade).get_cache(cache_key="async") == {"value": 2} + + await binding.async_store_batch([request("first"), request("second")], [{"value": 3}, {"value": 4}]) + assert await binding.async_lookup_batch([request("second"), request("missing"), request("first")]) == { + "values": [{"value": 4}, None, {"value": 3}], + "missing_indices": [1], + } + await binding.async_flush() + assert [blob.name for blob in backend.container_client.list_blobs()] == [] + assert await binding.async_lookup(request("async")) is None + + async def test_redis_facade_buffers_native_async_writes(redis_url: str) -> None: parsed: Final = urlparse(redis_url) with rebound(litellm, "default_redis_ttl", 60): From b5548082c621f536b6a9ab243508f40d9da34fad Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 21:12:07 +0000 Subject: [PATCH 12/15] build(rust): raise native extension size limit to 30 MB for the Azure Blob SDK Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/scripts/verify_linux_native_wheel.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/scripts/verify_linux_native_wheel.py b/.github/scripts/verify_linux_native_wheel.py index 4fb8f068eb0..0adbc015ad0 100644 --- a/.github/scripts/verify_linux_native_wheel.py +++ b/.github/scripts/verify_linux_native_wheel.py @@ -205,7 +205,7 @@ def main( native_module: Final = load_native_module(native_path) native_module_loads: Final = native_module is not None panic_test_hook_absent: Final = native_module is not None and not hasattr(native_module, "_panic_for_test") - native_size_limit: Final = 25_000_000 + native_size_limit: Final = 30_000_000 native_size_within_limit: Final = native_member.file_size <= native_size_limit validations: Final = ( (f"Python tag is {EXPECTED_PYTHON_TAG}", python_tag == EXPECTED_PYTHON_TAG), @@ -222,7 +222,7 @@ def main( ("Python extension entry point is present", extension_entry_point_present), ("Native module loads", native_module_loads), ("Production module omits the panic test hook", panic_test_hook_absent), - ("Native extension does not exceed 25 MB", native_size_within_limit), + ("Native extension does not exceed 30 MB", native_size_within_limit), ("Wheel contents are valid", not unexpected_members), ) @@ -267,7 +267,7 @@ def main( ), ( not native_size_within_limit, - f"native extension exceeds 20 MB: {native_member.file_size / 1_000_000:.2f} MB", + f"native extension exceeds 30 MB: {native_member.file_size / 1_000_000:.2f} MB", ), (bool(unexpected_members), f"wheel contains unexpected build artifacts: {', '.join(unexpected_members)}"), ) From 0e281d458fa71ae5b9c12006e5986115776f5f50 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 21:16:19 +0000 Subject: [PATCH 13/15] fix(rust): treat ConditionNotMet as an existing blob on sync Azure Blob writes Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../crates/cache-azure-blob/src/cache.rs | 23 ++++++++----- .../cache-azure-blob/src/cache/tests.rs | 34 ++++++++++++++++++- 2 files changed, 47 insertions(+), 10 deletions(-) diff --git a/litellm-rust/crates/cache-azure-blob/src/cache.rs b/litellm-rust/crates/cache-azure-blob/src/cache.rs index f28b8c9a641..efad0bc7f82 100644 --- a/litellm-rust/crates/cache-azure-blob/src/cache.rs +++ b/litellm-rust/crates/cache-azure-blob/src/cache.rs @@ -19,7 +19,6 @@ use url::Url; use crate::credential::AzureBlobCredential; -/// Synchronous methods block on `runtime` and therefore must run outside of it pub struct AzureBlobCache { container: BlobContainerClient, codec: C, @@ -54,14 +53,15 @@ impl AzureBlobCache { codec: C, runtime: Handle, ) -> Result { - let mut url = Url::parse(account_url).map_err(|_| Error::Unavailable)?; - let account_url = url.as_str().trim_end_matches('/').to_string(); - url.path_segments_mut() - .map_err(|()| Error::Unavailable)? - .pop_if_empty() - .push(container); + let account_url = Url::parse(account_url) + .map_err(|_| Error::Unavailable)? + .as_str() + .trim_end_matches('/') + .to_string(); + let container_url = + Url::parse(&format!("{account_url}/{container}")).map_err(|_| Error::Unavailable)?; let client = BlobContainerClient::new( - url, + container_url, credential, Some(BlobContainerClientOptions { client_options, @@ -108,7 +108,7 @@ impl AzureBlobCache { .await { Ok(_) => Ok(()), - Err(error) if is_storage_error(&error, StorageErrorCode::BlobAlreadyExists) => Ok(()), + Err(error) if !overwrite && is_already_present(&error) => Ok(()), Err(_) => Err(Error::Unavailable), } } @@ -153,6 +153,11 @@ impl AzureBlobCache { } } +fn is_already_present(error: &azure_core::Error) -> bool { + is_storage_error(error, StorageErrorCode::BlobAlreadyExists) + || is_storage_error(error, StorageErrorCode::ConditionNotMet) +} + fn is_storage_error(error: &azure_core::Error, code: StorageErrorCode) -> bool { matches!( error.kind(), diff --git a/litellm-rust/crates/cache-azure-blob/src/cache/tests.rs b/litellm-rust/crates/cache-azure-blob/src/cache/tests.rs index fd116a0e28f..580674450c3 100644 --- a/litellm-rust/crates/cache-azure-blob/src/cache/tests.rs +++ b/litellm-rust/crates/cache-azure-blob/src/cache/tests.rs @@ -39,6 +39,7 @@ struct FakeState { blobs: BTreeMap>, requests: Vec, failing: bool, + precondition_conflicts: bool, } #[derive(Clone, Default)] @@ -79,6 +80,10 @@ impl FakeBlobService { self.state.lock().unwrap().failing = failing; } + fn set_precondition_conflicts(&self, enabled: bool) { + self.state.lock().unwrap().precondition_conflicts = enabled; + } + fn requests(&self) -> Vec { self.state.lock().unwrap().requests.clone() } @@ -158,7 +163,15 @@ impl HttpClient for FakeBlobService { } (Method::Put, false, Some(name)) => { if if_none_match.as_deref() == Some("*") && state.blobs.contains_key(&name) { - Self::respond(StatusCode::Conflict, Some("BlobAlreadyExists"), Vec::new()) + if state.precondition_conflicts { + Self::respond( + StatusCode::PreconditionFailed, + Some("ConditionNotMet"), + Vec::new(), + ) + } else { + Self::respond(StatusCode::Conflict, Some("BlobAlreadyExists"), Vec::new()) + } } else { let bytes = match request.body() { Body::Bytes(bytes) => bytes.to_vec(), @@ -369,6 +382,25 @@ fn sync_set_does_not_overwrite_an_existing_blob() { ); } +#[test] +fn sync_set_treats_a_precondition_conflict_as_an_existing_blob() { + let fixture = Fixture::new(FakeBlobService::default()); + fixture.service.set_precondition_conflicts(true); + fixture + .cache + .set_cache("key", entry(json!({"v": "first"})), &no_ttl()) + .unwrap(); + fixture + .cache + .set_cache("key", entry(json!({"v": "second"})), &no_ttl()) + .unwrap(); + + assert_eq!( + fixture.stored_json("key")["response"], + json!({"v": "first"}) + ); +} + #[test] fn async_set_overwrites_an_existing_blob() { let fixture = Fixture::new(FakeBlobService::default()); From 2f8bee053dbb0074151589fc2d15214ce2f6f9cf Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 21 Sep 2026 14:16:37 -0700 Subject: [PATCH 14/15] test(alerting): inject the webhook client and extend the mapped test files --- .../SlackAlerting/slack_alerting.py | 6 +- .../SlackAlerting/test_batching_handler.py | 98 ------------------- .../SlackAlerting/test_ms_teams.py | 61 +++++++++--- .../SlackAlerting/test_slack_alerting.py | 94 +++++++++++++++++- 4 files changed, 147 insertions(+), 112 deletions(-) delete mode 100644 tests/test_litellm/integrations/SlackAlerting/test_batching_handler.py diff --git a/litellm/integrations/SlackAlerting/slack_alerting.py b/litellm/integrations/SlackAlerting/slack_alerting.py index 7b579546403..8d0d044ff93 100644 --- a/litellm/integrations/SlackAlerting/slack_alerting.py +++ b/litellm/integrations/SlackAlerting/slack_alerting.py @@ -33,6 +33,7 @@ from litellm.litellm_core_utils.exception_mapping_utils import ( _add_key_name_and_team_to_alert, ) from litellm.llms.custom_httpx.http_handler import ( + AsyncHTTPHandler, get_async_httpx_client, httpxSpecialProvider, ) @@ -99,6 +100,7 @@ class SlackAlerting(CustomBatchLogger): alerting_args={}, default_webhook_url: str | None = None, alert_type_config: dict[str, dict] | None = None, + async_http_handler: AsyncHTTPHandler | None = None, **kwargs, ): if alerting_threshold is None: @@ -107,7 +109,9 @@ class SlackAlerting(CustomBatchLogger): self.alerting = alerting self.alert_types = alert_types self.internal_usage_cache = internal_usage_cache or DualCache() - self.async_http_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback) + self.async_http_handler = async_http_handler or get_async_httpx_client( + llm_provider=httpxSpecialProvider.LoggingCallback + ) self.alert_to_webhook_url = process_slack_alerting_variables(alert_to_webhook_url=alert_to_webhook_url) self.is_running = False self.alerting_args = SlackAlertingArgs(**alerting_args) diff --git a/tests/test_litellm/integrations/SlackAlerting/test_batching_handler.py b/tests/test_litellm/integrations/SlackAlerting/test_batching_handler.py deleted file mode 100644 index 9052cb8bb5d..00000000000 --- a/tests/test_litellm/integrations/SlackAlerting/test_batching_handler.py +++ /dev/null @@ -1,98 +0,0 @@ -import json -from typing import Final -from unittest.mock import AsyncMock, MagicMock - -import pytest - -from litellm.integrations.SlackAlerting.ms_teams import MS_TEAMS_WEBHOOK_URL_ENV -from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting -from litellm.types.integrations.slack_alerting import AlertQueueItem, AlertType - -SLACK_WEBHOOK_URL: Final = "https://hooks.slack.com/services/test" -THRESHOLD_ALERT: Final = "User Budget: 15% or less of budget remaining\n\n*user_id:* `user-a`" -CROSSED_ALERT: Final = "User Budget: Budget Crossed\n\n*user_id:* `user-b`" - - -def _slack_alerting_recording_posts(alerting: list[str]) -> SlackAlerting: - slack_alerting: Final = SlackAlerting(alerting=alerting) - slack_alerting.periodic_started = True - response: Final = MagicMock() - response.status_code = 200 - slack_alerting.async_http_handler = MagicMock() - slack_alerting.async_http_handler.post = AsyncMock(return_value=response) - return slack_alerting - - -def _queued_slack_alert(text: str) -> AlertQueueItem: - return { - "url": SLACK_WEBHOOK_URL, - "headers": {"Content-type": "application/json"}, - "payload": {"text": text}, - "alert_type": AlertType.budget_alerts, - } - - -def _posted_bodies(slack_alerting: SlackAlerting) -> tuple[dict, ...]: - return tuple(json.loads(call.kwargs["data"]) for call in slack_alerting.async_http_handler.post.call_args_list) - - -async def _send_budget_alert(slack_alerting: SlackAlerting, message: str) -> None: - await slack_alerting.send_alert( - message=message, - level="High", - alert_type=AlertType.budget_alerts, - alerting_metadata={}, - ) - - -@pytest.mark.asyncio -async def test_async_send_batch_delivers_every_distinct_alert_queued_in_one_flush(monkeypatch): - monkeypatch.setenv("SLACK_WEBHOOK_URL", SLACK_WEBHOOK_URL) - slack_alerting: Final = _slack_alerting_recording_posts(["slack"]) - await _send_budget_alert(slack_alerting, THRESHOLD_ALERT) - await _send_budget_alert(slack_alerting, CROSSED_ALERT) - - await slack_alerting.async_send_batch() - - posted_texts: Final = tuple(body["text"] for body in _posted_bodies(slack_alerting)) - assert len(posted_texts) == 2 - assert THRESHOLD_ALERT in posted_texts[0] - assert CROSSED_ALERT in posted_texts[1] - assert not any(text.startswith("[Num Alerts") for text in posted_texts) - assert slack_alerting.log_queue == [] - - -@pytest.mark.asyncio -async def test_async_send_batch_collapses_only_identical_alerts(): - slack_alerting: Final = _slack_alerting_recording_posts(["slack"]) - slack_alerting.log_queue.extend( - ( - _queued_slack_alert(THRESHOLD_ALERT), - _queued_slack_alert(CROSSED_ALERT), - _queued_slack_alert(THRESHOLD_ALERT), - ) - ) - - await slack_alerting.async_send_batch() - - assert _posted_bodies(slack_alerting) == ( - {"text": f"[Num Alerts: 2]\n\n{THRESHOLD_ALERT}"}, - {"text": CROSSED_ALERT}, - ) - - -@pytest.mark.asyncio -async def test_async_send_batch_delivers_every_distinct_ms_teams_alert(monkeypatch): - monkeypatch.setenv(MS_TEAMS_WEBHOOK_URL_ENV, "https://teams.example/webhook") - slack_alerting: Final = _slack_alerting_recording_posts(["ms_teams"]) - await _send_budget_alert(slack_alerting, THRESHOLD_ALERT) - await _send_budget_alert(slack_alerting, CROSSED_ALERT) - - await slack_alerting.async_send_batch() - - card_texts: Final = tuple( - body["attachments"][0]["content"]["body"][0]["text"] for body in _posted_bodies(slack_alerting) - ) - assert len(card_texts) == 2 - assert THRESHOLD_ALERT in card_texts[0] - assert CROSSED_ALERT in card_texts[1] diff --git a/tests/test_litellm/integrations/SlackAlerting/test_ms_teams.py b/tests/test_litellm/integrations/SlackAlerting/test_ms_teams.py index 41b7f3b969b..44426c00628 100644 --- a/tests/test_litellm/integrations/SlackAlerting/test_ms_teams.py +++ b/tests/test_litellm/integrations/SlackAlerting/test_ms_teams.py @@ -2,18 +2,39 @@ import json from typing import Final from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest +from pydantic import TypeAdapter from litellm.integrations.SlackAlerting.batching_handler import send_to_webhook from litellm.integrations.SlackAlerting.ms_teams import ( MS_TEAMS_ALERTING_DESTINATION, MS_TEAMS_WEBHOOK_URL_ENV, + MSTeamsMessage, build_ms_teams_payload, get_ms_teams_webhook_url, ) from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.proxy._types import AlertType +_MS_TEAMS_MESSAGE: Final = TypeAdapter(MSTeamsMessage) + + +def _webhook_accepting_posts() -> AsyncMock: + response: Final = MagicMock(spec=httpx.Response) + response.status_code = 200 + http_handler: Final = AsyncMock(spec=AsyncHTTPHandler) + http_handler.post.return_value = response + return http_handler + + +def _posted_card_texts(http_handler: AsyncMock) -> tuple[str, ...]: + return tuple( + _MS_TEAMS_MESSAGE.validate_json(call.kwargs["data"])["attachments"][0]["content"]["body"][0]["text"] + for call in http_handler.post.call_args_list + ) + def test_build_ms_teams_payload_wraps_text_in_adaptive_card(): payload: Final = build_ms_teams_payload("hello alert") @@ -80,11 +101,8 @@ async def test_send_alert_slack_and_ms_teams_enqueue_both(monkeypatch): @pytest.mark.asyncio async def test_send_to_webhook_posts_adaptive_card_for_ms_teams_items(): - slack_alerting: Final = SlackAlerting(alerting=["ms_teams"]) - mock_response: Final = MagicMock() - mock_response.status_code = 200 - slack_alerting.async_http_handler = MagicMock() - slack_alerting.async_http_handler.post = AsyncMock(return_value=mock_response) + http_handler: Final = _webhook_accepting_posts() + slack_alerting: Final = SlackAlerting(alerting=["ms_teams"], async_http_handler=http_handler) item: Final = { "url": "https://teams.example/webhook", @@ -95,7 +113,7 @@ async def test_send_to_webhook_posts_adaptive_card_for_ms_teams_items(): } await send_to_webhook(slackAlertingInstance=slack_alerting, item=item, count=1) - call_kwargs: Final = slack_alerting.async_http_handler.post.call_args.kwargs + call_kwargs: Final = http_handler.post.call_args.kwargs assert call_kwargs["url"] == "https://teams.example/webhook" sent_body: Final = json.loads(call_kwargs["data"]) assert sent_body["type"] == "message" @@ -104,11 +122,8 @@ async def test_send_to_webhook_posts_adaptive_card_for_ms_teams_items(): @pytest.mark.asyncio async def test_send_to_webhook_keeps_slack_payload_shape(): - slack_alerting: Final = SlackAlerting(alerting=["slack"]) - mock_response: Final = MagicMock() - mock_response.status_code = 200 - slack_alerting.async_http_handler = MagicMock() - slack_alerting.async_http_handler.post = AsyncMock(return_value=mock_response) + http_handler: Final = _webhook_accepting_posts() + slack_alerting: Final = SlackAlerting(alerting=["slack"], async_http_handler=http_handler) item: Final = { "url": "https://hooks.slack.com/services/test", @@ -118,5 +133,27 @@ async def test_send_to_webhook_keeps_slack_payload_shape(): } await send_to_webhook(slackAlertingInstance=slack_alerting, item=item, count=1) - call_kwargs: Final = slack_alerting.async_http_handler.post.call_args.kwargs + call_kwargs: Final = http_handler.post.call_args.kwargs assert json.loads(call_kwargs["data"]) == {"text": "alert body"} + + +@pytest.mark.asyncio +async def test_async_send_batch_delivers_every_distinct_ms_teams_alert(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv(MS_TEAMS_WEBHOOK_URL_ENV, "https://teams.example/webhook") + http_handler: Final = _webhook_accepting_posts() + slack_alerting: Final = SlackAlerting(alerting=["ms_teams"], async_http_handler=http_handler) + slack_alerting.periodic_started = True + for message in ("User Budget: 15% or less of budget remaining", "User Budget: Budget Crossed"): + await slack_alerting.send_alert( + message=message, + level="High", + alert_type=AlertType.budget_alerts, + alerting_metadata={}, + ) + + await slack_alerting.async_send_batch() + + card_texts: Final = _posted_card_texts(http_handler) + assert len(card_texts) == 2 + assert "User Budget: 15% or less of budget remaining" in card_texts[0] + assert "User Budget: Budget Crossed" in card_texts[1] diff --git a/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py b/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py index 4bc6c08bd63..2d5eb78950c 100644 --- a/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py +++ b/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py @@ -6,14 +6,18 @@ import unittest from typing import Final, List, Optional, Tuple from unittest.mock import ANY, AsyncMock, MagicMock, Mock, patch +import httpx import pytest +from pydantic import TypeAdapter +from typing_extensions import ReadOnly, TypedDict import litellm from litellm.caching.caching import DualCache from litellm.integrations.SlackAlerting.budget_alert_types import get_budget_alert_type from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.proxy._types import CallInfo, Litellm_EntityType -from litellm.types.integrations.slack_alerting import AlertType, SlackAlertingCacheKeys +from litellm.types.integrations.slack_alerting import AlertQueueItem, AlertType, SlackAlertingCacheKeys class TestSlackAlerting(unittest.TestCase): @@ -434,3 +438,91 @@ async def test_send_alert_raises_when_no_webhook_url_configured(monkeypatch): alert_type=AlertType.budget_alerts, alerting_metadata={}, ) + + +SLACK_WEBHOOK_URL: Final = "https://hooks.slack.com/services/test" +THRESHOLD_ALERT: Final = "User Budget: 15% or less of budget remaining\n\n*user_id:* `user-a`" +CROSSED_ALERT: Final = "User Budget: Budget Crossed\n\n*user_id:* `user-b`" + + +class _SlackWebhookBody(TypedDict): + text: ReadOnly[str] + + +_SLACK_WEBHOOK_BODY: Final = TypeAdapter(_SlackWebhookBody) + + +def _webhook_accepting_posts() -> AsyncMock: + response: Final = MagicMock(spec=httpx.Response) + response.status_code = 200 + http_handler: Final = AsyncMock(spec=AsyncHTTPHandler) + http_handler.post.return_value = response + return http_handler + + +def _slack_alerting_flushing_to(http_handler: AsyncHTTPHandler) -> SlackAlerting: + slack_alerting: Final = SlackAlerting(alerting=["slack"], async_http_handler=http_handler) + slack_alerting.periodic_started = True + return slack_alerting + + +def _queued_slack_alert(text: str) -> AlertQueueItem: + return { + "url": SLACK_WEBHOOK_URL, + "headers": {"Content-type": "application/json"}, + "payload": {"text": text}, + "alert_type": AlertType.budget_alerts, + } + + +def _posted_slack_bodies(http_handler: AsyncMock) -> tuple[_SlackWebhookBody, ...]: + return tuple(_SLACK_WEBHOOK_BODY.validate_json(call.kwargs["data"]) for call in http_handler.post.call_args_list) + + +async def _send_budget_alert(slack_alerting: SlackAlerting, message: str) -> None: + await slack_alerting.send_alert( + message=message, + level="High", + alert_type=AlertType.budget_alerts, + alerting_metadata={}, + ) + + +@pytest.mark.asyncio +async def test_async_send_batch_delivers_every_distinct_alert_queued_in_one_flush( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("SLACK_WEBHOOK_URL", SLACK_WEBHOOK_URL) + http_handler: Final = _webhook_accepting_posts() + slack_alerting: Final = _slack_alerting_flushing_to(http_handler) + await _send_budget_alert(slack_alerting, THRESHOLD_ALERT) + await _send_budget_alert(slack_alerting, CROSSED_ALERT) + + await slack_alerting.async_send_batch() + + posted_texts: Final = tuple(body["text"] for body in _posted_slack_bodies(http_handler)) + assert len(posted_texts) == 2 + assert THRESHOLD_ALERT in posted_texts[0] + assert CROSSED_ALERT in posted_texts[1] + assert not any(text.startswith("[Num Alerts") for text in posted_texts) + assert slack_alerting.log_queue == [] + + +@pytest.mark.asyncio +async def test_async_send_batch_collapses_only_identical_alerts() -> None: + http_handler: Final = _webhook_accepting_posts() + slack_alerting: Final = _slack_alerting_flushing_to(http_handler) + slack_alerting.log_queue.extend( + ( + _queued_slack_alert(THRESHOLD_ALERT), + _queued_slack_alert(CROSSED_ALERT), + _queued_slack_alert(THRESHOLD_ALERT), + ) + ) + + await slack_alerting.async_send_batch() + + assert _posted_slack_bodies(http_handler) == ( + {"text": f"[Num Alerts: 2]\n\n{THRESHOLD_ALERT}"}, + {"text": CROSSED_ALERT}, + ) From 406365816802151bb9b9b7d99d6bd28a7ff79d22 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 21:22:02 +0000 Subject: [PATCH 15/15] build(rust): drop native extension size limit bump, deferred to #42300 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/scripts/verify_linux_native_wheel.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/scripts/verify_linux_native_wheel.py b/.github/scripts/verify_linux_native_wheel.py index 0adbc015ad0..4fb8f068eb0 100644 --- a/.github/scripts/verify_linux_native_wheel.py +++ b/.github/scripts/verify_linux_native_wheel.py @@ -205,7 +205,7 @@ def main( native_module: Final = load_native_module(native_path) native_module_loads: Final = native_module is not None panic_test_hook_absent: Final = native_module is not None and not hasattr(native_module, "_panic_for_test") - native_size_limit: Final = 30_000_000 + native_size_limit: Final = 25_000_000 native_size_within_limit: Final = native_member.file_size <= native_size_limit validations: Final = ( (f"Python tag is {EXPECTED_PYTHON_TAG}", python_tag == EXPECTED_PYTHON_TAG), @@ -222,7 +222,7 @@ def main( ("Python extension entry point is present", extension_entry_point_present), ("Native module loads", native_module_loads), ("Production module omits the panic test hook", panic_test_hook_absent), - ("Native extension does not exceed 30 MB", native_size_within_limit), + ("Native extension does not exceed 25 MB", native_size_within_limit), ("Wheel contents are valid", not unexpected_members), ) @@ -267,7 +267,7 @@ def main( ), ( not native_size_within_limit, - f"native extension exceeds 30 MB: {native_member.file_size / 1_000_000:.2f} MB", + f"native extension exceeds 20 MB: {native_member.file_size / 1_000_000:.2f} MB", ), (bool(unexpected_members), f"wheel contains unexpected build artifacts: {', '.join(unexpected_members)}"), )