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/24] 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/24] 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/24] 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/24] 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 4e2d4b5ff9f163b6fb121487db1e2079e0fd8d44 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:23:54 +0000 Subject: [PATCH 05/24] feat(rust): add CyberArk Conjur 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 + .../crates/secrets-cyberark/Cargo.toml | 25 + .../crates/secrets-cyberark/src/error.rs | 27 + .../crates/secrets-cyberark/src/lib.rs | 7 + .../secrets-cyberark/src/secret_manager.rs | 315 +++++++++++ .../tests/fixtures/parity.json | 32 ++ .../secrets-cyberark/tests/secret_manager.rs | 488 ++++++++++++++++++ litellm-rust/crates/secrets/Cargo.toml | 2 + litellm-rust/crates/secrets/src/error.rs | 3 + litellm-rust/crates/secrets/src/handler.rs | 10 + litellm-rust/crates/secrets/src/lib.rs | 2 + litellm-rust/crates/secrets/tests/handler.rs | 53 ++ .../test_cyberark_secret_manager.py | 119 +++++ 15 files changed, 1106 insertions(+), 1 deletion(-) create mode 100644 litellm-rust/crates/secrets-cyberark/Cargo.toml create mode 100644 litellm-rust/crates/secrets-cyberark/src/error.rs create mode 100644 litellm-rust/crates/secrets-cyberark/src/lib.rs create mode 100644 litellm-rust/crates/secrets-cyberark/src/secret_manager.rs create mode 100644 litellm-rust/crates/secrets-cyberark/tests/fixtures/parity.json create mode 100644 litellm-rust/crates/secrets-cyberark/tests/secret_manager.rs create mode 100644 tests/test_litellm/secret_managers/test_cyberark_secret_manager.py diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml index 278fa7c425f..161353708f0 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 cyberark aws,google aws,google,cyberark; 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..07c9eb8c437 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-cyberark", "litellm-secrets-google", "litellm-secrets-types", "moka", @@ -2731,6 +2732,26 @@ dependencies = [ "wiremock", ] +[[package]] +name = "litellm-secrets-cyberark" +version = "0.1.0" +dependencies = [ + "base64 0.22.1", + "litellm-core-utils", + "litellm-secrets-types", + "moka", + "percent-encoding", + "reqwest 0.12.28", + "rstest", + "serde", + "serde_json", + "thiserror 2.0.19", + "tokio", + "tracing", + "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..d35df1eafb8 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-cyberark = { path = "crates/secrets-cyberark" } litellm-http = { path = "crates/http" } litellm-llms = { path = "crates/llms" } litellm-types = { path = "crates/types" } diff --git a/litellm-rust/crates/secrets-cyberark/Cargo.toml b/litellm-rust/crates/secrets-cyberark/Cargo.toml new file mode 100644 index 00000000000..de8ae5b71dc --- /dev/null +++ b/litellm-rust/crates/secrets-cyberark/Cargo.toml @@ -0,0 +1,25 @@ +[package] +name = "litellm-secrets-cyberark" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +litellm-secrets-types.workspace = true +litellm-core-utils.workspace = true +base64.workspace = true +moka.workspace = true +reqwest.workspace = true +serde_json.workspace = true +thiserror.workspace = true +veil.workspace = true +tracing = "0.1" +percent-encoding = "2.3" + +[dev-dependencies] +rstest.workspace = true +tokio.workspace = true +wiremock = "0.6.5" +serde.workspace = true +serde_json.workspace = true diff --git a/litellm-rust/crates/secrets-cyberark/src/error.rs b/litellm-rust/crates/secrets-cyberark/src/error.rs new file mode 100644 index 00000000000..5a14f4f3db8 --- /dev/null +++ b/litellm-rust/crates/secrets-cyberark/src/error.rs @@ -0,0 +1,27 @@ +#[derive(thiserror::Error, veil::Redact)] +pub enum Error { + #[error("CyberArk Conjur HTTP request failed")] + Http( + #[from] + #[redact] + reqwest::Error, + ), + #[error("CyberArk Conjur authentication returned HTTP {0}")] + AuthStatus(u16), + #[error("CyberArk Conjur returned HTTP {0}")] + Status(u16), + #[error( + "CyberArk credentials are missing: set CYBERARK_API_KEY or both CYBERARK_CLIENT_CERT and CYBERARK_CLIENT_KEY" + )] + MissingCredentials, + #[error("CyberArk client certificate could not be loaded")] + ClientCertificate, + #[error("invalid refresh interval")] + RefreshInterval, + #[error("invalid CyberArk Conjur endpoint")] + Endpoint, + #[error("CyberArk secret manager requires an enterprise license")] + EnterpriseRequired, + #[error(transparent)] + Operation(#[from] litellm_secrets_types::Error), +} diff --git a/litellm-rust/crates/secrets-cyberark/src/lib.rs b/litellm-rust/crates/secrets-cyberark/src/lib.rs new file mode 100644 index 00000000000..5288f8116b1 --- /dev/null +++ b/litellm-rust/crates/secrets-cyberark/src/lib.rs @@ -0,0 +1,7 @@ +#![forbid(unsafe_code)] + +mod error; +mod secret_manager; + +pub use error::Error; +pub use secret_manager::{CyberArkSecretManager, DeleteOutcome}; diff --git a/litellm-rust/crates/secrets-cyberark/src/secret_manager.rs b/litellm-rust/crates/secrets-cyberark/src/secret_manager.rs new file mode 100644 index 00000000000..ade257a5872 --- /dev/null +++ b/litellm-rust/crates/secrets-cyberark/src/secret_manager.rs @@ -0,0 +1,315 @@ +use std::{fs, sync::Arc, time::Duration}; + +use base64::{Engine, engine::general_purpose::STANDARD}; +use litellm_core_utils::settings::Lookup; +use litellm_secrets_types::{BaseSecretManager, SecretValue, validate_secret_name}; +use moka::future::Cache; +use percent_encoding::{AsciiSet, NON_ALPHANUMERIC, utf8_percent_encode}; + +use crate::Error; + +const CYBERARK_API_BASE: &str = "CYBERARK_API_BASE"; +const CYBERARK_ACCOUNT: &str = "CYBERARK_ACCOUNT"; +const CYBERARK_USERNAME: &str = "CYBERARK_USERNAME"; +const CYBERARK_API_KEY: &str = "CYBERARK_API_KEY"; +const CYBERARK_CLIENT_CERT: &str = "CYBERARK_CLIENT_CERT"; +const CYBERARK_CLIENT_KEY: &str = "CYBERARK_CLIENT_KEY"; +const CYBERARK_SSL_VERIFY: &str = "CYBERARK_SSL_VERIFY"; +const CYBERARK_REFRESH_INTERVAL: &str = "CYBERARK_REFRESH_INTERVAL"; +const DEFAULT_API_BASE: &str = "http://127.0.0.1:8080"; +const DEFAULT_ACCOUNT: &str = "default"; +const DEFAULT_USERNAME: &str = "admin"; +const DEFAULT_REFRESH_INTERVAL: Duration = Duration::from_secs(300); +const SECRET_NAME_SAFE: &AsciiSet = &NON_ALPHANUMERIC + .remove(b'-') + .remove(b'_') + .remove(b'.') + .remove(b'~'); + +#[derive(Clone)] +pub struct CyberArkSecretManager { + client: reqwest::Client, + endpoint: reqwest::Url, + account: String, + username: String, + api_key: SecretValue, + token: Cache<(), SecretValue>, + secrets: Cache, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum DeleteOutcome { + NotSupported, +} + +impl CyberArkSecretManager { + pub fn with_client( + client: reqwest::Client, + endpoint: reqwest::Url, + account: String, + username: String, + api_key: SecretValue, + refresh_interval: Option, + ) -> Self { + let endpoint = normalize_endpoint(endpoint); + let ttl = refresh_interval + .filter(|interval| !interval.is_zero()) + .unwrap_or(DEFAULT_REFRESH_INTERVAL); + let token = Cache::builder().time_to_live(ttl).build(); + let secrets = Cache::builder().time_to_live(ttl).build(); + Self { + client, + endpoint, + account, + username, + api_key, + token, + secrets, + } + } + + pub fn new( + environment: Arc, + enterprise_enabled: bool, + ) -> Result { + let api_key = environment.get(CYBERARK_API_KEY).unwrap_or_default(); + let cert = environment.get(CYBERARK_CLIENT_CERT).unwrap_or_default(); + let key = environment.get(CYBERARK_CLIENT_KEY).unwrap_or_default(); + if api_key.is_empty() && (cert.is_empty() || key.is_empty()) { + return Err(Error::MissingCredentials); + } + if !enterprise_enabled { + return Err(Error::EnterpriseRequired); + } + let verify = environment + .get(CYBERARK_SSL_VERIFY) + .map(|value| !value.trim().eq_ignore_ascii_case("false")) + .unwrap_or(true); + let mut builder = reqwest::Client::builder(); + if !verify { + tracing::warn!( + "CyberArk SSL verification is disabled. This is insecure and should only be used for testing with self-signed certificates." + ); + builder = builder.danger_accept_invalid_certs(true); + } + if !cert.is_empty() && !key.is_empty() { + let certificate = fs::read(cert).map_err(|_| Error::ClientCertificate)?; + let private_key = fs::read(key).map_err(|_| Error::ClientCertificate)?; + let identity = reqwest::Identity::from_pem(&[certificate, private_key].concat()) + .map_err(|_| Error::ClientCertificate)?; + builder = builder.identity(identity); + } + let client = builder.build()?; + let endpoint = reqwest::Url::parse( + &environment + .get(CYBERARK_API_BASE) + .unwrap_or_else(|| DEFAULT_API_BASE.to_owned()), + ) + .map_err(|_| Error::Endpoint)?; + let account = environment + .get(CYBERARK_ACCOUNT) + .unwrap_or_else(|| DEFAULT_ACCOUNT.to_owned()); + let username = environment + .get(CYBERARK_USERNAME) + .unwrap_or_else(|| DEFAULT_USERNAME.to_owned()); + let refresh_interval = environment + .get(CYBERARK_REFRESH_INTERVAL) + .map(|value| { + value + .parse::() + .map_err(|_| Error::RefreshInterval) + .map(|seconds| match seconds { + seconds if seconds < 0 => Duration::from_nanos(1), + 0 => DEFAULT_REFRESH_INTERVAL, + seconds => Duration::from_secs(seconds as u64), + }) + }) + .transpose()?; + Ok(Self::with_client( + client, + endpoint, + account, + username, + SecretValue::new(api_key), + refresh_interval, + )) + } + + fn secret_url(&self, name: &str) -> Result { + let encoded = utf8_percent_encode(name, SECRET_NAME_SAFE); + self.endpoint + .join(&format!("secrets/{}/variable/{}", self.account, encoded)) + .map_err(|_| Error::Endpoint) + } + + async fn authenticate(&self) -> Result { + if let Some(token) = self.token.get(&()).await { + return Ok(token); + } + let url = self + .endpoint + .join(&format!( + "authn/{}/{}/authenticate", + self.account, self.username + )) + .map_err(|_| Error::Endpoint)?; + let response = self + .client + .post(url) + .body(self.api_key.expose().to_owned()) + .send() + .await?; + if !response.status().is_success() { + return Err(Error::AuthStatus(response.status().as_u16())); + } + let token = SecretValue::new(STANDARD.encode(response.text().await?)); + self.token.insert((), token.clone()).await; + Ok(token) + } + + async fn authorization_header(&self) -> Result { + Ok(format!( + "Token token=\"{}\"", + self.authenticate().await?.expose() + )) + } + + pub async fn async_read_secret(&self, name: &str) -> Result, Error> { + if let Some(value) = self.secrets.get(name).await { + return Ok(Some(value)); + } + let response = self + .client + .get(self.secret_url(name)?) + .header("Authorization", self.authorization_header().await?) + .send() + .await?; + if response.status() == reqwest::StatusCode::NOT_FOUND { + return Ok(None); + } + if !response.status().is_success() { + return Err(Error::Status(response.status().as_u16())); + } + let value = SecretValue::new(response.text().await?); + self.secrets.insert(name.to_owned(), value.clone()).await; + Ok(Some(value)) + } + + pub async fn async_write_secret( + &self, + name: &str, + value: &SecretValue, + _description: Option<&str>, + ) -> Result<(), Error> { + validate_secret_name(name)?; + self.ensure_variable_exists(name).await; + let response = self + .client + .post(self.secret_url(name)?) + .header("Authorization", self.authorization_header().await?) + .body(value.expose().to_owned()) + .send() + .await?; + if !response.status().is_success() { + return Err(Error::Status(response.status().as_u16())); + } + self.secrets.insert(name.to_owned(), value.clone()).await; + Ok(()) + } + + async fn ensure_variable_exists(&self, name: &str) { + let policy_url = self + .endpoint + .join(&format!("policies/{}/policy/root", self.account)); + let Ok(policy_url) = policy_url else { + tracing::warn!("Could not build CyberArk policy endpoint"); + return; + }; + let Ok(authorization) = self.authorization_header().await else { + tracing::warn!("Could not authenticate while ensuring CyberArk variable exists"); + return; + }; + let body = format!( + "- !variable {}\n", + serde_json::to_string(name).expect("serializing a string cannot fail") + ); + let response = self + .client + .post(policy_url) + .header("Authorization", authorization) + .header("Content-Type", "application/x-yaml") + .body(body) + .send() + .await; + match response { + Ok(response) if response.status().is_success() => {} + Ok(response) + if matches!( + response.status(), + reqwest::StatusCode::CONFLICT | reqwest::StatusCode::UNPROCESSABLE_ENTITY + ) => + { + tracing::debug!( + "CyberArk variable policy already exists or conflicts: {}", + response.status() + ); + } + Ok(response) => { + tracing::warn!( + "Could not ensure CyberArk variable exists: {}", + response.status() + ); + } + Err(error) => { + tracing::warn!("Error ensuring CyberArk variable exists: {error}"); + } + } + } + + pub async fn async_delete_secret( + &self, + name: &str, + _recovery_window_in_days: i64, + ) -> Result { + tracing::warn!( + "CyberArk Conjur does not support direct secret deletion. Secrets must be removed through policy updates." + ); + self.secrets.invalidate(name).await; + Ok(DeleteOutcome::NotSupported) + } +} + +impl BaseSecretManager for CyberArkSecretManager { + type Error = Error; + type WriteResponse = (); + type DeleteResponse = DeleteOutcome; + + async fn async_read_secret(&self, name: &str) -> Result, Error> { + self.async_read_secret(name).await + } + + async fn async_write_secret( + &self, + name: &str, + value: &SecretValue, + description: Option<&str>, + ) -> Result<(), Error> { + self.async_write_secret(name, value, description).await + } + + async fn async_delete_secret( + &self, + name: &str, + recovery_window_in_days: i64, + ) -> Result { + self.async_delete_secret(name, recovery_window_in_days) + .await + } +} + +fn normalize_endpoint(mut endpoint: reqwest::Url) -> reqwest::Url { + if !endpoint.path().ends_with('/') { + endpoint.set_path(&format!("{}/", endpoint.path())); + } + endpoint +} diff --git a/litellm-rust/crates/secrets-cyberark/tests/fixtures/parity.json b/litellm-rust/crates/secrets-cyberark/tests/fixtures/parity.json new file mode 100644 index 00000000000..b7aab572985 --- /dev/null +++ b/litellm-rust/crates/secrets-cyberark/tests/fixtures/parity.json @@ -0,0 +1,32 @@ +{ + "endpoint": "http://conjur.test:8080", + "account": "acct", + "username": "admin", + "api_key": "k3y", + "authenticate_path": "/authn/acct/admin/authenticate", + "token_json": "{\"protected\":\"p\",\"payload\":\"q\",\"signature\":\"s\"}", + "authorization_header": "Token token=\"eyJwcm90ZWN0ZWQiOiJwIiwicGF5bG9hZCI6InEiLCJzaWduYXR1cmUiOiJzIn0=\"", + "policy_path": "/policies/acct/policy/root", + "secrets": [ + { + "name": "OPENAI_API_KEY", + "path": "/secrets/acct/variable/OPENAI_API_KEY", + "policy_body": "- !variable \"OPENAI_API_KEY\"\n" + }, + { + "name": "team/app/key", + "path": "/secrets/acct/variable/team%2Fapp%2Fkey", + "policy_body": "- !variable \"team/app/key\"\n" + }, + { + "name": "a b+c.d-e_f~g", + "path": "/secrets/acct/variable/a%20b%2Bc.d-e_f~g", + "policy_body": "- !variable \"a b+c.d-e_f~g\"\n" + }, + { + "name": "needs \"quote\"", + "path": "/secrets/acct/variable/needs%20%22quote%22", + "policy_body": "- !variable \"needs \\\"quote\\\"\"\n" + } + ] +} diff --git a/litellm-rust/crates/secrets-cyberark/tests/secret_manager.rs b/litellm-rust/crates/secrets-cyberark/tests/secret_manager.rs new file mode 100644 index 00000000000..e8158c46945 --- /dev/null +++ b/litellm-rust/crates/secrets-cyberark/tests/secret_manager.rs @@ -0,0 +1,488 @@ +use std::{sync::Arc, time::Duration}; + +use base64::{Engine, engine::general_purpose::STANDARD}; +use litellm_secrets_cyberark::{CyberArkSecretManager, DeleteOutcome, Error}; +use litellm_secrets_types::{SecretValue, validate_secret_name}; +use serde::Deserialize; +use wiremock::{ + Match, Mock, MockServer, Request, ResponseTemplate, + matchers::{body_string, header, method, path}, +}; + +const TOKEN_JSON: &str = r#"{"protected":"p","payload":"q","signature":"s"}"#; + +#[derive(Deserialize)] +struct ParityFixture { + endpoint: String, + account: String, + username: String, + api_key: String, + authenticate_path: String, + token_json: String, + authorization_header: String, + policy_path: String, + secrets: Vec, +} + +#[derive(Deserialize)] +struct ParitySecret { + name: String, + path: String, + policy_body: String, +} + +#[derive(Debug)] +struct RawPath(String); + +impl Match for RawPath { + fn matches(&self, request: &Request) -> bool { + request.url.path() == self.0 + } +} + +fn fixture() -> ParityFixture { + serde_json::from_str(include_str!("fixtures/parity.json")).unwrap() +} + +fn manager(server: &MockServer, ttl: Duration) -> CyberArkSecretManager { + CyberArkSecretManager::with_client( + reqwest::Client::new(), + server.uri().parse().unwrap(), + "acct".into(), + "admin".into(), + SecretValue::new("k3y"), + Some(ttl), + ) +} + +async fn mount_auth(server: &MockServer, expected: u64) { + Mock::given(method("POST")) + .and(path("/authn/acct/admin/authenticate")) + .and(body_string("k3y")) + .respond_with(ResponseTemplate::new(200).set_body_string(TOKEN_JSON)) + .expect(expected) + .mount(server) + .await; +} + +#[tokio::test] +async fn successful_reads_cache_auth_secret_and_redact_values() { + let server = MockServer::start().await; + mount_auth(&server, 1).await; + let token = STANDARD.encode(TOKEN_JSON); + Mock::given(path("/secrets/acct/variable/OPENAI_API_KEY")) + .and(header("authorization", format!("Token token=\"{token}\""))) + .respond_with(ResponseTemplate::new(200).set_body_string("sk-live")) + .expect(1) + .mount(&server) + .await; + let manager = manager(&server, Duration::from_secs(60)); + + for _ in 0..2 { + let value = manager + .async_read_secret("OPENAI_API_KEY") + .await + .unwrap() + .unwrap(); + assert_eq!(value.expose(), "sk-live"); + assert!(!format!("{value:?}").contains("sk-live")); + } +} + +#[rstest::rstest] +#[case::not_found(404)] +#[case::unauthorized(401)] +#[case::forbidden(403)] +#[case::server_error(500)] +#[tokio::test] +async fn failed_reads_are_not_cached(#[case] status: u16) { + let server = MockServer::start().await; + mount_auth(&server, 1).await; + let failing = Mock::given(path("/secrets/acct/variable/key")) + .respond_with(ResponseTemplate::new(status)) + .expect(1) + .mount_as_scoped(&server) + .await; + let manager = manager(&server, Duration::from_secs(60)); + let result = manager.async_read_secret("key").await; + if status == 404 { + assert_eq!(result.unwrap(), None); + } else { + assert!(matches!(result, Err(Error::Status(actual)) if actual == status)); + } + drop(failing); + Mock::given(path("/secrets/acct/variable/key")) + .respond_with(ResponseTemplate::new(200).set_body_string("recovered")) + .expect(1) + .mount(&server) + .await; + for _ in 0..2 { + assert_eq!( + manager + .async_read_secret("key") + .await + .unwrap() + .unwrap() + .expose(), + "recovered" + ); + } +} + +#[tokio::test] +async fn failed_authentication_is_not_cached_and_does_not_read_secret() { + let server = MockServer::start().await; + let failing = Mock::given(path("/authn/acct/admin/authenticate")) + .respond_with(ResponseTemplate::new(401)) + .expect(1) + .mount_as_scoped(&server) + .await; + let unused_secret = Mock::given(path("/secrets/acct/variable/key")) + .respond_with(ResponseTemplate::new(200).set_body_string("value")) + .expect(0) + .mount_as_scoped(&server) + .await; + let manager = manager(&server, Duration::from_secs(60)); + assert!(matches!( + manager.async_read_secret("key").await, + Err(Error::AuthStatus(401)) + )); + drop(unused_secret); + drop(failing); + mount_auth(&server, 1).await; + Mock::given(path("/secrets/acct/variable/key")) + .respond_with(ResponseTemplate::new(200).set_body_string("value")) + .expect(1) + .mount(&server) + .await; + assert_eq!( + manager + .async_read_secret("key") + .await + .unwrap() + .unwrap() + .expose(), + "value" + ); +} + +#[tokio::test] +async fn expired_tokens_and_secrets_are_fetched_again() { + let server = MockServer::start().await; + mount_auth(&server, 2).await; + Mock::given(path("/secrets/acct/variable/key")) + .respond_with(ResponseTemplate::new(200).set_body_string("value")) + .expect(2) + .mount(&server) + .await; + let manager = manager(&server, Duration::from_millis(1)); + for _ in 0..2 { + assert!(manager.async_read_secret("key").await.unwrap().is_some()); + tokio::time::sleep(Duration::from_millis(5)).await; + } +} + +#[rstest::rstest] +#[tokio::test] +async fn secret_names_use_python_quote_encoding( + #[values("OPENAI_API_KEY", "team/app/key", "a b+c.d-e_f~g", "needs \"quote\"")] name: &str, +) { + let fixture = fixture(); + let secret = fixture + .secrets + .iter() + .find(|secret| secret.name == name) + .unwrap(); + let server = MockServer::start().await; + mount_auth(&server, 1).await; + Mock::given(RawPath(secret.path.clone())) + .respond_with(ResponseTemplate::new(200).set_body_string("value")) + .expect(1) + .mount(&server) + .await; + assert_eq!( + manager(&server, Duration::from_secs(60)) + .async_read_secret(name) + .await + .unwrap() + .unwrap() + .expose(), + "value" + ); +} + +#[rstest::rstest] +#[case(201)] +#[case(409)] +#[case(422)] +#[case(500)] +#[tokio::test] +async fn writes_tolerate_policy_status_and_cache_value(#[case] policy_status: u16) { + let server = MockServer::start().await; + mount_auth(&server, 1).await; + Mock::given(path("/policies/acct/policy/root")) + .and(header("content-type", "application/x-yaml")) + .and(body_string("- !variable \"team/app\"\n")) + .respond_with(ResponseTemplate::new(policy_status)) + .expect(1) + .mount(&server) + .await; + Mock::given(path("/secrets/acct/variable/team%2Fapp")) + .and(body_string("v")) + .respond_with(ResponseTemplate::new(200)) + .expect(1) + .mount(&server) + .await; + let manager = manager(&server, Duration::from_secs(60)); + manager + .async_write_secret("team/app", &SecretValue::new("v"), None) + .await + .unwrap(); + assert_eq!( + manager + .async_read_secret("team/app") + .await + .unwrap() + .unwrap() + .expose(), + "v" + ); +} + +#[tokio::test] +async fn failed_value_write_is_not_cached() { + let server = MockServer::start().await; + mount_auth(&server, 1).await; + Mock::given(path("/policies/acct/policy/root")) + .respond_with(ResponseTemplate::new(409)) + .mount(&server) + .await; + Mock::given(path("/secrets/acct/variable/key")) + .and(body_string("v")) + .respond_with(ResponseTemplate::new(403)) + .expect(1) + .mount(&server) + .await; + Mock::given(path("/secrets/acct/variable/key")) + .respond_with(ResponseTemplate::new(200).set_body_string("recovered")) + .expect(1) + .mount(&server) + .await; + let manager = manager(&server, Duration::from_secs(60)); + assert!(matches!( + manager + .async_write_secret("key", &SecretValue::new("v"), None) + .await, + Err(Error::Status(403)) + )); + assert_eq!( + manager + .async_read_secret("key") + .await + .unwrap() + .unwrap() + .expose(), + "recovered" + ); +} + +#[tokio::test] +async fn unsafe_names_fail_before_http_calls() { + let server = MockServer::start().await; + let manager = manager(&server, Duration::from_secs(60)); + assert!(matches!( + manager + .async_write_secret("../etc", &SecretValue::new("v"), None) + .await, + Err(Error::Operation( + litellm_secrets_types::Error::UnsafeSecretName + )) + )); +} + +#[tokio::test] +async fn delete_invalidates_cache_and_reports_not_supported() { + let server = MockServer::start().await; + mount_auth(&server, 1).await; + Mock::given(path("/secrets/acct/variable/key")) + .respond_with(ResponseTemplate::new(200).set_body_string("v")) + .expect(2) + .mount(&server) + .await; + let manager = manager(&server, Duration::from_secs(60)); + assert_eq!( + manager + .async_read_secret("key") + .await + .unwrap() + .unwrap() + .expose(), + "v" + ); + assert_eq!( + manager.async_delete_secret("key", 7).await.unwrap(), + DeleteOutcome::NotSupported + ); + assert_eq!( + manager + .async_read_secret("key") + .await + .unwrap() + .unwrap() + .expose(), + "v" + ); +} + +#[test] +fn new_validates_credentials_before_license_and_configuration() { + let empty: Arc = + Arc::new(|_: &str| None); + assert!(matches!( + CyberArkSecretManager::new(empty, true), + Err(Error::MissingCredentials) + )); + assert!(matches!( + CyberArkSecretManager::new( + Arc::new(|name: &str| (name == "CYBERARK_API_KEY").then(|| "k3y".into())), + false + ), + Err(Error::EnterpriseRequired) + )); + assert!(matches!( + CyberArkSecretManager::new( + Arc::new(|name: &str| (name == "CYBERARK_CLIENT_CERT").then(|| "cert".into())), + true + ), + Err(Error::MissingCredentials) + )); + assert!(matches!( + CyberArkSecretManager::new( + Arc::new(|name: &str| match name { + "CYBERARK_API_KEY" => Some("k3y".into()), + "CYBERARK_REFRESH_INTERVAL" => Some("abc".into()), + _ => None, + }), + true + ), + Err(Error::RefreshInterval) + )); + assert!(matches!( + CyberArkSecretManager::new( + Arc::new(|name: &str| match name { + "CYBERARK_API_KEY" => Some("k3y".into()), + "CYBERARK_API_BASE" => Some("not a url".into()), + _ => None, + }), + true + ), + Err(Error::Endpoint) + )); +} + +#[tokio::test] +async fn new_reads_environment_defaults_end_to_end() { + let server = MockServer::start().await; + Mock::given(path("/authn/default/admin/authenticate")) + .and(body_string("k3y")) + .respond_with(ResponseTemplate::new(200).set_body_string(TOKEN_JSON)) + .mount(&server) + .await; + Mock::given(path("/secrets/default/variable/key")) + .respond_with(ResponseTemplate::new(200).set_body_string("value")) + .mount(&server) + .await; + let endpoint = server.uri(); + let manager = CyberArkSecretManager::new( + Arc::new(move |name: &str| match name { + "CYBERARK_API_BASE" => Some(endpoint.clone()), + "CYBERARK_API_KEY" => Some("k3y".into()), + _ => None, + }), + true, + ) + .unwrap(); + assert_eq!( + manager + .async_read_secret("key") + .await + .unwrap() + .unwrap() + .expose(), + "value" + ); +} + +#[test] +fn new_reports_missing_client_certificate_files() { + assert!(matches!( + CyberArkSecretManager::new( + Arc::new(|name: &str| match name { + "CYBERARK_CLIENT_CERT" => Some("/missing/cert".into()), + "CYBERARK_CLIENT_KEY" => Some("/missing/key".into()), + _ => None, + }), + true + ), + Err(Error::ClientCertificate) + )); +} + +#[test] +fn secret_name_validation_matches_write_guard() { + assert!(validate_secret_name("../etc").is_err()); +} + +#[tokio::test] +async fn trailing_slash_endpoint_preserves_base_path() { + let server = MockServer::start().await; + Mock::given(path("/prefix/authn/acct/admin/authenticate")) + .and(body_string("k3y")) + .respond_with(ResponseTemplate::new(200).set_body_string(TOKEN_JSON)) + .expect(1) + .mount(&server) + .await; + Mock::given(path("/prefix/secrets/acct/variable/key")) + .respond_with(ResponseTemplate::new(200).set_body_string("value")) + .mount(&server) + .await; + let endpoint = format!("{}/prefix/", server.uri()).parse().unwrap(); + let manager = CyberArkSecretManager::with_client( + reqwest::Client::new(), + endpoint, + "acct".into(), + "admin".into(), + SecretValue::new("k3y"), + Some(Duration::from_secs(60)), + ); + assert_eq!( + manager + .async_read_secret("key") + .await + .unwrap() + .unwrap() + .expose(), + "value" + ); +} + +#[test] +fn parity_fixture_matches_authentication_contract() { + let fixture = fixture(); + assert_eq!(fixture.endpoint, "http://conjur.test:8080"); + assert_eq!(fixture.account, "acct"); + assert_eq!(fixture.username, "admin"); + assert_eq!(fixture.api_key, "k3y"); + assert_eq!(fixture.authenticate_path, "/authn/acct/admin/authenticate"); + assert_eq!(fixture.token_json, TOKEN_JSON); + assert_eq!( + fixture.authorization_header, + format!("Token token=\"{}\"", STANDARD.encode(TOKEN_JSON)) + ); + assert_eq!(fixture.policy_path, "/policies/acct/policy/root"); + assert_eq!(fixture.secrets.len(), 4); + assert_eq!( + fixture.secrets[1].policy_body, + "- !variable \"team/app/key\"\n" + ); +} diff --git a/litellm-rust/crates/secrets/Cargo.toml b/litellm-rust/crates/secrets/Cargo.toml index a7e7ec80636..1298cbd33fe 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"] +cyberark = ["dep:litellm-secrets-cyberark"] [dependencies] litellm-secrets-types.workspace = true litellm-secrets-aws = { workspace = true, optional = true } litellm-secrets-google = { workspace = true, optional = true } +litellm-secrets-cyberark = { 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..cf08c7c2abb 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 = "cyberark")] + #[error(transparent)] + Cyberark(#[from] litellm_secrets_cyberark::Error), } diff --git a/litellm-rust/crates/secrets/src/handler.rs b/litellm-rust/crates/secrets/src/handler.rs index 943ffdf6158..190a67b7671 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 = "cyberark")] + Cyberark(crate::cyberark::CyberArkSecretManager), } impl SecretManager { @@ -27,6 +29,8 @@ impl SecretManager { Self::GoogleKms(_) => KeyManagementSystem::GoogleKms, #[cfg(feature = "google")] Self::GoogleSecretManager(_) => KeyManagementSystem::GoogleSecretManager, + #[cfg(feature = "cyberark")] + Self::Cyberark(_) => KeyManagementSystem::Cyberark, } } } @@ -78,6 +82,12 @@ pub async fn get_secret_from_manager( .get_secret_from_google_secret_manager(secret_name) .await .map_err(Error::from), + #[cfg(feature = "cyberark")] + SecretManager::Cyberark(client) => client + .async_read_secret(secret_name) + .await + .map(|value| value.map(Secret::String)) + .map_err(Error::from), } } diff --git a/litellm-rust/crates/secrets/src/lib.rs b/litellm-rust/crates/secrets/src/lib.rs index ff2e95f7b2f..e132627345f 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 = "cyberark")] +pub use litellm_secrets_cyberark as cyberark; #[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..8f996cad63f 100644 --- a/litellm-rust/crates/secrets/tests/handler.rs +++ b/litellm-rust/crates/secrets/tests/handler.rs @@ -105,3 +105,56 @@ async fn google_handler_requires_canonical_base64_and_preserves_plaintext_whites Err(Error::MissingCiphertext) )); } + +#[cfg(feature = "cyberark")] +#[tokio::test] +async fn cyberark_handler_reads_values_and_surfaces_errors() { + use std::time::Duration; + + use litellm_secrets::{ + Error, KeyManagementSettings, SecretManager, SecretValue, cyberark::CyberArkSecretManager, + get_secret_from_manager, + }; + use wiremock::{ + Mock, MockServer, ResponseTemplate, + matchers::{body_string, path}, + }; + + let server = MockServer::start().await; + Mock::given(path("/authn/acct/admin/authenticate")) + .and(body_string("k3y")) + .respond_with(ResponseTemplate::new(200).set_body_string("token")) + .mount(&server) + .await; + Mock::given(path("/secrets/acct/variable/KEY")) + .respond_with(ResponseTemplate::new(200).set_body_string("value")) + .mount(&server) + .await; + let manager = SecretManager::Cyberark(CyberArkSecretManager::with_client( + reqwest::Client::new(), + server.uri().parse().unwrap(), + "acct".into(), + "admin".into(), + SecretValue::new("k3y"), + Some(Duration::from_secs(60)), + )); + assert_eq!( + manager.system(), + litellm_secrets::KeyManagementSystem::Cyberark + ); + 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")); + + Mock::given(path("/secrets/acct/variable/ERROR")) + .respond_with(ResponseTemplate::new(500)) + .mount(&server) + .await; + assert!(matches!( + get_secret_from_manager(&manager, "ERROR", &settings, &|_: &str| None).await, + Err(Error::Cyberark(_)) + )); +} diff --git a/tests/test_litellm/secret_managers/test_cyberark_secret_manager.py b/tests/test_litellm/secret_managers/test_cyberark_secret_manager.py new file mode 100644 index 00000000000..3f3669ab9ef --- /dev/null +++ b/tests/test_litellm/secret_managers/test_cyberark_secret_manager.py @@ -0,0 +1,119 @@ +import json +from pathlib import Path +from typing import Final, TypedDict, cast + +import pytest +import respx + +import litellm +import litellm.proxy.proxy_server +from litellm.secret_managers.cyberark_secret_manager import CyberArkSecretManager + +FIXTURE_PATH: Final = Path(__file__).resolve().parents[3] / "litellm-rust/crates/secrets-cyberark/tests/fixtures/parity.json" + + +class ParitySecret(TypedDict): + name: str + path: str + policy_body: str + + +class ParityFixture(TypedDict): + endpoint: str + account: str + username: str + api_key: str + authenticate_path: str + token_json: str + authorization_header: str + policy_path: str + secrets: list[ParitySecret] + + +def _fixture() -> ParityFixture: + return cast(ParityFixture, json.loads(FIXTURE_PATH.read_text())) + + +def _configure_manager(monkeypatch: pytest.MonkeyPatch, fixture: ParityFixture) -> CyberArkSecretManager: + monkeypatch.setattr(litellm.proxy.proxy_server, "premium_user", True) + monkeypatch.setenv("CYBERARK_API_BASE", fixture["endpoint"]) + monkeypatch.setenv("CYBERARK_ACCOUNT", fixture["account"]) + monkeypatch.setenv("CYBERARK_USERNAME", fixture["username"]) + monkeypatch.setenv("CYBERARK_API_KEY", fixture["api_key"]) + monkeypatch.setenv("CYBERARK_REFRESH_INTERVAL", "300") + monkeypatch.delenv("CYBERARK_CLIENT_CERT", raising=False) + monkeypatch.delenv("CYBERARK_CLIENT_KEY", raising=False) + return CyberArkSecretManager() + + +def _respond( + route: respx.Route, + *, + status_code: int = 200, + content: str | bytes | None = None, + text: str | None = None, +) -> respx.Route: + return route.respond( # pyright: ignore[reportUnknownMemberType] # respx route stubs leave response builder partially unknown + status_code=status_code, + content=content, + text=text, + ) + + +@respx.mock +def test_sync_read_matches_parity_fixture(monkeypatch: pytest.MonkeyPatch) -> None: + fixture: Final = _fixture() + manager: Final = _configure_manager(monkeypatch, fixture) + endpoint: Final = fixture["endpoint"] + token_json: Final = fixture["token_json"] + auth_route: Final = _respond( + respx.post(endpoint + fixture["authenticate_path"]), + content=token_json.encode(), + ) + routes: Final = [ + _respond(respx.get(endpoint + secret["path"]), text="value") + for secret in fixture["secrets"] + ] + + for secret in fixture["secrets"]: + assert manager.sync_read_secret(secret["name"]) == "value" # pyright: ignore[reportUnknownMemberType] # legacy secret manager API is untyped + + expected_authorization: Final = fixture["authorization_header"] + assert auth_route.calls.last.request.content == fixture["api_key"].encode() + assert all(route.calls.last.request.headers["Authorization"] == expected_authorization for route in routes) + assert all( + route.calls.last.request.url.raw_path.decode() == secret["path"] + for route, secret in zip(routes, fixture["secrets"], strict=True) + ) + + +@pytest.mark.asyncio +@respx.mock +async def test_async_write_matches_parity_fixture(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + fixture: Final = _fixture() + manager: Final = _configure_manager(monkeypatch, fixture) + secret: Final = fixture["secrets"][0] + endpoint: Final = fixture["endpoint"] + token_json: Final = fixture["token_json"] + _respond(respx.post(endpoint + fixture["authenticate_path"]), content=token_json.encode()) + policy_route: Final = _respond(respx.post(endpoint + fixture["policy_path"]), status_code=201) + value_route: Final = _respond(respx.post(endpoint + secret["path"]), status_code=201) + + await manager.async_write_secret(secret["name"], "v") # pyright: ignore[reportUnknownMemberType] # legacy secret manager API is untyped + + assert policy_route.calls.last.request.content.decode() == secret["policy_body"] + assert policy_route.calls.last.request.headers["Content-Type"] == "application/x-yaml" + assert value_route.calls.last.request.content == b"v" + + +def test_missing_credentials_raise_value_error(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm.proxy.proxy_server, "premium_user", True) + for name in ( + "CYBERARK_API_KEY", + "CYBERARK_CLIENT_CERT", + "CYBERARK_CLIENT_KEY", + ): + monkeypatch.delenv(name, raising=False) + with pytest.raises(ValueError, match="Missing CyberArk credentials"): + CyberArkSecretManager() From fe6804ea74a7ec28bd3767b2c54647bee9d5fe5d Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:25:07 +0000 Subject: [PATCH 06/24] refactor(rust): reject negative CyberArk refresh intervals Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../crates/secrets-cyberark/src/secret_manager.rs | 8 ++------ .../crates/secrets-cyberark/tests/secret_manager.rs | 7 +------ 2 files changed, 3 insertions(+), 12 deletions(-) diff --git a/litellm-rust/crates/secrets-cyberark/src/secret_manager.rs b/litellm-rust/crates/secrets-cyberark/src/secret_manager.rs index ade257a5872..ee4bded4b29 100644 --- a/litellm-rust/crates/secrets-cyberark/src/secret_manager.rs +++ b/litellm-rust/crates/secrets-cyberark/src/secret_manager.rs @@ -116,13 +116,9 @@ impl CyberArkSecretManager { .get(CYBERARK_REFRESH_INTERVAL) .map(|value| { value - .parse::() + .parse::() + .map(Duration::from_secs) .map_err(|_| Error::RefreshInterval) - .map(|seconds| match seconds { - seconds if seconds < 0 => Duration::from_nanos(1), - 0 => DEFAULT_REFRESH_INTERVAL, - seconds => Duration::from_secs(seconds as u64), - }) }) .transpose()?; Ok(Self::with_client( diff --git a/litellm-rust/crates/secrets-cyberark/tests/secret_manager.rs b/litellm-rust/crates/secrets-cyberark/tests/secret_manager.rs index e8158c46945..9ea8da63383 100644 --- a/litellm-rust/crates/secrets-cyberark/tests/secret_manager.rs +++ b/litellm-rust/crates/secrets-cyberark/tests/secret_manager.rs @@ -2,7 +2,7 @@ use std::{sync::Arc, time::Duration}; use base64::{Engine, engine::general_purpose::STANDARD}; use litellm_secrets_cyberark::{CyberArkSecretManager, DeleteOutcome, Error}; -use litellm_secrets_types::{SecretValue, validate_secret_name}; +use litellm_secrets_types::SecretValue; use serde::Deserialize; use wiremock::{ Match, Mock, MockServer, Request, ResponseTemplate, @@ -428,11 +428,6 @@ fn new_reports_missing_client_certificate_files() { )); } -#[test] -fn secret_name_validation_matches_write_guard() { - assert!(validate_secret_name("../etc").is_err()); -} - #[tokio::test] async fn trailing_slash_endpoint_preserves_base_path() { let server = MockServer::start().await; From ae69a8c79a82d54d50436c1eb06ab9cca3625574 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:30:48 +0000 Subject: [PATCH 07/24] 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 08/24] 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 7d597b2dd4232c9ad6ab59d7db760feb2be18410 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:36:59 +0000 Subject: [PATCH 09/24] fix(rust): coalesce CyberArk authentication Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../crates/secrets-cyberark/Cargo.toml | 1 + .../secrets-cyberark/src/secret_manager.rs | 6 ++++ .../secrets-cyberark/tests/secret_manager.rs | 33 +++++++++++++++++++ 3 files changed, 40 insertions(+) diff --git a/litellm-rust/crates/secrets-cyberark/Cargo.toml b/litellm-rust/crates/secrets-cyberark/Cargo.toml index de8ae5b71dc..3c1159c40be 100644 --- a/litellm-rust/crates/secrets-cyberark/Cargo.toml +++ b/litellm-rust/crates/secrets-cyberark/Cargo.toml @@ -16,6 +16,7 @@ thiserror.workspace = true veil.workspace = true tracing = "0.1" percent-encoding = "2.3" +tokio = { workspace = true, features = ["sync"] } [dev-dependencies] rstest.workspace = true diff --git a/litellm-rust/crates/secrets-cyberark/src/secret_manager.rs b/litellm-rust/crates/secrets-cyberark/src/secret_manager.rs index ee4bded4b29..9d6eaaf1c4e 100644 --- a/litellm-rust/crates/secrets-cyberark/src/secret_manager.rs +++ b/litellm-rust/crates/secrets-cyberark/src/secret_manager.rs @@ -35,6 +35,7 @@ pub struct CyberArkSecretManager { api_key: SecretValue, token: Cache<(), SecretValue>, secrets: Cache, + authentication_lock: Arc>, } #[derive(Clone, Copy, Debug, PartialEq, Eq)] @@ -65,6 +66,7 @@ impl CyberArkSecretManager { api_key, token, secrets, + authentication_lock: Arc::new(tokio::sync::Mutex::new(())), } } @@ -139,6 +141,10 @@ impl CyberArkSecretManager { } async fn authenticate(&self) -> Result { + if let Some(token) = self.token.get(&()).await { + return Ok(token); + } + let _guard = self.authentication_lock.lock().await; if let Some(token) = self.token.get(&()).await { return Ok(token); } diff --git a/litellm-rust/crates/secrets-cyberark/tests/secret_manager.rs b/litellm-rust/crates/secrets-cyberark/tests/secret_manager.rs index 9ea8da63383..fd7198b70fb 100644 --- a/litellm-rust/crates/secrets-cyberark/tests/secret_manager.rs +++ b/litellm-rust/crates/secrets-cyberark/tests/secret_manager.rs @@ -89,6 +89,39 @@ async fn successful_reads_cache_auth_secret_and_redact_values() { } } +#[tokio::test] +async fn concurrent_reads_share_authentication_request() { + let server = MockServer::start().await; + Mock::given(path("/authn/acct/admin/authenticate")) + .and(body_string("k3y")) + .respond_with( + ResponseTemplate::new(200) + .set_body_string(TOKEN_JSON) + .set_delay(Duration::from_millis(20)), + ) + .expect(1) + .mount(&server) + .await; + Mock::given(path("/secrets/acct/variable/key")) + .and(header( + "authorization", + format!("Token token=\"{}\"", STANDARD.encode(TOKEN_JSON)), + )) + .respond_with(ResponseTemplate::new(200).set_body_string("value")) + .expect(2) + .mount(&server) + .await; + let manager = manager(&server, Duration::from_secs(60)); + + let (first, second) = tokio::join!( + manager.async_read_secret("key"), + manager.async_read_secret("key") + ); + + assert_eq!(first.unwrap().unwrap().expose(), "value"); + assert_eq!(second.unwrap().unwrap().expose(), "value"); +} + #[rstest::rstest] #[case::not_found(404)] #[case::unauthorized(401)] From df8966591939e6aa4f747429b762ec33a6721543 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:37:57 +0000 Subject: [PATCH 10/24] 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 11/24] 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 12/24] 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 13/24] 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 14/24] 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 15/24] 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 16/24] 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 17/24] 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 18/24] 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)}"), ) From a41061be43e8c77ec57cf208f3dbf2ecd46a9553 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 12:49:42 -0700 Subject: [PATCH 19/24] docs(rust): plan Python interop foundation --- PYTHON_INTEROP_PLAN.md | 88 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 88 insertions(+) create mode 100644 PYTHON_INTEROP_PLAN.md diff --git a/PYTHON_INTEROP_PLAN.md b/PYTHON_INTEROP_PLAN.md new file mode 100644 index 00000000000..b456d795cfd --- /dev/null +++ b/PYTHON_INTEROP_PLAN.md @@ -0,0 +1,88 @@ +# Python interop foundation PR plan + +Proposed implementation PR title: `fix(rust): preserve Python settings semantics at the native boundary` + +Base: `main` at `457b01e96d131f88df8cace8832f8e044ee5f167`. Planning branch: `litellm_python_interop_foundation` + +This plan follows `migration/tdd/sections/rust/python-interop/{pyo3-contract,boundary,coercion}.typ` in the sibling `litellm-typst` checkout. The first PR establishes conversion contracts and exercises them through existing HTTP, URL policy, and OCR provider-default consumers. Implementation, runtime validation, and PR creation remain future work + +**What already exists** + +`host-python/src/marshal.rs` already uses `pythonize` directly, separates internal conversion from public argument errors, and contains serializer panics in `Pythonized`. Keep those entrypoints. `Pythonized` currently stringifies conversion errors, unlike `from_py` and `to_py`, so its error transfer needs correction + +`core-utils/src/serde_compat.rs` already provides composable `LaxI64` and `FiniteF64` adapters. They first deserialize into `serde_json::Value`. Preserve their accepted input contracts while moving scalar decoding to Serde visitors, avoiding an intermediate JSON representation and making behavior through `pythonize` explicit + +`python-bridge/src/http.rs` currently uses strict derived Boolean and string extraction for mutable settings, converts extraction errors into `RustBridgeDeclined`, and silently drops unsupported `ssl_verify` values. OCR provider defaults have the same strict-extraction/decline pattern. `python_settings.json` checks field names only. These are the initial production consumers and regressions for the foundation + +**Ownership and placement** + +Keep interpreter attachment, structured-data conversion, panic containment, and execution adapters in `host-python`. Keep field paths, configuration error policy, named settings coercion, and product-specific tagged inputs in `python-bridge`. Keep pure parsing and Serde adapters in `core-utils`. This requires no new crate and no domain dependency from `host-python` into `core-utils` + +Add one focused `python-bridge/src/coercion.rs` module containing the field reader, semantic wrappers, and projection errors. Use `core-utils/src/serde_compat.rs` for shared token/numeric parsing and Serde decoding initially, splitting only if its size warrants it. The bridge can call those pure helpers through its existing dependency. Do not introduce a generic coercion registry, runtime manifest dispatch, or a new conversion framework + +Settings snapshots hold raw `Bound<'py, PyAny>` values only during attached projection. Successful adapters return owned values. Python snapshot dataclass annotations use `object` for arbitrary mutable globals, retaining strict annotations for accessor-owned fields such as `user_agent` and `readable`. No live settings object, iterator, or borrowed Python value enters native HTTP state + +**Serde and structured conversion** + +Keep `from_py` and `to_py` as the internal, exception-preserving conversion path. Make `Pythonized` use the same standard `PythonizeError` to `PyErr` conversion without losing its panic guard. Keep the explicitly argument-focused `ValueError` contract of `from_py_argument`; settings projection never passes through that helper + +Implement scalar visitors behind the existing `LaxI64` and `FiniteF64` adapter names. Preserve integer precision beyond the exact f64 range, signed bounds, supported decimal/underscore strings, Boolean numeric behavior, fractional rejection for integers, and nonfinite rejection for floats. Preserve composition inside `Option` and sequences, missing/null behavior at the field boundary, and ordinary numeric serialization. Use the existing [serde_with DeserializeAs contract](https://docs.rs/serde_with/3.16.1/serde_with/trait.DeserializeAs.html) + +Add the pure `parse_str_bool(&str) -> Option` parser used by bridge `StrBool`, HTTP `SslVerify::parse`, and the environment switch helper. It recognizes trimmed, case-insensitive true/false only. The environment helper retains its separate rule that only `Some(true)` contributes an enabled layer. Numeric helpers are shared only where a second actual consumer needs them + +Run common ordinary-value fixtures through JSON deserialization and direct Python-to-typed-Serde conversion. Restrict equivalence claims to their overlapping input domain. Live descriptors, identity, truthiness, iteration, and stringification are exercised through PyO3 separately. Do not add unused Serde counterparts for Python-only semantics or convert live settings through JSON text, `serde_json::Value`, `repr`, or `py_literal` + +**Named settings semantics** + +| Adapter | Contract and first consumer | +| --- | --- | +| `Truthy` | Execute Python truth testing and preserve its exception; IPv4, URL validation, and trust-env globals | +| `ExactTrue` | Compare identity with the True singleton without equality or truth testing; HTTP2, transport disable, and token refresh | +| `StrBool` | None or an actual string parsed by the shared parser; no arbitrary stringification | +| `OptionalStrictString` | None or an actual string, including the empty string; client certificate | +| `FalsyOptionalString` | Test truthiness first, treat falsey as absent, reject truthy non-strings; provider project and location | +| `TuningString` | Test truthiness first, then retain actual strings and ignore other values; TLS tuning | +| `StringCollection` | Apply the field's explicit container/member policy and return owned strings; URL allowlist | +| `SslVerifyInput` | Classify None, actual Boolean, Boolean string, CA path, live SSLContext, and invalid types separately | + +A single generic Boolean or optional-string coercer cannot implement these contracts. Accept string subclasses by their Unicode contents without calling overridden convenience methods. Use no `.ok()` or default value to discard an error from a Python protocol operation + +For URL hosts, a direct string represents one host. Otherwise test container truthiness and iterate, test member truthiness, skip falsey members, and reject truthy non-string members. Normalize with the existing URL-policy rules, deduplicate, and sort only where membership makes order irrelevant. Keep origin/scheme/port parsing in its existing domain helper rather than applying hostname normalization blindly to arbitrary URL strings + +**Errors and first production adoption** + +Represent projection failures as a tagged result separating original Python exceptions, invalid configuration, unsupported live objects, and internal accessor/schema failures. Map it once at the bridge: preserve original `PyErr`, use field-focused `ValueError` for invalid or unsupported configuration, and `RuntimeError` for genuine internal schema failures. Diagnostics contain group, field, expected forms, and actual type, never the supplied value or its representation + +Attribute, truthiness, iteration, and explicit stringification exceptions retain identity, traceback, cause, and context. In particular, do not relabel an AttributeError deliberately raised by a descriptor as a missing-field schema failure. Contract validation must distinguish schema drift from errors executing Python behavior + +Convert HTTP and URL snapshots, per-call `ssl_verify`, and OCR provider defaults through the named adapters. Preserve existing call/environment/global precedence and projection timing. Finish projection before constructing the native client or starting provider I/O. Invalid values and live SSLContext must raise configuration errors rather than disappear or authorize fallback + +Keep certificate paths through projection and validate empty, missing, or unusable client-certificate paths before I/O. The current native layer filters empty client-certificate paths, so correcting that narrow downstream behavior is part of adoption. Keep the existing CA-bundle missing-file policy explicit and separately tested; do not silently conflate it with client-certificate validation + +Classify the existing Secret Manager `readable` field as a strict accessor Boolean. Full Secret Manager client/system/settings snapshots and callback execution remain a separate PR. The existing readable-manager capability gap must be documented and must not be reported as fixed by this foundation + +**Semantic manifest** + +Extend `python_settings.json` with a stable group version and field records containing adapter ID, requiredness, precedence role, sensitivity, and specialized accepted/unsupported shapes. Include only the snapshot fields that exist in this PR. Update the Python contract test and a static Rust `SettingSpec` table to agree with the manifest + +The manifest checks declared contracts; behavioral tests prove the adapters implement them. Projection stays direct typed code. A manifest row alone is never evidence that a coercion works + +**Behavioral validation** + +Extend the existing mapped Python settings tests and Rust marshal/HTTP tests. A new coercion module may have its own focused Rust tests. Use the existing installed-extension OCR suites for public-route regressions. Do not add source-text assertions or class-attribute monkeypatching + +The acceptance matrix covers None, Boolean values, integer zero/one, strings, containers, subclasses, and arbitrary objects. Protocol fixtures raise pre-created exceptions from descriptors, `__bool__`, `__len__`, `__iter__`, and `__next__`; assert identity and exception chains. Verify ExactTrue never invokes hostile equality/truthiness. Verify falsey provider defaults preserve fallback, HTTP2 does not accept integer one, and a false/unknown environment token cannot switch off a true global + +Cover a real SSLContext, unsupported objects, Boolean strings, certificate path failures, direct-string hosts, sets, generators, duplicates, mixed members, and protocol failures. Mutate globals and source collections between calls: the next snapshot observes changes and an already projected value stays unchanged. At the installed public OCR boundary, projection failures must cause zero provider requests and zero Python fallback calls under required-native execution + +Run focused crate tests first, then the workspace Rust checks used by CI, `make test-rust-extension`, relevant Python settings tests, and `make check`. Use the fresh installed wheel and verify native provenance. Review the saved `make check` log rather than rerunning it to inspect output + +Target mutation tests at truthiness versus strict extraction, identity versus equality, swallowed versus preserved exceptions, string-as-one versus character iteration, normalization/deduplication, numeric bounds, and terminal errors versus fallback. Aim for more than 90% killed non-equivalent mutants in the changed coercion paths + +Before opening the implementation PR for maintainer review, provide a reproducible localhost proxy curl request with a real provider and positive native-execution evidence. Record the configured settings and user-visible result without credentials. Unit tests belong in validation, not the proof-of-fix section. Require the current tip's CI and coverage, Greptile confidence of at least 4/5, and acceptable Veria/Bugbot results; pending or unavailable results remain explicitly unresolved + +**Commit sequence and follow-ups** + +Start with the Serde visitor/error-transfer changes and their regression tests. Follow with field adapters and the shared token parser. Adopt them in HTTP, URL policy, and provider defaults together with the semantic manifest and installed-extension regressions. Keep these as reviewable commits in one foundational implementation PR + +Follow-up PRs can add `OptionalRedisBool`, cache-specific stringification and collection rules, and full Secret Manager bindings using the same field/error machinery. Redis accepts a different token set from StrBool, so do not share their Boolean semantics. Runtime redesign, callback lifecycle changes, free-threaded support, wholesale request serialization, and unrelated cache work are outside this PR From 52216df1ff3dffd8c97ee8ea459ebcd4cd20b17f Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 13:01:57 -0700 Subject: [PATCH 20/24] fix(rust): preserve structured conversion semantics --- litellm-rust/Cargo.lock | 1 + .../crates/core-utils/src/serde_compat.rs | 112 +++++++++++++++--- .../crates/host-python/src/marshal.rs | 15 ++- litellm-rust/crates/python-bridge/Cargo.toml | 2 + .../crates/python-bridge/src/marshal.rs | 54 +++++++++ 5 files changed, 164 insertions(+), 20 deletions(-) diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index ea4f3d848c4..911cc571df9 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2688,6 +2688,7 @@ dependencies = [ "rstest", "serde", "serde_json", + "serde_with", "tokio", "tokio-tungstenite", ] diff --git a/litellm-rust/crates/core-utils/src/serde_compat.rs b/litellm-rust/crates/core-utils/src/serde_compat.rs index bb2648eb0be..c767c709f50 100644 --- a/litellm-rust/crates/core-utils/src/serde_compat.rs +++ b/litellm-rust/crates/core-utils/src/serde_compat.rs @@ -1,33 +1,91 @@ -use serde::{Deserialize, Deserializer, de::Error}; -use serde_json::Value; +use serde::{ + Deserializer, + de::{Error, Visitor}, +}; use serde_with::DeserializeAs; pub struct LaxI64; pub struct FiniteF64; +pub fn parse_str_bool(value: &str) -> Option { + let token = value.trim_matches(|character: char| { + character.is_whitespace() || matches!(character, '\u{1c}'..='\u{1f}') + }); + if token.eq_ignore_ascii_case("true") { + return Some(true); + } + token.eq_ignore_ascii_case("false").then_some(false) +} + impl<'de> DeserializeAs<'de, i64> for LaxI64 { fn deserialize_as>(deserializer: D) -> Result { - match Value::deserialize(deserializer)? { - Value::Number(number) if number.is_f64() => number.as_f64().and_then(integral_float), - Value::Number(number) => number.as_i64(), - Value::String(value) => integer_string(value.trim()), - Value::Bool(value) => Some(i64::from(value)), - _ => None, - } - .ok_or_else(|| D::Error::custom("expected an integer in the i64 range")) + deserializer.deserialize_any(Self) + } +} + +impl<'de> Visitor<'de> for LaxI64 { + type Value = i64; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("an integer in the i64 range") + } + + fn visit_i64(self, value: i64) -> Result { + Ok(value) + } + + fn visit_u64(self, value: u64) -> Result { + i64::try_from(value).map_err(E::custom) + } + + fn visit_f64(self, value: f64) -> Result { + integral_float(value).ok_or_else(|| E::custom("expected an integer in the i64 range")) + } + + fn visit_str(self, value: &str) -> Result { + integer_string(value.trim()) + .ok_or_else(|| E::custom("expected an integer in the i64 range")) + } + + fn visit_bool(self, value: bool) -> Result { + Ok(i64::from(value)) } } impl<'de> DeserializeAs<'de, f64> for FiniteF64 { fn deserialize_as>(deserializer: D) -> Result { - match Value::deserialize(deserializer)? { - Value::Number(number) => number.as_f64(), - Value::String(value) => value.trim().parse::().ok(), - Value::Bool(value) => Some(f64::from(value)), - _ => None, - } - .filter(|value| value.is_finite()) - .ok_or_else(|| D::Error::custom("expected a finite number")) + deserializer.deserialize_any(Self) + } +} + +impl<'de> Visitor<'de> for FiniteF64 { + type Value = f64; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("a finite number") + } + + fn visit_i64(self, value: i64) -> Result { + Ok(value as f64) + } + + fn visit_u64(self, value: u64) -> Result { + Ok(value as f64) + } + + fn visit_f64(self, value: f64) -> Result { + value + .is_finite() + .then_some(value) + .ok_or_else(|| E::custom("expected a finite number")) + } + + fn visit_str(self, value: &str) -> Result { + self.visit_f64(value.trim().parse::().map_err(E::custom)?) + } + + fn visit_bool(self, value: bool) -> Result { + Ok(f64::from(value)) } } @@ -66,7 +124,7 @@ fn integral_float(value: f64) -> Option { #[cfg(test)] mod tests { - use serde::Serialize; + use serde::{Deserialize, Serialize}; use serde_json::json; use serde_with::serde_as; @@ -81,6 +139,22 @@ mod tests { float: Option, } + #[test] + fn boolean_tokens_follow_python_string_trimming_without_redis_tokens() { + for (input, expected) in [ + (" True ", Some(true)), + ("\u{1c}TRUE\u{1f}", Some(true)), + ("\u{a0}False\u{2003}", Some(false)), + ("true\u{200b}", None), + ("yes", None), + ("1", None), + ("", None), + ("unknown", None), + ] { + assert_eq!(parse_str_bool(input), expected, "{input:?}"); + } + } + #[test] fn adapters_compose_and_serialize_as_numbers() { let numbers: Numbers = serde_json::from_value(json!({ diff --git a/litellm-rust/crates/host-python/src/marshal.rs b/litellm-rust/crates/host-python/src/marshal.rs index 881ad0e0389..8f284abf9dd 100644 --- a/litellm-rust/crates/host-python/src/marshal.rs +++ b/litellm-rust/crates/host-python/src/marshal.rs @@ -45,7 +45,7 @@ where fn into_pyobject(self, py: Python<'py>) -> PyResult { catch_unwind(AssertUnwindSafe(|| pythonize::pythonize(py, &self.0))) .map_err(panic_to_pyerr)? - .map_err(|error| PyValueError::new_err(error.to_string())) + .map_err(PyErr::from) } } @@ -87,6 +87,19 @@ mod tests { }); } + #[test] + fn pythonized_preserves_python_serialization_error_types() { + crate::initialize_python(); + Python::attach(|py| { + let value = std::collections::BTreeMap::from([(vec![1], "value")]); + let direct = to_py(py, &value).unwrap_err(); + let wrapped = Pythonized(value).into_pyobject(py).unwrap_err(); + assert!(direct.is_instance_of::(py)); + assert!(wrapped.is_instance_of::(py)); + assert_eq!(wrapped.to_string(), direct.to_string()); + }); + } + #[test] fn pythonized_maps_serializer_panics_to_a_base_exception() { crate::initialize_python(); diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 1eb2ec28036..1bba83922f9 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -41,6 +41,8 @@ serde_json.workspace = true tokio = { workspace = true, features = ["sync"] } [dev-dependencies] +serde.workspace = true +serde_with.workspace = true criterion.workspace = true futures-util.workspace = true rstest.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/marshal.rs b/litellm-rust/crates/python-bridge/src/marshal.rs index 2aba51cc4ff..fe5d551a931 100644 --- a/litellm-rust/crates/python-bridge/src/marshal.rs +++ b/litellm-rust/crates/python-bridge/src/marshal.rs @@ -172,6 +172,60 @@ mod tests { request_input_sources(&kwargs, names.iter().copied()) } + #[serde_with::serde_as] + #[derive(Debug, serde::Deserialize, serde::Serialize, PartialEq)] + struct Numbers { + #[serde_as(deserialize_as = "Option>")] + integers: Option>, + #[serde_as(deserialize_as = "Option")] + float: Option, + } + + #[test] + fn numeric_adapters_agree_across_json_and_python_boundaries() { + Python::initialize(); + Python::attach(|py| { + for input in [ + json!({}), + json!({"integers": null, "float": null}), + json!({"integers": [i64::MIN, i64::MAX, "9007199254740993.0", " +1_000.00 ", true, 3.0], "float": " 1.25 "}), + json!({"integers": [u64::MAX]}), + json!({"integers": ["1.0000000000000001"]}), + json!({"integers": [2.5]}), + json!({"float": "NaN"}), + json!({"float": "inf"}), + json!({"float": "1e999"}), + json!({"float": true}), + json!({"float": u64::MAX}), + ] { + let expected = serde_json::from_value::(input.clone()); + let python = litellm_host_python::to_py(py, &input).unwrap(); + let actual = from_py::(python.bind(py)); + match (expected, actual) { + (Ok(expected), Ok(actual)) => { + assert_eq!(actual, expected); + let serialized = litellm_host_python::to_py(py, &actual).unwrap(); + assert_eq!( + from_py::(serialized.bind(py)).unwrap(), + serde_json::to_value(expected).unwrap() + ); + } + (Err(_), Err(_)) => {} + mismatch => panic!("boundary mismatch for {input}: {mismatch:?}"), + } + } + for source in [ + c"{'float': float('nan')}", + c"{'float': float('inf')}", + c"{'integers': [float('inf')]}", + c"{'integers': [2 ** 100]}", + ] { + let value = py.eval(source, None, None).unwrap(); + assert!(from_py::(&value).is_err()); + } + }); + } + #[test] fn argument_converters_keep_nested_values_and_accept_explicit_none() { Python::initialize(); From 5aeb367d2a42ec4050c7db17813dd95b1e9e5839 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 13:01:57 -0700 Subject: [PATCH 21/24] fix(rust): preserve Python settings coercion at the native boundary --- PYTHON_INTEROP_PLAN.md | 4 +- .../crates/core-utils/src/settings.rs | 4 +- litellm-rust/crates/http/src/media.rs | 2 +- litellm-rust/crates/http/src/settings.rs | 19 +- .../crates/python-bridge/python_settings.json | 176 +++++++-- .../crates/python-bridge/src/coercion.rs | 231 +++++++++++ .../python-bridge/src/coercion/tests.rs | 372 ++++++++++++++++++ litellm-rust/crates/python-bridge/src/http.rs | 192 +++++---- litellm-rust/crates/python-bridge/src/lib.rs | 1 + .../python-bridge/src/python_settings.rs | 216 ++++++++-- .../python-bridge/src/routes/ocr/mod.rs | 73 ++-- litellm/rust_bridge/settings.py | 29 +- .../test_litellm/rust_bridge/test_settings.py | 29 +- tests/test_litellm_rust/ocr/test_requests.py | 150 ++++++- 14 files changed, 1309 insertions(+), 189 deletions(-) create mode 100644 litellm-rust/crates/python-bridge/src/coercion.rs create mode 100644 litellm-rust/crates/python-bridge/src/coercion/tests.rs diff --git a/PYTHON_INTEROP_PLAN.md b/PYTHON_INTEROP_PLAN.md index b456d795cfd..a7783f81472 100644 --- a/PYTHON_INTEROP_PLAN.md +++ b/PYTHON_INTEROP_PLAN.md @@ -4,9 +4,9 @@ Proposed implementation PR title: `fix(rust): preserve Python settings semantics Base: `main` at `457b01e96d131f88df8cace8832f8e044ee5f167`. Planning branch: `litellm_python_interop_foundation` -This plan follows `migration/tdd/sections/rust/python-interop/{pyo3-contract,boundary,coercion}.typ` in the sibling `litellm-typst` checkout. The first PR establishes conversion contracts and exercises them through existing HTTP, URL policy, and OCR provider-default consumers. Implementation, runtime validation, and PR creation remain future work +This plan follows `migration/tdd/sections/rust/python-interop/{pyo3-contract,boundary,coercion}.typ` in the sibling `litellm-typst` checkout. The first PR establishes conversion contracts and exercises them through existing HTTP, URL policy, and OCR provider-default consumers. The foundation is implemented on this branch. No PR is being created for this task -**What already exists** +**Baseline before implementation** `host-python/src/marshal.rs` already uses `pythonize` directly, separates internal conversion from public argument errors, and contains serializer panics in `Pythonized`. Keep those entrypoints. `Pythonized` currently stringifies conversion errors, unlike `from_py` and `to_py`, so its error transfer needs correction diff --git a/litellm-rust/crates/core-utils/src/settings.rs b/litellm-rust/crates/core-utils/src/settings.rs index 59c76ce3015..293dc71871c 100644 --- a/litellm-rust/crates/core-utils/src/settings.rs +++ b/litellm-rust/crates/core-utils/src/settings.rs @@ -1,5 +1,7 @@ use std::str::FromStr; +use crate::serde_compat::parse_str_bool; + pub trait Lookup { fn get(&self, name: &str) -> Option; @@ -9,7 +11,7 @@ pub trait Lookup { fn enabled(&self, name: &str) -> Option { self.get(name) - .is_some_and(|value| value.trim().eq_ignore_ascii_case("true")) + .is_some_and(|value| parse_str_bool(&value) == Some(true)) .then_some(true) } diff --git a/litellm-rust/crates/http/src/media.rs b/litellm-rust/crates/http/src/media.rs index 3b29c9e28a7..753f25c29de 100644 --- a/litellm-rust/crates/http/src/media.rs +++ b/litellm-rust/crates/http/src/media.rs @@ -62,7 +62,7 @@ impl UrlPolicy { } } -fn normalize_host(host: &str) -> String { +pub fn normalize_host(host: &str) -> String { host.to_ascii_lowercase().trim_end_matches('.').to_owned() } diff --git a/litellm-rust/crates/http/src/settings.rs b/litellm-rust/crates/http/src/settings.rs index a6397f1e8e3..e1edc6d37e1 100644 --- a/litellm-rust/crates/http/src/settings.rs +++ b/litellm-rust/crates/http/src/settings.rs @@ -3,7 +3,10 @@ use std::{ time::Duration, }; -use litellm_core_utils::settings::{Layer, Lookup, merge}; +use litellm_core_utils::{ + serde_compat::parse_str_bool, + settings::{Layer, Lookup, merge}, +}; use crate::proxy::EnvironmentProxies; @@ -16,9 +19,9 @@ pub enum SslVerify { impl SslVerify { pub fn parse(value: &str) -> Self { - match value.trim().to_ascii_lowercase().as_str() { - "true" => Self::Enabled, - "false" => Self::Disabled, + match parse_str_bool(value) { + Some(true) => Self::Enabled, + Some(false) => Self::Disabled, _ => Self::CaBundle(PathBuf::from(value)), } } @@ -152,9 +155,7 @@ impl HttpSettings { Self { ssl_verify: merged.ssl_verify, ssl_cert_file: merged.ssl_cert_file, - ssl_certificate: merged - .ssl_certificate - .filter(|path| !path.as_os_str().is_empty()), + ssl_certificate: merged.ssl_certificate, ssl_security_level: merged.ssl_security_level.filter(|level| !level.is_empty()), ssl_ecdh_curve: merged.ssl_ecdh_curve.filter(|curve| !curve.is_empty()), force_ipv4: merged.force_ipv4.unwrap_or(defaults.force_ipv4), @@ -287,7 +288,7 @@ mod tests { } #[test] - fn empty_environment_values_clear_the_setting_like_python_truthiness() { + fn empty_certificate_is_retained_for_validation_while_empty_tuning_is_absent() { let configured = HttpSettingsLayer { ssl_certificate: Some("/configured/client.pem".into()), ssl_security_level: Some("configured".into()), @@ -300,7 +301,7 @@ mod tests { ("SSL_ECDH_CURVE", ""), ])); let settings = HttpSettings::from_layers([environment, configured]); - assert_eq!(settings.ssl_certificate, None); + assert_eq!(settings.ssl_certificate, Some(PathBuf::new())); assert_eq!(settings.ssl_security_level, None); assert_eq!(settings.ssl_ecdh_curve, None); } diff --git a/litellm-rust/crates/python-bridge/python_settings.json b/litellm-rust/crates/python-bridge/python_settings.json index 0af55083bef..ea53d1d2025 100644 --- a/litellm-rust/crates/python-bridge/python_settings.json +++ b/litellm-rust/crates/python-bridge/python_settings.json @@ -1,26 +1,154 @@ { - "http_settings": [ - "ssl_verify", - "ssl_certificate", - "ssl_security_level", - "ssl_ecdh_curve", - "force_ipv4", - "http2", - "aiohttp_trust_env", - "disable_aiohttp_trust_env", - "disable_aiohttp_transport", - "user_agent" - ], - "url_policy": [ - "user_url_validation", - "user_url_allowed_hosts" - ], - "provider_defaults": [ - "vertex_project", - "vertex_location", - "enable_azure_ad_token_refresh" - ], - "secret_manager": [ - "readable" - ] + "http_settings": { + "version": 1, + "fields": { + "ssl_verify": { + "adapter": "SslVerifyInput", + "required": true, + "precedence": "module_global", + "sensitive": false, + "shapes": [ + "none", + "bool", + "str" + ], + "unsupported_live": "configuration_error" + }, + "ssl_certificate": { + "adapter": "OptionalStrictString", + "required": true, + "precedence": "module_global", + "sensitive": false, + "shapes": [], + "unsupported_live": null + }, + "ssl_security_level": { + "adapter": "TuningString", + "required": true, + "precedence": "module_global", + "sensitive": false, + "shapes": [], + "unsupported_live": null + }, + "ssl_ecdh_curve": { + "adapter": "TuningString", + "required": true, + "precedence": "module_global", + "sensitive": false, + "shapes": [], + "unsupported_live": null + }, + "force_ipv4": { + "adapter": "Truthy", + "required": true, + "precedence": "module_global", + "sensitive": false, + "shapes": [], + "unsupported_live": null + }, + "http2": { + "adapter": "ExactTrue", + "required": true, + "precedence": "module_global", + "sensitive": false, + "shapes": [], + "unsupported_live": null + }, + "aiohttp_trust_env": { + "adapter": "Truthy", + "required": true, + "precedence": "module_global", + "sensitive": false, + "shapes": [], + "unsupported_live": null + }, + "disable_aiohttp_trust_env": { + "adapter": "Truthy", + "required": true, + "precedence": "module_global", + "sensitive": false, + "shapes": [], + "unsupported_live": null + }, + "disable_aiohttp_transport": { + "adapter": "ExactTrue", + "required": true, + "precedence": "module_global", + "sensitive": false, + "shapes": [], + "unsupported_live": null + }, + "user_agent": { + "adapter": "StrictString", + "required": true, + "precedence": "accessor", + "sensitive": false, + "shapes": [], + "unsupported_live": null + } + } + }, + "url_policy": { + "version": 1, + "fields": { + "user_url_validation": { + "adapter": "Truthy", + "required": true, + "precedence": "module_global", + "sensitive": false, + "shapes": [], + "unsupported_live": null + }, + "user_url_allowed_hosts": { + "adapter": "HostCollection", + "required": true, + "precedence": "module_global", + "sensitive": false, + "shapes": [], + "unsupported_live": null + } + } + }, + "provider_defaults": { + "version": 1, + "fields": { + "vertex_project": { + "adapter": "FalsyOptionalString", + "required": true, + "precedence": "module_global", + "sensitive": true, + "shapes": [], + "unsupported_live": null + }, + "vertex_location": { + "adapter": "FalsyOptionalString", + "required": true, + "precedence": "module_global", + "sensitive": true, + "shapes": [], + "unsupported_live": null + }, + "enable_azure_ad_token_refresh": { + "adapter": "ExactTrue", + "required": true, + "precedence": "module_global", + "sensitive": false, + "shapes": [], + "unsupported_live": null + } + } + }, + "secret_manager": { + "version": 1, + "fields": { + "readable": { + "adapter": "StrictBool", + "required": true, + "precedence": "accessor", + "sensitive": false, + "shapes": [], + "unsupported_live": null + } + } + } } diff --git a/litellm-rust/crates/python-bridge/src/coercion.rs b/litellm-rust/crates/python-bridge/src/coercion.rs new file mode 100644 index 00000000000..bb5b8b2d454 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/coercion.rs @@ -0,0 +1,231 @@ +use std::collections::BTreeSet; + +use litellm_core_utils::serde_compat::parse_str_bool; +use litellm_http::SslVerify; +use pyo3::{ + exceptions::{PyAttributeError, PyRuntimeError, PyValueError}, + prelude::*, + types::{PyBool, PyString}, +}; + +#[derive(Debug)] +pub(crate) enum ProjectionError { + Python(PyErr), + InvalidConfiguration(String), + UnsupportedLiveObject(String), + InternalSchemaFailure(String), +} + +impl From for ProjectionError { + fn from(error: PyErr) -> Self { + Self::Python(error) + } +} + +impl From for PyErr { + fn from(error: ProjectionError) -> Self { + match error { + ProjectionError::Python(error) => error, + ProjectionError::InvalidConfiguration(message) + | ProjectionError::UnsupportedLiveObject(message) => PyValueError::new_err(message), + ProjectionError::InternalSchemaFailure(message) => PyRuntimeError::new_err(message), + } + } +} + +pub(crate) struct Truthy(pub bool); +pub(crate) struct ExactTrue(pub bool); +pub(crate) struct StrBool(pub Option); +pub(crate) struct OptionalStrictString(pub Option); +pub(crate) struct FalsyOptionalString(pub Option); +pub(crate) struct TuningString(pub Option); +pub(crate) struct StringCollection(pub Vec); +pub(crate) struct SslVerifyInput(pub Option); + +pub(crate) struct Field<'py> { + path: &'static str, + value: Bound<'py, PyAny>, +} + +impl<'py> Field<'py> { + pub(crate) fn new(path: &'static str, value: Bound<'py, PyAny>) -> Self { + Self { path, value } + } + + pub(crate) fn read( + snapshot: &Bound<'py, PyAny>, + path: &'static str, + ) -> Result { + let name = path.rsplit('.').next().unwrap_or(path); + match snapshot.getattr(name) { + Ok(value) => Ok(Self::new(path, value)), + Err(error) if error.is_instance_of::(snapshot.py()) => { + match Self::missing_field(snapshot, name) { + Ok(true) => Err(ProjectionError::InternalSchemaFailure(format!( + "{path}: missing snapshot field" + ))), + _ => Err(error.into()), + } + } + Err(error) => Err(error.into()), + } + } + + fn missing_field(snapshot: &Bound<'_, PyAny>, name: &str) -> PyResult { + let py = snapshot.py(); + let object = py.import("builtins")?.getattr("object")?; + let missing = object.call0()?; + let lookup = py.import("inspect")?.getattr("getattr_static")?; + let declared = lookup.call1((snapshot, name, &missing))?; + let fallback = lookup.call1((snapshot.get_type(), "__getattr__", &missing))?; + let getter = lookup.call1((snapshot.get_type(), "__getattribute__"))?; + Ok(declared.is(&missing) + && fallback.is(&missing) + && getter.is(object.getattr("__getattribute__")?)) + } + + fn expected(&self, expected: &'static str) -> Result { + Ok(format!( + "{}: expected {expected}, got {}", + self.path, + self.value.get_type().name()? + )) + } + + fn invalid(&self, expected: &'static str) -> ProjectionError { + match self.expected(expected) { + Ok(message) => ProjectionError::InvalidConfiguration(message), + Err(error) => error, + } + } + + pub(crate) fn truthy(&self) -> Result { + Ok(Truthy(self.value.is_truthy()?)) + } + + pub(crate) fn exact_true(&self) -> ExactTrue { + ExactTrue(self.value.is(PyBool::new(self.value.py(), true))) + } + + pub(crate) fn strict_string(&self) -> Result { + let value = self + .value + .cast::() + .map_err(|_| self.invalid("a string"))?; + Ok(value.to_str()?.to_owned()) + } + + pub(crate) fn schema_string(&self) -> Result { + if !self.value.is_instance_of::() { + return Err(ProjectionError::InternalSchemaFailure( + self.expected("a string")?, + )); + } + self.strict_string() + } + + pub(crate) fn schema_bool(&self) -> Result { + if !self.value.is_instance_of::() { + return Err(ProjectionError::InternalSchemaFailure( + self.expected("a Boolean")?, + )); + } + Ok(self.exact_true().0) + } + + pub(crate) fn str_bool(&self) -> Result { + if self.value.is_none() { + return Ok(StrBool(None)); + } + Ok(StrBool(parse_str_bool(&self.strict_string()?))) + } + + pub(crate) fn optional_strict_string(&self) -> Result { + if self.value.is_none() { + return Ok(OptionalStrictString(None)); + } + self.strict_string().map(Some).map(OptionalStrictString) + } + + pub(crate) fn falsy_optional_string(&self) -> Result { + if !self.truthy()?.0 { + return Ok(FalsyOptionalString(None)); + } + self.strict_string().map(Some).map(FalsyOptionalString) + } + + pub(crate) fn tuning_string(&self) -> Result { + if !self.truthy()?.0 || !self.value.is_instance_of::() { + return Ok(TuningString(None)); + } + self.strict_string().map(Some).map(TuningString) + } + + pub(crate) fn string_collection(&self) -> Result { + if !self.truthy()?.0 { + return Ok(StringCollection(Vec::new())); + } + if self.value.is_instance_of::() { + return self + .strict_string() + .map(|value| StringCollection(vec![value])); + } + let values = self + .value + .try_iter()? + .filter_map(|item| { + let member = match item { + Ok(value) => Self::new(self.path, value), + Err(error) => return Some(Err(error.into())), + }; + match member.truthy() { + Ok(Truthy(false)) => None, + Ok(Truthy(true)) => Some(member.strict_string()), + Err(error) => Some(Err(error)), + } + }) + .collect::, ProjectionError>>()?; + Ok(StringCollection(values)) + } + + pub(crate) fn host_collection(&self) -> Result { + let values = self + .string_collection()? + .0 + .into_iter() + .map(|host| litellm_http::media::normalize_host(&host)) + .collect::>(); + Ok(StringCollection(values.into_iter().collect())) + } + + pub(crate) fn ssl_verify(&self) -> Result { + if self.value.is_none() { + return Ok(SslVerifyInput(None)); + } + if self.value.is_instance_of::() { + return Ok(SslVerifyInput(Some(if self.exact_true().0 { + SslVerify::Enabled + } else { + SslVerify::Disabled + }))); + } + if self.value.is_instance_of::() { + let parsed = match self.str_bool()?.0 { + Some(true) => SslVerify::Enabled, + Some(false) => SslVerify::Disabled, + None => SslVerify::CaBundle(self.strict_string()?.into()), + }; + return Ok(SslVerifyInput(Some(parsed))); + } + let context = self.value.py().import("ssl")?.getattr("SSLContext")?; + if self.value.is_instance(&context)? { + return Err(ProjectionError::UnsupportedLiveObject(self.expected( + "a Boolean, Boolean string, CA path, or None; live SSLContext is unsupported", + )?)); + } + Err(self.invalid("a Boolean, Boolean string, CA path, or None")) + } +} + +#[cfg(test)] +mod tests; diff --git a/litellm-rust/crates/python-bridge/src/coercion/tests.rs b/litellm-rust/crates/python-bridge/src/coercion/tests.rs new file mode 100644 index 00000000000..5ed237c3c64 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/coercion/tests.rs @@ -0,0 +1,372 @@ +use std::ffi::CString; + +use pyo3::{ + exceptions::{PyLookupError, PyRuntimeError, PyValueError}, + types::PyDict, +}; +use rstest::rstest; + +use super::*; + +fn evaluate<'py>(py: Python<'py>, source: &str) -> Bound<'py, PyAny> { + py.eval(&CString::new(source).unwrap(), None, None).unwrap() +} + +#[rstest] +#[case("None", false, false)] +#[case("False", false, false)] +#[case("True", true, true)] +#[case("0", false, false)] +#[case("1", true, false)] +#[case("''", false, false)] +#[case("'false'", true, false)] +#[case("[]", false, false)] +#[case("[0]", true, false)] +#[case("{}", false, false)] +#[case("object()", true, false)] +fn boolean_operations_have_distinct_python_semantics( + #[case] source: &str, + #[case] truth: bool, + #[case] exact: bool, +) { + Python::initialize(); + Python::attach(|py| { + let value = evaluate(py, source); + let field = Field::new("test.flag", value.clone()); + assert_eq!(field.truthy().unwrap().0, truth); + assert_eq!(field.exact_true().0, exact); + assert_eq!( + field.truthy().unwrap().0, + py.import("builtins") + .unwrap() + .getattr("bool") + .unwrap() + .call1((value,)) + .unwrap() + .extract::() + .unwrap() + ); + }); +} + +#[rstest] +#[case("None", Ok(None), Ok(None), Ok(None))] +#[case("''", Ok(Some("")), Ok(None), Ok(None))] +#[case( + "' value '", + Ok(Some(" value ")), + Ok(Some(" value ")), + Ok(Some(" value ")) +)] +#[case("[]", Err(()), Ok(None), Ok(None))] +#[case("0", Err(()), Ok(None), Ok(None))] +#[case("1", Err(()), Err(()), Ok(None))] +#[case("object()", Err(()), Err(()), Ok(None))] +fn string_operations_do_not_conflate_absence_and_type_checks( + #[case] source: &str, + #[case] strict: Result, ()>, + #[case] fallback: Result, ()>, + #[case] tuning: Result, ()>, +) { + Python::initialize(); + Python::attach(|py| { + let field = Field::new("test.string", evaluate(py, source)); + let owned = + |expected: Result, ()>| expected.map(|value| value.map(str::to_owned)); + assert_eq!( + field + .optional_strict_string() + .map(|value| value.0) + .map_err(|_| ()), + owned(strict) + ); + assert_eq!( + field + .falsy_optional_string() + .map(|value| value.0) + .map_err(|_| ()), + owned(fallback) + ); + assert_eq!( + field.tuning_string().map(|value| value.0).map_err(|_| ()), + owned(tuning) + ); + }); +} + +#[rstest] +#[case("None", None)] +#[case("' True '", Some(true))] +#[case("' fAlSe '", Some(false))] +#[case("'yes'", None)] +#[case("'1'", None)] +#[case("'unknown'", None)] +fn string_boolean_tokens_remain_separate_from_truthiness( + #[case] source: &str, + #[case] expected: Option, +) { + Python::initialize(); + Python::attach(|py| { + assert_eq!( + Field::new("test.flag", evaluate(py, source)) + .str_bool() + .unwrap() + .0, + expected + ); + }); +} + +#[rstest] +#[case("'EXAMPLE.TEST.'", vec!["example.test"])] +#[case("['B.test', '', None, 0, [], 'A.test.', 'b.test']", vec!["a.test", "b.test"])] +#[case("('B.test', 'a.test')", vec!["a.test", "b.test"])] +#[case("{'B.test', 'a.test'}", vec!["a.test", "b.test"])] +#[case("(host for host in ['B.test', 'a.test'])", vec!["a.test", "b.test"])] +#[case("None", vec![])] +#[case("False", vec![])] +fn host_collection_is_owned_normalized_and_deterministic( + #[case] source: &str, + #[case] expected: Vec<&str>, +) { + Python::initialize(); + Python::attach(|py| { + assert_eq!( + Field::new("url_policy.user_url_allowed_hosts", evaluate(py, source)) + .host_collection() + .unwrap() + .0, + expected + ); + }); +} + +#[test] +fn protocol_errors_preserve_exception_identity_traceback_cause_and_context() { + Python::initialize(); + Python::attach(|py| { + let locals = PyDict::new(py); + py.run( + c" +failure = LookupError('protocol failed') +cause = ValueError('cause') +context = RuntimeError('context') +def fail(): + try: + raise context + except RuntimeError: + raise failure from cause +class Bool: + def __bool__(self): return fail() +class Length: + def __len__(self): return fail() +class Iter: + def __iter__(self): return fail() +class Next: + def __iter__(self): return self + def __next__(self): return fail() +class Descriptor: + @property + def flag(self): return fail() +values = (Bool(), Length(), Iter(), Next(), [Bool()]) +descriptor = Descriptor() +", + Some(&locals), + Some(&locals), + ) + .unwrap(); + let values = locals.get_item("values").unwrap().unwrap(); + for value in values.try_iter().unwrap() { + let error = Field::new("test.flag", value.unwrap()) + .host_collection() + .err() + .unwrap(); + let error = PyErr::from(error); + assert!( + error + .value(py) + .is(locals.get_item("failure").unwrap().unwrap()) + ); + assert!(error.is_instance_of::(py)); + assert!(error.traceback(py).is_some()); + assert!( + error + .value(py) + .getattr("__cause__") + .unwrap() + .is(locals.get_item("cause").unwrap().unwrap()) + ); + assert!( + error + .value(py) + .getattr("__context__") + .unwrap() + .is(locals.get_item("context").unwrap().unwrap()) + ); + } + let error = Field::read( + &locals.get_item("descriptor").unwrap().unwrap(), + "test.flag", + ) + .err() + .unwrap(); + assert!( + PyErr::from(error) + .value(py) + .is(locals.get_item("failure").unwrap().unwrap()) + ); + }); +} + +#[test] +fn identity_and_string_contents_do_not_invoke_unrelated_protocols() { + Python::initialize(); + Python::attach(|py| { + let locals = PyDict::new(py); + py.run( + c" +class Hostile: + def __bool__(self): raise AssertionError('bool called') + def __eq__(self, other): raise AssertionError('eq called') + def __str__(self): raise AssertionError('str called') +class Text(str): + def __str__(self): raise AssertionError('str called') + def strip(self): raise AssertionError('strip called') + def lower(self): raise AssertionError('lower called') +hostile = Hostile() +text = Text(' False ') +", + Some(&locals), + Some(&locals), + ) + .unwrap(); + let hostile = Field::new("test.flag", locals.get_item("hostile").unwrap().unwrap()); + assert!(!hostile.exact_true().0); + assert!(matches!( + hostile.strict_string(), + Err(ProjectionError::InvalidConfiguration(_)) + )); + let text = Field::new("test.flag", locals.get_item("text").unwrap().unwrap()); + assert_eq!(text.strict_string().unwrap(), " False "); + assert_eq!(text.str_bool().unwrap().0, Some(false)); + }); +} + +#[test] +fn missing_snapshot_fields_and_descriptor_attribute_errors_are_distinct() { + Python::initialize(); + Python::attach(|py| { + let locals = PyDict::new(py); + py.run( + c" +failure = AttributeError('descriptor failed') +class Snapshot: + @property + def flag(self): raise failure +snapshot = Snapshot() +class Dynamic: + def __getattr__(self, name): raise failure +class Intercepted: + def __getattribute__(self, name): raise failure +dynamic = Dynamic() +intercepted = Intercepted() +", + Some(&locals), + Some(&locals), + ) + .unwrap(); + let snapshot = locals.get_item("snapshot").unwrap().unwrap(); + let descriptor = PyErr::from(Field::read(&snapshot, "test.flag").err().unwrap()); + assert!( + descriptor + .value(py) + .is(locals.get_item("failure").unwrap().unwrap()) + ); + for name in ["dynamic", "intercepted"] { + let value = locals.get_item(name).unwrap().unwrap(); + let error = PyErr::from(Field::read(&value, "test.flag").err().unwrap()); + assert!( + error + .value(py) + .is(locals.get_item("failure").unwrap().unwrap()) + ); + } + let missing = PyErr::from(Field::read(&snapshot, "test.missing").err().unwrap()); + assert!(missing.is_instance_of::(py)); + assert!(missing.to_string().contains("test.missing")); + }); +} + +#[test] +fn configuration_errors_name_fields_without_exposing_values() { + Python::initialize(); + Python::attach(|py| { + for source in [ + "{'secret': 'do-not-print'}", + "['host.test', {'secret': 'do-not-print'}]", + ] { + let field = Field::new("test.setting", evaluate(py, source)); + let error = PyErr::from(field.falsy_optional_string().err().unwrap()); + assert!(error.is_instance_of::(py)); + assert!(error.to_string().contains("test.setting")); + assert!(!error.to_string().contains("do-not-print")); + } + let hosts = Field::new( + "url_policy.user_url_allowed_hosts", + evaluate(py, "['host.test', 1]"), + ); + assert!(matches!( + hosts.host_collection(), + Err(ProjectionError::InvalidConfiguration(_)) + )); + assert!(matches!( + Field::new("test.flag", evaluate(py, "1")).str_bool(), + Err(ProjectionError::InvalidConfiguration(_)) + )); + }); +} + +#[test] +fn projection_releases_the_source_collection() { + Python::initialize(); + Python::attach(|py| { + let source = evaluate(py, "['A.test']"); + let projected = Field::new("test.hosts", source.clone()) + .host_collection() + .unwrap() + .0; + source.call_method1("append", ("b.test",)).unwrap(); + assert_eq!(projected, ["a.test"]); + assert_eq!( + Field::new("test.hosts", source) + .host_collection() + .unwrap() + .0, + ["a.test", "b.test"] + ); + }); +} + +#[rstest] +#[case("True", Some(true))] +#[case("False", Some(false))] +#[case("1", None)] +#[case("None", None)] +#[case("[]", None)] +fn accessor_booleans_are_strict_schema_values( + #[case] source: &str, + #[case] expected: Option, +) { + Python::initialize(); + Python::attach(|py| { + let result = Field::new("secret_manager.readable", evaluate(py, source)).schema_bool(); + match expected { + Some(expected) => assert_eq!(result.unwrap(), expected), + None => { + let error = PyErr::from(result.unwrap_err()); + assert!(error.is_instance_of::(py)); + assert!(error.to_string().contains("secret_manager.readable")); + } + } + }); +} diff --git a/litellm-rust/crates/python-bridge/src/http.rs b/litellm-rust/crates/python-bridge/src/http.rs index 7e9a5f093b4..859d579129f 100644 --- a/litellm-rust/crates/python-bridge/src/http.rs +++ b/litellm-rust/crates/python-bridge/src/http.rs @@ -10,9 +10,9 @@ use litellm_http::{ Unsupported, media::{PublicDnsResolver, UrlPolicy}, }; -use pyo3::{prelude::*, types::PyDict}; +use pyo3::{exceptions::PyValueError, prelude::*, types::PyDict}; -use crate::{errors::RustBridgeDeclined, python_settings::PythonSettings}; +use crate::{coercion::Field, python_settings::PythonSettings}; static POOL: LazyLock = LazyLock::new(|| HttpClientPool::new(Arc::new(PublicDnsResolver))); @@ -41,6 +41,22 @@ pub(crate) fn call_config( Ok(resolution.config) } +pub(crate) fn client_error(error: litellm_http::Error, config: &HttpClientConfig) -> PyErr { + match error { + litellm_http::Error::Read { path, .. } | litellm_http::Error::InvalidPem { path, .. } + if config.client_certificate.as_ref() == Some(&path) => + { + PyValueError::new_err( + "http_settings.ssl_certificate: expected a readable PEM certificate and private key", + ) + } + litellm_http::Error::Read { .. } | litellm_http::Error::InvalidPem { .. } => { + PyValueError::new_err("http_settings.ssl_verify: expected a readable PEM CA bundle") + } + _ => PyValueError::new_err("http_settings: native HTTP client configuration is invalid"), + } +} + fn unreported( reported: &Mutex>, unsupported: Vec, @@ -53,25 +69,25 @@ fn unreported( } pub(crate) fn url_policy(py: Python<'_>) -> PyResult { - let policy: PythonUrlPolicy = - PythonSettings::UrlPolicy - .read(py)? - .extract() - .map_err(|error: PyErr| { - RustBridgeDeclined::new_err(format!( - "litellm URL policy cannot be used by the Rust route: {error}" - )) - })?; + project_url_policy(&PythonSettings::UrlPolicy.read(py)?) +} + +fn project_url_policy(value: &Bound<'_, PyAny>) -> PyResult { Ok(UrlPolicy { - validate: policy.user_url_validation, - allowed_hosts: policy.user_url_allowed_hosts, + validate: Field::read(value, "url_policy.user_url_validation")? + .truthy()? + .0, + allowed_hosts: Field::read(value, "url_policy.user_url_allowed_hosts")? + .host_collection()? + .0, }) } fn call_ssl_verify(kwargs: &Bound<'_, PyDict>) -> PyResult> { - Ok(kwargs - .get_item("ssl_verify")? - .and_then(|value| ssl_verify(&value))) + match kwargs.get_item("ssl_verify")? { + Some(value) => Ok(Field::new("request.ssl_verify", value).ssl_verify()?.0), + None => Ok(None), + } } fn for_call(call_ssl_verify: Option, asynchronous: bool) -> HttpSettingsLayer { @@ -82,64 +98,47 @@ fn for_call(call_ssl_verify: Option, asynchronous: bool) -> HttpSetti } } -#[derive(FromPyObject)] -struct PythonUrlPolicy { - user_url_validation: bool, - user_url_allowed_hosts: Vec, -} - -#[derive(FromPyObject)] -struct PythonHttpSettings<'py> { - ssl_verify: Bound<'py, PyAny>, - ssl_certificate: Option, - ssl_security_level: Option, - ssl_ecdh_curve: Option, - force_ipv4: bool, - http2: bool, - aiohttp_trust_env: bool, - disable_aiohttp_trust_env: bool, - disable_aiohttp_transport: bool, - user_agent: String, -} - fn configured(value: &Bound<'_, PyAny>) -> PyResult { - let python: PythonHttpSettings = value.extract().map_err(|error: PyErr| { - RustBridgeDeclined::new_err(format!( - "litellm HTTP settings cannot be used by the Rust route: {error}" - )) - })?; Ok(HttpSettingsLayer { - ssl_verify: ssl_verify(&python.ssl_verify), - ssl_certificate: python.ssl_certificate.map(PathBuf::from), - ssl_security_level: python.ssl_security_level, - ssl_ecdh_curve: python.ssl_ecdh_curve, - force_ipv4: Some(python.force_ipv4), - http2: Some(python.http2), - aiohttp_trust_env: Some(python.aiohttp_trust_env), - disable_aiohttp_trust_env: Some(python.disable_aiohttp_trust_env), - disable_aiohttp_transport: Some(python.disable_aiohttp_transport), - user_agent: Some(python.user_agent), + ssl_verify: Field::read(value, "http_settings.ssl_verify")? + .ssl_verify()? + .0, + ssl_certificate: Field::read(value, "http_settings.ssl_certificate")? + .optional_strict_string()? + .0 + .map(PathBuf::from), + ssl_security_level: Field::read(value, "http_settings.ssl_security_level")? + .tuning_string()? + .0, + ssl_ecdh_curve: Field::read(value, "http_settings.ssl_ecdh_curve")? + .tuning_string()? + .0, + force_ipv4: Some(Field::read(value, "http_settings.force_ipv4")?.truthy()?.0), + http2: Some(Field::read(value, "http_settings.http2")?.exact_true().0), + aiohttp_trust_env: Some( + Field::read(value, "http_settings.aiohttp_trust_env")? + .truthy()? + .0, + ), + disable_aiohttp_trust_env: Some( + Field::read(value, "http_settings.disable_aiohttp_trust_env")? + .truthy()? + .0, + ), + disable_aiohttp_transport: Some( + Field::read(value, "http_settings.disable_aiohttp_transport")? + .exact_true() + .0, + ), + user_agent: Some(Field::read(value, "http_settings.user_agent")?.schema_string()?), ..HttpSettingsLayer::default() }) } -fn ssl_verify(value: &Bound<'_, PyAny>) -> Option { - if let Ok(enabled) = value.extract::() { - return Some(if enabled { - SslVerify::Enabled - } else { - SslVerify::Disabled - }); - } - value - .extract::() - .ok() - .map(|path| SslVerify::parse(&path)) -} - #[cfg(test)] mod tests { use litellm_http::Verify; + use pyo3::exceptions::PyRuntimeError; use rstest::rstest; use super::*; @@ -163,7 +162,7 @@ defaults = dict( user_agent='litellm/test', ) defaults.update(dict({overrides})) -settings = types.SimpleNamespace(**{{name: defaults[name] for name in json.loads(contract)['http_settings']}}) +settings = types.SimpleNamespace(**{{name: defaults[name] for name in json.loads(contract)['http_settings']['fields']}}) " ); let locals = PyDict::new(py); @@ -259,12 +258,16 @@ user_agent='litellm/9.9.9', }); } - #[test] - fn ssl_context_global_is_ignored_so_environment_and_defaults_apply() { + #[rstest] + #[case("ssl_verify=object()")] + #[case("ssl_verify=__import__('ssl').SSLContext(__import__('ssl').PROTOCOL_TLS_CLIENT)")] + #[case("ssl_certificate=1")] + fn invalid_http_configuration_is_terminal(#[case] overrides: &str) { Python::initialize(); Python::attach(|py| { - let layer = configured(&python_settings(py, "ssl_verify=object()")).unwrap(); - assert_eq!(layer.ssl_verify, None); + let error = configured(&python_settings(py, overrides)).unwrap_err(); + assert!(error.is_instance_of::(py)); + assert!(error.to_string().contains("http_settings.ssl_")); }); } @@ -281,11 +284,21 @@ user_agent='litellm/9.9.9', } #[test] - fn mistyped_python_settings_decline_instead_of_raising() { + fn mutable_globals_use_their_consumer_operations() { Python::initialize(); Python::attach(|py| { - let error = configured(&python_settings(py, "force_ipv4='yes'")).unwrap_err(); - assert!(error.is_instance_of::(py)); + let layer = configured(&python_settings(py, + "force_ipv4='yes', http2=1, disable_aiohttp_transport=1, aiohttp_trust_env=[1], disable_aiohttp_trust_env=[], ssl_security_level=1, ssl_ecdh_curve=[]" + )).unwrap(); + assert_eq!(layer.force_ipv4, Some(true)); + assert_eq!(layer.http2, Some(false)); + assert_eq!(layer.disable_aiohttp_transport, Some(false)); + assert_eq!(layer.aiohttp_trust_env, Some(true)); + assert_eq!(layer.disable_aiohttp_trust_env, Some(false)); + assert_eq!(layer.ssl_security_level, None); + assert_eq!(layer.ssl_ecdh_curve, None); + let error = configured(&python_settings(py, "user_agent=1")).unwrap_err(); + assert!(error.is_instance_of::(py)); }); } @@ -323,17 +336,36 @@ user_agent='litellm/9.9.9', } #[test] - fn live_ssl_context_argument_is_ignored_so_the_configured_value_applies() { + fn live_ssl_context_argument_raises_instead_of_using_another_layer() { Python::initialize(); Python::attach(|py| { let kwargs = PyDict::new(py); - kwargs - .set_item("ssl_verify", py.eval(c"object()", None, None).unwrap()) + let ssl = py.import("ssl").unwrap(); + let context = ssl + .getattr("SSLContext") + .unwrap() + .call1((ssl.getattr("PROTOCOL_TLS_CLIENT").unwrap(),)) .unwrap(); - let call = for_call(call_ssl_verify(&kwargs).unwrap(), true); - let settings = - HttpSettings::from_layers([call, configured_ssl_verify(SslVerify::Disabled)]); - assert_eq!(settings.ssl_verify, Some(SslVerify::Disabled)); + kwargs.set_item("ssl_verify", context).unwrap(); + let error = call_ssl_verify(&kwargs).unwrap_err(); + assert!(error.is_instance_of::(py)); + assert!(error.to_string().contains("request.ssl_verify")); + assert!(error.to_string().contains("SSLContext")); + }); + } + + #[test] + fn url_policy_uses_truthiness_and_normalized_owned_hosts() { + Python::initialize(); + Python::attach(|py| { + let value = py.eval(c"__import__('types').SimpleNamespace(user_url_validation=[], user_url_allowed_hosts=['B.test', 'a.test.', 'b.test'])", None, None).unwrap(); + assert_eq!( + project_url_policy(&value).unwrap(), + UrlPolicy { + validate: false, + allowed_hosts: vec!["a.test".into(), "b.test".into()], + } + ); }); } diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index bd62c5aadf1..f13a3ad433f 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -1,4 +1,5 @@ mod cache; +mod coercion; mod credentials; mod diagnostics; mod errors; diff --git a/litellm-rust/crates/python-bridge/src/python_settings.rs b/litellm-rust/crates/python-bridge/src/python_settings.rs index 7ac23a05542..bdc6d14356d 100644 --- a/litellm-rust/crates/python-bridge/src/python_settings.rs +++ b/litellm-rust/crates/python-bridge/src/python_settings.rs @@ -43,32 +43,204 @@ pub(crate) const CONTRACT: &str = include_str!("../python_settings.json"); #[cfg(test)] mod tests { - use std::{collections::BTreeSet, ffi::CString}; - - use pyo3::{prelude::*, types::PyDict}; - use super::{CONTRACT, PythonSettings}; + use pyo3::prelude::*; + use serde_json::{Value, json}; + + struct SettingSpec { + group: &'static str, + name: &'static str, + adapter: &'static str, + precedence: &'static str, + sensitive: bool, + shapes: &'static [&'static str], + unsupported_live: Option<&'static str>, + } + + const SETTINGS: &[SettingSpec] = &[ + SettingSpec { + group: "http_settings", + name: "ssl_verify", + adapter: "SslVerifyInput", + precedence: "module_global", + sensitive: false, + shapes: &["none", "bool", "str"], + unsupported_live: Some("configuration_error"), + }, + SettingSpec { + group: "http_settings", + name: "ssl_certificate", + adapter: "OptionalStrictString", + precedence: "module_global", + sensitive: false, + shapes: &[], + unsupported_live: None, + }, + SettingSpec { + group: "http_settings", + name: "ssl_security_level", + adapter: "TuningString", + precedence: "module_global", + sensitive: false, + shapes: &[], + unsupported_live: None, + }, + SettingSpec { + group: "http_settings", + name: "ssl_ecdh_curve", + adapter: "TuningString", + precedence: "module_global", + sensitive: false, + shapes: &[], + unsupported_live: None, + }, + SettingSpec { + group: "http_settings", + name: "force_ipv4", + adapter: "Truthy", + precedence: "module_global", + sensitive: false, + shapes: &[], + unsupported_live: None, + }, + SettingSpec { + group: "http_settings", + name: "http2", + adapter: "ExactTrue", + precedence: "module_global", + sensitive: false, + shapes: &[], + unsupported_live: None, + }, + SettingSpec { + group: "http_settings", + name: "aiohttp_trust_env", + adapter: "Truthy", + precedence: "module_global", + sensitive: false, + shapes: &[], + unsupported_live: None, + }, + SettingSpec { + group: "http_settings", + name: "disable_aiohttp_trust_env", + adapter: "Truthy", + precedence: "module_global", + sensitive: false, + shapes: &[], + unsupported_live: None, + }, + SettingSpec { + group: "http_settings", + name: "disable_aiohttp_transport", + adapter: "ExactTrue", + precedence: "module_global", + sensitive: false, + shapes: &[], + unsupported_live: None, + }, + SettingSpec { + group: "http_settings", + name: "user_agent", + adapter: "StrictString", + precedence: "accessor", + sensitive: false, + shapes: &[], + unsupported_live: None, + }, + SettingSpec { + group: "url_policy", + name: "user_url_validation", + adapter: "Truthy", + precedence: "module_global", + sensitive: false, + shapes: &[], + unsupported_live: None, + }, + SettingSpec { + group: "url_policy", + name: "user_url_allowed_hosts", + adapter: "HostCollection", + precedence: "module_global", + sensitive: false, + shapes: &[], + unsupported_live: None, + }, + SettingSpec { + group: "provider_defaults", + name: "vertex_project", + adapter: "FalsyOptionalString", + precedence: "module_global", + sensitive: true, + shapes: &[], + unsupported_live: None, + }, + SettingSpec { + group: "provider_defaults", + name: "vertex_location", + adapter: "FalsyOptionalString", + precedence: "module_global", + sensitive: true, + shapes: &[], + unsupported_live: None, + }, + SettingSpec { + group: "provider_defaults", + name: "enable_azure_ad_token_refresh", + adapter: "ExactTrue", + precedence: "module_global", + sensitive: false, + shapes: &[], + unsupported_live: None, + }, + SettingSpec { + group: "secret_manager", + name: "readable", + adapter: "StrictBool", + precedence: "accessor", + sensitive: false, + shapes: &[], + unsupported_live: None, + }, + ]; #[test] - fn every_settings_group_is_in_the_python_contract() { - Python::initialize(); - Python::attach(|py| { - let locals = PyDict::new(py); - locals.set_item("contract", CONTRACT).unwrap(); - let source = CString::new("import json\nkeys = list(json.loads(contract))").unwrap(); - py.run(&source, Some(&locals), Some(&locals)).unwrap(); - let declared: BTreeSet = locals - .get_item("keys") + fn settings_manifest_matches_the_semantic_contract() { + pyo3::Python::initialize(); + let manifest: Value = pyo3::Python::attach(|py| { + let value = py + .import("json") .unwrap() - .unwrap() - .extract::>() - .unwrap() - .into_iter() - .collect(); - let read: BTreeSet = PythonSettings::ALL - .map(|group| group.name().to_owned()) - .into(); - assert_eq!(read, declared); + .call_method1("loads", (CONTRACT,)) + .unwrap(); + litellm_host_python::from_py(&value).unwrap() }); + let expected: serde_json::Map = PythonSettings::ALL + .into_iter() + .map(|group| { + let fields: serde_json::Map = SETTINGS + .iter() + .filter(|spec| spec.group == group.name()) + .map(|spec| { + ( + spec.name.to_owned(), + json!({ + "adapter": spec.adapter, + "required": true, + "precedence": spec.precedence, + "sensitive": spec.sensitive, + "shapes": spec.shapes, + "unsupported_live": spec.unsupported_live, + }), + ) + }) + .collect(); + ( + group.name().to_owned(), + json!({"version": 1, "fields": fields}), + ) + }) + .collect(); + assert_eq!(manifest, Value::Object(expected)); } } diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs index e518f972bac..ce6f04c321b 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs @@ -19,7 +19,7 @@ use pyo3::{ types::{PyDict, PyTuple}, }; -use crate::{errors::RustBridgeDeclined, http, python_settings::PythonSettings}; +use crate::{coercion::Field, errors::RustBridgeDeclined, http, python_settings::PythonSettings}; const SURFACE: LegacySurface = LegacySurface { call_type: "ocr", @@ -51,7 +51,7 @@ fn run_ocr( ocr_settings(py)?, secrets, ) - .map_err(|error| RustBridgeDeclined::new_err(error.to_string()))?; + .map_err(|error| http::client_error(error, &config))?; run_legacy_call( py, if asynchronous { ASYNC_SURFACE } else { SURFACE }, @@ -62,14 +62,8 @@ fn run_ocr( ) } -#[derive(FromPyObject)] -struct PythonSecretManager { - readable: bool, -} - fn process_environment_secrets(secret_manager: &Bound<'_, PyAny>) -> PyResult { - let manager: PythonSecretManager = secret_manager.extract()?; - if manager.readable { + if Field::read(secret_manager, "secret_manager.readable")?.schema_bool()? { return Err(RustBridgeDeclined::new_err( "a readable secret manager is configured and the Rust route only reads the process environment", )); @@ -77,26 +71,24 @@ fn process_environment_secrets(secret_manager: &Bound<'_, PyAny>) -> PyResult, - vertex_location: Option, - enable_azure_ad_token_refresh: Option, +fn ocr_settings(py: Python<'_>) -> PyResult { + project_provider_defaults(&PythonSettings::ProviderDefaults.read(py)?) } -fn ocr_settings(py: Python<'_>) -> PyResult { - let defaults: PythonProviderDefaults = PythonSettings::ProviderDefaults - .read(py)? - .extract() - .map_err(|error: PyErr| { - RustBridgeDeclined::new_err(format!( - "litellm provider defaults cannot be used by the Rust route: {error}" - )) - })?; +fn project_provider_defaults(value: &Bound<'_, PyAny>) -> PyResult { Ok(OcrSettings { - vertex_project: defaults.vertex_project, - vertex_location: defaults.vertex_location, - enable_azure_ad_token_refresh: defaults.enable_azure_ad_token_refresh == Some(true), + vertex_project: Field::read(value, "provider_defaults.vertex_project")? + .falsy_optional_string()? + .0, + vertex_location: Field::read(value, "provider_defaults.vertex_location")? + .falsy_optional_string()? + .0, + enable_azure_ad_token_refresh: Field::read( + value, + "provider_defaults.enable_azure_ad_token_refresh", + )? + .exact_true() + .0, ..OcrSettings::from_environment(&ProcessEnvironment) }) } @@ -140,6 +132,35 @@ mod tests { locals.get_item("manager").unwrap().unwrap() } + #[test] + fn provider_defaults_distinguish_falsey_values_and_exact_true() { + Python::initialize(); + Python::attach(|py| { + let value = py.eval(c"__import__('types').SimpleNamespace(vertex_project=[], vertex_location=0, enable_azure_ad_token_refresh=1)", None, None).unwrap(); + let projected = super::project_provider_defaults(&value).unwrap(); + assert_eq!(projected.vertex_project, None); + assert_eq!(projected.vertex_location, None); + assert!(!projected.enable_azure_ad_token_refresh); + value.setattr("vertex_project", "project").unwrap(); + value.setattr("vertex_location", "region").unwrap(); + value + .setattr("enable_azure_ad_token_refresh", true) + .unwrap(); + let next = super::project_provider_defaults(&value).unwrap(); + assert_eq!(next.vertex_project.as_deref(), Some("project")); + assert_eq!(next.vertex_location.as_deref(), Some("region")); + assert!(next.enable_azure_ad_token_refresh); + value.setattr("vertex_project", 1).unwrap(); + let error = super::project_provider_defaults(&value).err().unwrap(); + assert!(error.is_instance_of::(py)); + assert!( + error + .to_string() + .contains("provider_defaults.vertex_project") + ); + }); + } + #[test] fn a_readable_secret_manager_sends_the_call_back_to_python() { Python::initialize(); diff --git a/litellm/rust_bridge/settings.py b/litellm/rust_bridge/settings.py index 3aa2d742862..9a5cf49f298 100644 --- a/litellm/rust_bridge/settings.py +++ b/litellm/rust_bridge/settings.py @@ -1,34 +1,33 @@ from __future__ import annotations -from collections.abc import Sequence from dataclasses import dataclass @dataclass(frozen=True, slots=True) class HttpSettings: - ssl_verify: bool | str - ssl_certificate: str | None - ssl_security_level: str | None - ssl_ecdh_curve: str | None - force_ipv4: bool - http2: bool - aiohttp_trust_env: bool - disable_aiohttp_trust_env: bool - disable_aiohttp_transport: bool + ssl_verify: object + ssl_certificate: object + ssl_security_level: object + ssl_ecdh_curve: object + force_ipv4: object + http2: object + aiohttp_trust_env: object + disable_aiohttp_trust_env: object + disable_aiohttp_transport: object user_agent: str @dataclass(frozen=True, slots=True) class UrlPolicy: - user_url_validation: bool - user_url_allowed_hosts: Sequence[str] + user_url_validation: object + user_url_allowed_hosts: object @dataclass(frozen=True, slots=True) class ProviderDefaults: - vertex_project: str | None - vertex_location: str | None - enable_azure_ad_token_refresh: bool | None + vertex_project: object + vertex_location: object + enable_azure_ad_token_refresh: object @dataclass(frozen=True, slots=True) diff --git a/tests/test_litellm/rust_bridge/test_settings.py b/tests/test_litellm/rust_bridge/test_settings.py index 6b78ddad44b..023f02cffbb 100644 --- a/tests/test_litellm/rust_bridge/test_settings.py +++ b/tests/test_litellm/rust_bridge/test_settings.py @@ -6,6 +6,7 @@ from typing import Final import httpx import pytest from pydantic import TypeAdapter +from typing_extensions import ReadOnly, TypedDict import litellm from litellm.integrations.custom_secret_manager import CustomSecretManager @@ -17,14 +18,28 @@ from litellm.types.secret_managers.main import KeyManagementSettings, KeyManagem CONTRACT_PATH: Final = Path(__file__).parents[3] / "litellm-rust/crates/python-bridge/python_settings.json" -def test_the_rust_contract_matches_the_returned_fields() -> None: - contract: Final = TypeAdapter(dict[str, list[str]]).validate_json(CONTRACT_PATH.read_text()) +class SettingSpec(TypedDict): + adapter: ReadOnly[str] + required: ReadOnly[bool] + precedence: ReadOnly[str] + sensitive: ReadOnly[bool] + shapes: ReadOnly[list[str]] + unsupported_live: ReadOnly[str | None] - assert contract == { - "http_settings": [field.name for field in dataclasses.fields(settings.http_settings())], - "url_policy": [field.name for field in dataclasses.fields(settings.url_policy())], - "provider_defaults": [field.name for field in dataclasses.fields(settings.provider_defaults())], - "secret_manager": [field.name for field in dataclasses.fields(settings.secret_manager())], + +class SettingsGroup(TypedDict): + version: ReadOnly[int] + fields: ReadOnly[dict[str, SettingSpec]] + + +def test_the_rust_contract_matches_the_returned_fields() -> None: + contract: Final = TypeAdapter(dict[str, SettingsGroup]).validate_json(CONTRACT_PATH.read_text()) + + assert {name: tuple(group["fields"]) for name, group in contract.items()} == { + "http_settings": tuple(field.name for field in dataclasses.fields(settings.http_settings())), + "url_policy": tuple(field.name for field in dataclasses.fields(settings.url_policy())), + "provider_defaults": tuple(field.name for field in dataclasses.fields(settings.provider_defaults())), + "secret_manager": tuple(field.name for field in dataclasses.fields(settings.secret_manager())), } diff --git a/tests/test_litellm_rust/ocr/test_requests.py b/tests/test_litellm_rust/ocr/test_requests.py index 5e9d2c78808..51815651eb4 100644 --- a/tests/test_litellm_rust/ocr/test_requests.py +++ b/tests/test_litellm_rust/ocr/test_requests.py @@ -3,14 +3,13 @@ from collections.abc import Callable from dataclasses import dataclass from io import BytesIO from pathlib import Path -from typing import Final +from typing import Final, NoReturn import httpx import pytest from pydantic import JsonValue import litellm -from litellm.llms.base_llm.ocr.transformation import OCRResponse from tests.test_litellm_rust.support.callback_recorder import RecordingLogger from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec from tests.test_litellm_rust.support.requests import ( @@ -503,3 +502,150 @@ async def test_native_failures_raise_the_public_exception_class( assert len(ocr_server.requests) == failure.provider_requests if failure.cause is not None: assert isinstance(caught.value.__context__, failure.cause) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize( + "name,value", + [ + ("ssl_verify", object()), + ("ssl_certificate", 1), + ("ssl_certificate", ""), + ("vertex_project", 1), + ("vertex_location", ["region"]), + ("user_url_allowed_hosts", ["example.test", 1]), + ], +) +async def test_native_settings_fail_before_provider_io( + ocr_server: RecordingServer, + monkeypatch: pytest.MonkeyPatch, + asynchronous: bool, + name: str, + value: object, +) -> None: + ocr_server.expected_requests = 0 + monkeypatch.setattr(litellm, name, value) + with pytest.raises(ValueError, match=r"http_settings|provider_defaults|url_policy"): + await call_native(ocr_server, asynchronous, num_retries=0) + assert ocr_server.requests == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +async def test_native_ssl_context_is_terminal_configuration(ocr_server: RecordingServer, asynchronous: bool) -> None: + import ssl + + ocr_server.expected_requests = 0 + context: Final = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) + with pytest.raises(ValueError, match=r"request\.ssl_verify.*SSLContext"): + await call_native(ocr_server, asynchronous, ssl_verify=context, num_retries=0) + assert ocr_server.requests == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +async def test_native_settings_preserve_protocol_failures( + ocr_server: RecordingServer, monkeypatch: pytest.MonkeyPatch, asynchronous: bool +) -> None: + ocr_server.expected_requests = 0 + failure: Final = LookupError("settings truth test failed") + cause: Final = RuntimeError("settings cause") + + class RaisesBool: + def __bool__(self) -> bool: + raise failure from cause + + monkeypatch.setattr(litellm, "force_ipv4", RaisesBool()) + with pytest.raises(LookupError) as caught: + await call_native(ocr_server, asynchronous, num_retries=0) + assert caught.value is failure + assert caught.value.__cause__ is cause + assert caught.value.__traceback__ is not None + assert ocr_server.requests == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +async def test_native_settings_observe_mutation_between_calls( + ocr_server: RecordingServer, monkeypatch: pytest.MonkeyPatch, asynchronous: bool +) -> None: + monkeypatch.setattr(litellm, "force_ipv4", "yes") + monkeypatch.setattr(litellm, "http2", 1) + monkeypatch.setattr(litellm, "vertex_project", []) + monkeypatch.setattr(litellm, "vertex_location", 0) + monkeypatch.setattr(litellm, "user_url_allowed_hosts", "EXAMPLE.TEST.") + response: Final = await call_native(ocr_server, asynchronous, num_retries=0) + assert response.pages[0].markdown == "native OCR response" + assert_native_request(ocr_server) + monkeypatch.setattr(litellm, "ssl_certificate", 1) + with pytest.raises(ValueError, match=r"http_settings\.ssl_certificate"): + await call_native(ocr_server, asynchronous, num_retries=0) + assert len(ocr_server.requests) == 1 + + +@pytest.mark.parametrize("required", [False, True]) +@pytest.mark.parametrize("failure", ["invalid", "live", "schema"]) +def test_native_projection_errors_never_select_python( + ocr_server: RecordingServer, monkeypatch: pytest.MonkeyPatch, required: bool, failure: str +) -> None: + import dataclasses + import ssl + + from litellm.rust_bridge import runtime, settings + from litellm.rust_bridge.catalog import Context, Route, Rule + from litellm.rust_bridge.configuration import Rollout + from litellm.rust_bridge.ocr.entrypoints import NATIVE_OCR, LiteLLMOcrRequest + + ocr_server.expected_requests = 0 + snapshot: Final = dataclasses.replace(settings.http_settings(), user_agent=1) + if failure == "schema": + monkeypatch.setattr(settings, "http_settings", lambda: snapshot) + else: + monkeypatch.setattr( + litellm, "ssl_verify", ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) if failure == "live" else object() + ) + request: Final = LiteLLMOcrRequest( + model="mistral/mistral-ocr-latest", + document=OCR_DOCUMENT, + api_key="test-key", + api_base=ocr_server.base_url, + timeout=None, + custom_llm_provider="mistral", + extra_headers=None, + kwargs={}, + ) + + def python_fallback() -> NoReturn: + pytest.fail("projection failures must not select Python") + + with pytest.raises(RuntimeError if failure == "schema" else ValueError, match="http_settings"): + runtime.run( + Context(Route.OCR, provider="mistral"), + binding=NATIVE_OCR, + native=lambda native: native(request, (), {}), + python=python_fallback, + rules=(Rule(Route.OCR, Rollout.RUST_REQUIRED if required else Rollout.RUST_OPT_OUT),), + ) + assert ocr_server.requests == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True], ids=["sync", "async"]) +@pytest.mark.parametrize("present", [False, True], ids=["missing", "invalid-pem"]) +async def test_native_client_certificate_is_validated_before_io( + ocr_server: RecordingServer, + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + asynchronous: bool, + present: bool, +) -> None: + ocr_server.expected_requests = 0 + certificate: Final = tmp_path / "client.pem" + if present: + certificate.write_text("invalid certificate") + monkeypatch.setattr(litellm, "ssl_certificate", str(certificate)) + with pytest.raises(ValueError, match=r"http_settings\.ssl_certificate.*PEM") as caught: + await call_native(ocr_server, asynchronous, num_retries=0) + assert str(certificate) not in str(caught.value) + assert ocr_server.requests == [] From 27808b51a0d82602457fef61e8befb6ecb847694 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:09:50 +0000 Subject: [PATCH 22/24] chore(rust): drop interop planning note Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- PYTHON_INTEROP_PLAN.md | 88 ------------------------------------------ 1 file changed, 88 deletions(-) delete mode 100644 PYTHON_INTEROP_PLAN.md diff --git a/PYTHON_INTEROP_PLAN.md b/PYTHON_INTEROP_PLAN.md deleted file mode 100644 index a7783f81472..00000000000 --- a/PYTHON_INTEROP_PLAN.md +++ /dev/null @@ -1,88 +0,0 @@ -# Python interop foundation PR plan - -Proposed implementation PR title: `fix(rust): preserve Python settings semantics at the native boundary` - -Base: `main` at `457b01e96d131f88df8cace8832f8e044ee5f167`. Planning branch: `litellm_python_interop_foundation` - -This plan follows `migration/tdd/sections/rust/python-interop/{pyo3-contract,boundary,coercion}.typ` in the sibling `litellm-typst` checkout. The first PR establishes conversion contracts and exercises them through existing HTTP, URL policy, and OCR provider-default consumers. The foundation is implemented on this branch. No PR is being created for this task - -**Baseline before implementation** - -`host-python/src/marshal.rs` already uses `pythonize` directly, separates internal conversion from public argument errors, and contains serializer panics in `Pythonized`. Keep those entrypoints. `Pythonized` currently stringifies conversion errors, unlike `from_py` and `to_py`, so its error transfer needs correction - -`core-utils/src/serde_compat.rs` already provides composable `LaxI64` and `FiniteF64` adapters. They first deserialize into `serde_json::Value`. Preserve their accepted input contracts while moving scalar decoding to Serde visitors, avoiding an intermediate JSON representation and making behavior through `pythonize` explicit - -`python-bridge/src/http.rs` currently uses strict derived Boolean and string extraction for mutable settings, converts extraction errors into `RustBridgeDeclined`, and silently drops unsupported `ssl_verify` values. OCR provider defaults have the same strict-extraction/decline pattern. `python_settings.json` checks field names only. These are the initial production consumers and regressions for the foundation - -**Ownership and placement** - -Keep interpreter attachment, structured-data conversion, panic containment, and execution adapters in `host-python`. Keep field paths, configuration error policy, named settings coercion, and product-specific tagged inputs in `python-bridge`. Keep pure parsing and Serde adapters in `core-utils`. This requires no new crate and no domain dependency from `host-python` into `core-utils` - -Add one focused `python-bridge/src/coercion.rs` module containing the field reader, semantic wrappers, and projection errors. Use `core-utils/src/serde_compat.rs` for shared token/numeric parsing and Serde decoding initially, splitting only if its size warrants it. The bridge can call those pure helpers through its existing dependency. Do not introduce a generic coercion registry, runtime manifest dispatch, or a new conversion framework - -Settings snapshots hold raw `Bound<'py, PyAny>` values only during attached projection. Successful adapters return owned values. Python snapshot dataclass annotations use `object` for arbitrary mutable globals, retaining strict annotations for accessor-owned fields such as `user_agent` and `readable`. No live settings object, iterator, or borrowed Python value enters native HTTP state - -**Serde and structured conversion** - -Keep `from_py` and `to_py` as the internal, exception-preserving conversion path. Make `Pythonized` use the same standard `PythonizeError` to `PyErr` conversion without losing its panic guard. Keep the explicitly argument-focused `ValueError` contract of `from_py_argument`; settings projection never passes through that helper - -Implement scalar visitors behind the existing `LaxI64` and `FiniteF64` adapter names. Preserve integer precision beyond the exact f64 range, signed bounds, supported decimal/underscore strings, Boolean numeric behavior, fractional rejection for integers, and nonfinite rejection for floats. Preserve composition inside `Option` and sequences, missing/null behavior at the field boundary, and ordinary numeric serialization. Use the existing [serde_with DeserializeAs contract](https://docs.rs/serde_with/3.16.1/serde_with/trait.DeserializeAs.html) - -Add the pure `parse_str_bool(&str) -> Option` parser used by bridge `StrBool`, HTTP `SslVerify::parse`, and the environment switch helper. It recognizes trimmed, case-insensitive true/false only. The environment helper retains its separate rule that only `Some(true)` contributes an enabled layer. Numeric helpers are shared only where a second actual consumer needs them - -Run common ordinary-value fixtures through JSON deserialization and direct Python-to-typed-Serde conversion. Restrict equivalence claims to their overlapping input domain. Live descriptors, identity, truthiness, iteration, and stringification are exercised through PyO3 separately. Do not add unused Serde counterparts for Python-only semantics or convert live settings through JSON text, `serde_json::Value`, `repr`, or `py_literal` - -**Named settings semantics** - -| Adapter | Contract and first consumer | -| --- | --- | -| `Truthy` | Execute Python truth testing and preserve its exception; IPv4, URL validation, and trust-env globals | -| `ExactTrue` | Compare identity with the True singleton without equality or truth testing; HTTP2, transport disable, and token refresh | -| `StrBool` | None or an actual string parsed by the shared parser; no arbitrary stringification | -| `OptionalStrictString` | None or an actual string, including the empty string; client certificate | -| `FalsyOptionalString` | Test truthiness first, treat falsey as absent, reject truthy non-strings; provider project and location | -| `TuningString` | Test truthiness first, then retain actual strings and ignore other values; TLS tuning | -| `StringCollection` | Apply the field's explicit container/member policy and return owned strings; URL allowlist | -| `SslVerifyInput` | Classify None, actual Boolean, Boolean string, CA path, live SSLContext, and invalid types separately | - -A single generic Boolean or optional-string coercer cannot implement these contracts. Accept string subclasses by their Unicode contents without calling overridden convenience methods. Use no `.ok()` or default value to discard an error from a Python protocol operation - -For URL hosts, a direct string represents one host. Otherwise test container truthiness and iterate, test member truthiness, skip falsey members, and reject truthy non-string members. Normalize with the existing URL-policy rules, deduplicate, and sort only where membership makes order irrelevant. Keep origin/scheme/port parsing in its existing domain helper rather than applying hostname normalization blindly to arbitrary URL strings - -**Errors and first production adoption** - -Represent projection failures as a tagged result separating original Python exceptions, invalid configuration, unsupported live objects, and internal accessor/schema failures. Map it once at the bridge: preserve original `PyErr`, use field-focused `ValueError` for invalid or unsupported configuration, and `RuntimeError` for genuine internal schema failures. Diagnostics contain group, field, expected forms, and actual type, never the supplied value or its representation - -Attribute, truthiness, iteration, and explicit stringification exceptions retain identity, traceback, cause, and context. In particular, do not relabel an AttributeError deliberately raised by a descriptor as a missing-field schema failure. Contract validation must distinguish schema drift from errors executing Python behavior - -Convert HTTP and URL snapshots, per-call `ssl_verify`, and OCR provider defaults through the named adapters. Preserve existing call/environment/global precedence and projection timing. Finish projection before constructing the native client or starting provider I/O. Invalid values and live SSLContext must raise configuration errors rather than disappear or authorize fallback - -Keep certificate paths through projection and validate empty, missing, or unusable client-certificate paths before I/O. The current native layer filters empty client-certificate paths, so correcting that narrow downstream behavior is part of adoption. Keep the existing CA-bundle missing-file policy explicit and separately tested; do not silently conflate it with client-certificate validation - -Classify the existing Secret Manager `readable` field as a strict accessor Boolean. Full Secret Manager client/system/settings snapshots and callback execution remain a separate PR. The existing readable-manager capability gap must be documented and must not be reported as fixed by this foundation - -**Semantic manifest** - -Extend `python_settings.json` with a stable group version and field records containing adapter ID, requiredness, precedence role, sensitivity, and specialized accepted/unsupported shapes. Include only the snapshot fields that exist in this PR. Update the Python contract test and a static Rust `SettingSpec` table to agree with the manifest - -The manifest checks declared contracts; behavioral tests prove the adapters implement them. Projection stays direct typed code. A manifest row alone is never evidence that a coercion works - -**Behavioral validation** - -Extend the existing mapped Python settings tests and Rust marshal/HTTP tests. A new coercion module may have its own focused Rust tests. Use the existing installed-extension OCR suites for public-route regressions. Do not add source-text assertions or class-attribute monkeypatching - -The acceptance matrix covers None, Boolean values, integer zero/one, strings, containers, subclasses, and arbitrary objects. Protocol fixtures raise pre-created exceptions from descriptors, `__bool__`, `__len__`, `__iter__`, and `__next__`; assert identity and exception chains. Verify ExactTrue never invokes hostile equality/truthiness. Verify falsey provider defaults preserve fallback, HTTP2 does not accept integer one, and a false/unknown environment token cannot switch off a true global - -Cover a real SSLContext, unsupported objects, Boolean strings, certificate path failures, direct-string hosts, sets, generators, duplicates, mixed members, and protocol failures. Mutate globals and source collections between calls: the next snapshot observes changes and an already projected value stays unchanged. At the installed public OCR boundary, projection failures must cause zero provider requests and zero Python fallback calls under required-native execution - -Run focused crate tests first, then the workspace Rust checks used by CI, `make test-rust-extension`, relevant Python settings tests, and `make check`. Use the fresh installed wheel and verify native provenance. Review the saved `make check` log rather than rerunning it to inspect output - -Target mutation tests at truthiness versus strict extraction, identity versus equality, swallowed versus preserved exceptions, string-as-one versus character iteration, normalization/deduplication, numeric bounds, and terminal errors versus fallback. Aim for more than 90% killed non-equivalent mutants in the changed coercion paths - -Before opening the implementation PR for maintainer review, provide a reproducible localhost proxy curl request with a real provider and positive native-execution evidence. Record the configured settings and user-visible result without credentials. Unit tests belong in validation, not the proof-of-fix section. Require the current tip's CI and coverage, Greptile confidence of at least 4/5, and acceptable Veria/Bugbot results; pending or unavailable results remain explicitly unresolved - -**Commit sequence and follow-ups** - -Start with the Serde visitor/error-transfer changes and their regression tests. Follow with field adapters and the shared token parser. Adopt them in HTTP, URL policy, and provider defaults together with the semantic manifest and installed-extension regressions. Keep these as reviewable commits in one foundational implementation PR - -Follow-up PRs can add `OptionalRedisBool`, cache-specific stringification and collection rules, and full Secret Manager bindings using the same field/error machinery. Redis accepts a different token set from StrBool, so do not share their Boolean semantics. Runtime redesign, callback lifecycle changes, free-threaded support, wholesale request serialization, and unrelated cache work are outside this PR From 9c48e137dcf63853c4ae75f2f060945ef44e2839 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:21:23 +0000 Subject: [PATCH 23/24] fix(rust): preserve HTTP host and TLS error context Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/crates/http/src/config.rs | 13 ++++- litellm-rust/crates/http/src/error.rs | 18 +++++- litellm-rust/crates/http/src/lib.rs | 2 +- litellm-rust/crates/http/src/media.rs | 46 ++++++++++++++-- litellm-rust/crates/http/src/tls.rs | 49 ++++++++++------- litellm-rust/crates/python-bridge/src/http.rs | 55 +++++++++++++++---- .../python-bridge/src/routes/ocr/mod.rs | 2 +- 7 files changed, 146 insertions(+), 39 deletions(-) diff --git a/litellm-rust/crates/http/src/config.rs b/litellm-rust/crates/http/src/config.rs index bf8ecef85a8..cb0173369d5 100644 --- a/litellm-rust/crates/http/src/config.rs +++ b/litellm-rust/crates/http/src/config.rs @@ -129,6 +129,7 @@ mod tests { use rstest::rstest; use super::*; + use crate::TlsSource; fn settings(ssl_verify: Option, ssl_cert_file: Option<&str>) -> HttpSettings { HttpSettings { @@ -298,7 +299,11 @@ mod tests { }; assert!(matches!( reqwest::ClientBuilder::try_from(&config), - Err(Error::Read { path: reported, .. }) if reported == path + Err(Error::Read { + path: reported, + tls_source: TlsSource::CaBundle, + .. + }) if reported == path )); } @@ -315,7 +320,11 @@ mod tests { std::fs::remove_file(&path).unwrap(); assert!(matches!( result, - Err(Error::InvalidPem { path: reported, .. }) if reported == path + Err(Error::InvalidPem { + path: reported, + tls_source: TlsSource::CaBundle, + .. + }) if reported == path )); } } diff --git a/litellm-rust/crates/http/src/error.rs b/litellm-rust/crates/http/src/error.rs index e06f7c00cf5..eafb4d2976b 100644 --- a/litellm-rust/crates/http/src/error.rs +++ b/litellm-rust/crates/http/src/error.rs @@ -1,11 +1,25 @@ use std::path::PathBuf; +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum TlsSource { + CaBundle, + ClientIdentity, +} + #[derive(Clone, Debug, thiserror::Error, PartialEq, Eq)] pub enum Error { #[error("could not read {}: {message}", path.display())] - Read { path: PathBuf, message: String }, + Read { + path: PathBuf, + message: String, + tls_source: TlsSource, + }, #[error("{} is not a PEM file: {message}", path.display())] - InvalidPem { path: PathBuf, message: String }, + InvalidPem { + path: PathBuf, + message: String, + tls_source: TlsSource, + }, #[error("could not build the HTTP client: {0}")] Client(String), #[error("request body could not be serialized: {0}")] diff --git a/litellm-rust/crates/http/src/lib.rs b/litellm-rust/crates/http/src/lib.rs index 6f62a00175c..a1456208bb3 100644 --- a/litellm-rust/crates/http/src/lib.rs +++ b/litellm-rust/crates/http/src/lib.rs @@ -10,7 +10,7 @@ mod tls; pub mod transport; pub use config::{HttpClientConfig, Resolution, Verify}; -pub use error::Error; +pub use error::{Error, TlsSource}; pub use pool::{ClientVariant, HttpClientPool}; pub use proxy::EnvironmentProxies; pub use settings::{HttpSettings, HttpSettingsLayer, SslVerify, TcpKeepalive}; diff --git a/litellm-rust/crates/http/src/media.rs b/litellm-rust/crates/http/src/media.rs index 753f25c29de..1b9159973ef 100644 --- a/litellm-rust/crates/http/src/media.rs +++ b/litellm-rust/crates/http/src/media.rs @@ -54,16 +54,39 @@ impl Default for UrlPolicy { impl UrlPolicy { fn allows(&self, host: &str, port: u16) -> bool { let host = normalize_host(host); - let with_port = format!("{host}:{port}"); self.allowed_hosts .iter() - .map(|entry| normalize_host(entry)) - .any(|entry| entry == host || entry == with_port) + .filter_map(|entry| parse_allowed_host(entry)) + .any(|(entry_host, entry_port)| { + entry_host == host && entry_port.is_none_or(|entry_port| entry_port == port) + }) } } pub fn normalize_host(host: &str) -> String { - host.to_ascii_lowercase().trim_end_matches('.').to_owned() + let host = host.trim().trim_end_matches('.'); + let host = host + .strip_prefix('[') + .and_then(|host| host.strip_suffix(']')) + .unwrap_or(host); + host.to_ascii_lowercase() +} + +fn parse_allowed_host(entry: &str) -> Option<(String, Option)> { + let entry = entry.trim(); + if let Some(entry) = entry.strip_prefix('[') { + let (host, suffix) = entry.split_once(']')?; + let port = match suffix { + "" => None, + suffix => Some(suffix.strip_prefix(':')?.parse().ok()?), + }; + return Some((normalize_host(host), port)); + } + let (host, port) = match entry.rsplit_once(':') { + Some((host, port)) if !host.contains(':') => (host, Some(port.parse().ok()?)), + _ => (entry, None), + }; + Some((normalize_host(host), port)) } type ProxyMatch = Arc bool + Send + Sync>; @@ -670,6 +693,21 @@ mod tests { assert!(matches!(result, Err(Error::BlockedUrl))); } + #[test] + fn allowlist_matches_bracketed_ipv6_hosts_and_ports() { + let policy = UrlPolicy { + validate: true, + allowed_hosts: vec!["[2001:db8::1]".into(), "[2001:db8::1]:8443".into()], + }; + assert!(policy.allows("2001:db8::1", 443)); + assert!(policy.allows("2001:db8::1", 8443)); + let port_specific = UrlPolicy { + validate: true, + allowed_hosts: vec!["[2001:db8::1]:8443".into()], + }; + assert!(!port_specific.allows("2001:db8::1", 9443)); + } + #[tokio::test] async fn validation_off_fetches_private_hosts_and_follows_redirects() { let (url, server, _) = serve_named( diff --git a/litellm-rust/crates/http/src/tls.rs b/litellm-rust/crates/http/src/tls.rs index aaae2b659e3..e2e6d27cd54 100644 --- a/litellm-rust/crates/http/src/tls.rs +++ b/litellm-rust/crates/http/src/tls.rs @@ -9,7 +9,7 @@ use rustls::{ use crate::{ config::{HttpClientConfig, Verify}, - error::Error, + error::{Error, TlsSource}, }; #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] @@ -197,15 +197,17 @@ impl TryFrom<&HttpClientConfig> for ClientConfig { Verify::BuiltInRoots => builder.with_root_certificates(RootCertStore { roots: webpki_roots::TLS_SERVER_ROOTS.to_vec(), }), - Verify::CaBundle(path) => builder.with_root_certificates(bundle_roots(path)?), + Verify::CaBundle(path) => { + builder.with_root_certificates(bundle_roots(path, TlsSource::CaBundle)?) + } }; let mut tls = match &config.client_certificate { None => verified.with_no_client_auth(), Some(path) => { - let (chain, key) = identity(path)?; + let (chain, key) = identity(path, TlsSource::ClientIdentity)?; verified .with_client_auth_cert(chain, key) - .map_err(|error| invalid_pem(path, error))? + .map_err(|error| invalid_pem(path, TlsSource::ClientIdentity, error))? } }; tls.alpn_protocols = if config.http2 { @@ -217,47 +219,52 @@ impl TryFrom<&HttpClientConfig> for ClientConfig { } } -fn bundle_roots(path: &Path) -> Result { - let certificates = certificates(path)?; +fn bundle_roots(path: &Path, source: TlsSource) -> Result { + let certificates = certificates(path, source)?; if certificates.is_empty() { - return Err(invalid_pem(path, "no certificates found")); + return Err(invalid_pem(path, source, "no certificates found")); } let mut store = RootCertStore::empty(); for certificate in certificates { store .add(certificate) - .map_err(|error| invalid_pem(path, error))?; + .map_err(|error| invalid_pem(path, source, error))?; } Ok(store) } -fn identity(path: &Path) -> Result<(Vec>, PrivateKeyDer<'static>), Error> { - let chain = certificates(path)?; +fn identity( + path: &Path, + source: TlsSource, +) -> Result<(Vec>, PrivateKeyDer<'static>), Error> { + let chain = certificates(path, source)?; if chain.is_empty() { - return Err(invalid_pem(path, "no certificates found")); + return Err(invalid_pem(path, source, "no certificates found")); } - let key = - PrivateKeyDer::from_pem_slice(&read(path)?).map_err(|error| invalid_pem(path, error))?; + let key = PrivateKeyDer::from_pem_slice(&read(path, source)?) + .map_err(|error| invalid_pem(path, source, error))?; Ok((chain, key)) } -fn certificates(path: &Path) -> Result>, Error> { - CertificateDer::pem_slice_iter(&read(path)?) +fn certificates(path: &Path, source: TlsSource) -> Result>, Error> { + CertificateDer::pem_slice_iter(&read(path, source)?) .collect::>() - .map_err(|error| invalid_pem(path, error)) + .map_err(|error| invalid_pem(path, source, error)) } -fn read(path: &Path) -> Result, Error> { +fn read(path: &Path, source: TlsSource) -> Result, Error> { std::fs::read(path).map_err(|error| Error::Read { path: path.to_path_buf(), message: error.to_string(), + tls_source: source, }) } -fn invalid_pem(path: &Path, message: impl fmt::Display) -> Error { +fn invalid_pem(path: &Path, source: TlsSource, message: impl fmt::Display) -> Error { Error::InvalidPem { path: path.to_path_buf(), message: message.to_string(), + tls_source: source, } } @@ -405,7 +412,11 @@ mod tests { std::fs::remove_file(&path).unwrap(); assert!(matches!( result, - Err(Error::InvalidPem { path: reported, .. }) if reported == path + Err(Error::InvalidPem { + path: reported, + tls_source: TlsSource::ClientIdentity, + .. + }) if reported == path )); } } diff --git a/litellm-rust/crates/python-bridge/src/http.rs b/litellm-rust/crates/python-bridge/src/http.rs index 859d579129f..596a89a73d7 100644 --- a/litellm-rust/crates/python-bridge/src/http.rs +++ b/litellm-rust/crates/python-bridge/src/http.rs @@ -7,7 +7,7 @@ use std::{ use litellm_core_utils::settings::ProcessEnvironment; use litellm_http::{ HttpClientConfig, HttpClientPool, HttpSettings, HttpSettingsLayer, Resolution, SslVerify, - Unsupported, + TlsSource, Unsupported, media::{PublicDnsResolver, UrlPolicy}, }; use pyo3::{exceptions::PyValueError, prelude::*, types::PyDict}; @@ -41,18 +41,26 @@ pub(crate) fn call_config( Ok(resolution.config) } -pub(crate) fn client_error(error: litellm_http::Error, config: &HttpClientConfig) -> PyErr { +pub(crate) fn client_error(error: litellm_http::Error) -> PyErr { match error { - litellm_http::Error::Read { path, .. } | litellm_http::Error::InvalidPem { path, .. } - if config.client_certificate.as_ref() == Some(&path) => - { - PyValueError::new_err( - "http_settings.ssl_certificate: expected a readable PEM certificate and private key", - ) + litellm_http::Error::Read { + tls_source: TlsSource::ClientIdentity, + .. } - litellm_http::Error::Read { .. } | litellm_http::Error::InvalidPem { .. } => { - PyValueError::new_err("http_settings.ssl_verify: expected a readable PEM CA bundle") + | litellm_http::Error::InvalidPem { + tls_source: TlsSource::ClientIdentity, + .. + } => PyValueError::new_err( + "http_settings.ssl_certificate: expected a readable PEM certificate and private key", + ), + litellm_http::Error::Read { + tls_source: TlsSource::CaBundle, + .. } + | litellm_http::Error::InvalidPem { + tls_source: TlsSource::CaBundle, + .. + } => PyValueError::new_err("http_settings.ssl_verify: expected a readable PEM CA bundle"), _ => PyValueError::new_err("http_settings: native HTTP client configuration is invalid"), } } @@ -188,6 +196,33 @@ settings = types.SimpleNamespace(**{{name: defaults[name] for name in json.loads }); } + #[test] + fn client_error_uses_tls_source_when_paths_match() { + Python::initialize(); + Python::attach(|py| { + let path = PathBuf::from("/shared.pem"); + let ca_error = client_error(litellm_http::Error::InvalidPem { + path: path.clone(), + message: "invalid".into(), + tls_source: TlsSource::CaBundle, + }); + assert_eq!( + ca_error.to_string(), + "ValueError: http_settings.ssl_verify: expected a readable PEM CA bundle" + ); + let client_error = client_error(litellm_http::Error::InvalidPem { + path, + message: "invalid".into(), + tls_source: TlsSource::ClientIdentity, + }); + assert!(client_error.is_instance_of::(py)); + assert_eq!( + client_error.to_string(), + "ValueError: http_settings.ssl_certificate: expected a readable PEM certificate and private key" + ); + }); + } + #[test] fn python_settings_flow_into_the_configured_layer() { Python::initialize(); diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs index ce6f04c321b..d0b13e5056a 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/mod.rs @@ -51,7 +51,7 @@ fn run_ocr( ocr_settings(py)?, secrets, ) - .map_err(|error| http::client_error(error, &config))?; + .map_err(http::client_error)?; run_legacy_call( py, if asynchronous { ASYNC_SURFACE } else { SURFACE }, From c7c0afb1f0b0c737b6f561065c67563a929b5ae0 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 21:21:26 +0000 Subject: [PATCH 24/24] ci(rust): raise the native wheel size gate to 40 MB Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/scripts/verify_linux_native_wheel.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/.github/scripts/verify_linux_native_wheel.py b/.github/scripts/verify_linux_native_wheel.py index 0adbc015ad0..ea6d2401084 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 = 40_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), + (f"Native extension does not exceed {native_size_limit / 1_000_000:.0f} MB", native_size_within_limit), ("Wheel contents are valid", not unexpected_members), ) @@ -267,7 +267,8 @@ 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 {native_size_limit / 1_000_000:.0f} MB: " + f"{native_member.file_size / 1_000_000:.2f} MB", ), (bool(unexpected_members), f"wheel contains unexpected build artifacts: {', '.join(unexpected_members)}"), )