diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 962d1266fd7..ef88ae574fb 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -1,6 +1,6 @@ { "reportAny": { - "limit": 17271 + "limit": 17270 }, "reportArgumentType": { "limit": 2539 diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index cf2cee9b6ef..570b306d6df 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -36,6 +36,7 @@ from litellm.llms.base_llm.managed_resources.isolation import ( build_list_page, build_owner_filter, can_access_resource, + resolve_resource_owner_id, ) from litellm.proxy._types import ( CallTypes, @@ -222,7 +223,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): file_object=file_object, model_mappings=model_mappings, flat_model_file_ids=list(model_mappings.values()), - created_by=user_api_key_dict.user_id, + created_by=resolve_resource_owner_id(user_api_key_dict), team_id=user_api_key_dict.team_id, updated_by=user_api_key_dict.user_id, ) @@ -238,7 +239,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): "unified_file_id": file_id, "model_mappings": json.dumps(model_mappings), "flat_model_file_ids": list(model_mappings.values()), - "created_by": user_api_key_dict.user_id, + "created_by": resolve_resource_owner_id(user_api_key_dict), "team_id": user_api_key_dict.team_id, "updated_by": user_api_key_dict.user_id, } @@ -342,7 +343,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): "file_object": file_object.model_dump_json(), "model_object_id": model_object_id, "file_purpose": file_purpose, - "created_by": user_api_key_dict.user_id, + "created_by": resolve_resource_owner_id(user_api_key_dict), "team_id": user_api_key_dict.team_id, "updated_by": user_api_key_dict.user_id, "status": file_object.status, diff --git a/litellm/llms/base_llm/managed_resources/base_managed_resource.py b/litellm/llms/base_llm/managed_resources/base_managed_resource.py index 2a59eddf88a..cced330d873 100644 --- a/litellm/llms/base_llm/managed_resources/base_managed_resource.py +++ b/litellm/llms/base_llm/managed_resources/base_managed_resource.py @@ -12,6 +12,7 @@ from litellm.llms.base_llm.managed_resources.isolation import ( build_list_page, build_owner_filter, can_access_resource, + resolve_resource_owner_id, ) from litellm.proxy._types import UserAPIKeyAuth from litellm.types.utils import SpecialEnums @@ -157,7 +158,7 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): "resource_object": resource_object, "model_mappings": model_mappings, "flat_model_resource_ids": list(model_mappings.values()), - "created_by": user_api_key_dict.user_id, + "created_by": resolve_resource_owner_id(user_api_key_dict), "team_id": user_api_key_dict.team_id, "updated_by": user_api_key_dict.user_id, } @@ -179,7 +180,7 @@ class BaseManagedResource(ABC, Generic[ResourceObjectType]): "unified_resource_id": unified_resource_id, "model_mappings": json.dumps(model_mappings), "flat_model_resource_ids": list(model_mappings.values()), - "created_by": user_api_key_dict.user_id, + "created_by": resolve_resource_owner_id(user_api_key_dict), "team_id": user_api_key_dict.team_id, "updated_by": user_api_key_dict.user_id, } diff --git a/litellm/llms/base_llm/managed_resources/isolation.py b/litellm/llms/base_llm/managed_resources/isolation.py index e1b204214d7..6a71e8e9223 100644 --- a/litellm/llms/base_llm/managed_resources/isolation.py +++ b/litellm/llms/base_llm/managed_resources/isolation.py @@ -3,10 +3,11 @@ Tenant-isolation helpers for managed file/batch/vector-store resources. Returns a Prisma filter and an ownership check that scope managed resources to the caller's identity: proxy admins see everything, user-keyed callers -see records they created, and service-account keys (no user_id) fall back -to the resource's owning team. Callers with no admin role and no -identifying ids are denied so an empty user_id can never select an -unscoped query. +see records they created, service-account keys (no user_id) fall back to +the resource's owning team, and keys with neither a user_id nor a team_id +fall back to their own hashed token so they can still reach the resources +they created. Callers with no admin role and no identifying ids at all +are denied so an empty user_id can never select an unscoped query. """ from typing import Any, Final @@ -19,6 +20,32 @@ from litellm.proxy._types import ( ) +def resolve_resource_owner_id( + user_api_key_dict: UserAPIKeyAuth, +) -> str | None: + """Return the identity to stamp on (and match against) a managed + resource's ``created_by``. + + A key with neither a user_id nor a team_id would otherwise stamp + ``created_by=None`` and be locked out of its own resources, so it owns + them under its hashed token instead, using the ``key:`` scope prefix + already used by ``proxy/common_utils/resource_ownership.py``. ``None`` + means the caller has no usable identity of its own and must fall back + to team scoping, or be denied. + """ + if user_api_key_dict.user_id is not None: + return user_api_key_dict.user_id + + if user_api_key_dict.team_id is not None: + return None + + token: Final = user_api_key_dict.token or user_api_key_dict.api_key + if token: + return f"key:{token}" + + return None + + def build_list_page(items: list[Any], has_more: bool = False) -> dict[str, Any]: """Build the OpenAI-style paginated list response shape used by managed file/batch/vector-store listings. ``first_id`` and ``last_id`` are @@ -39,7 +66,8 @@ def build_owner_filter( to records the caller is allowed to see. - ``{}`` means no scoping (proxy admins). - - ``{"created_by": }`` for user-keyed callers. + - ``{"created_by": }`` for user-keyed callers, and for keys + with no user_id and no team_id (owner id is their hashed token). - ``{"team_id": }`` for service-account callers that have a team but no user_id. - ``{"OR": [...]}`` when the caller has both — listing must include @@ -62,12 +90,13 @@ def build_owner_filter( ] } - if user_id is not None: - return {"created_by": user_id} - if team_id is not None: return {"team_id": team_id} + owner_id: Final = resolve_resource_owner_id(user_api_key_dict) + if owner_id is not None: + return {"created_by": owner_id} + return None @@ -86,8 +115,8 @@ def can_access_resource( if _user_has_admin_view(user_api_key_dict): return True - user_id: Final = user_api_key_dict.user_id - if user_id is not None and created_by is not None and created_by == user_id: + owner_id: Final = resolve_resource_owner_id(user_api_key_dict) + if owner_id is not None and created_by is not None and created_by == owner_id: return True team_id: Final = user_api_key_dict.team_id diff --git a/litellm/main.py b/litellm/main.py index cafa1e4718f..c341db08155 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -8612,9 +8612,9 @@ def _joined_streamed_citations(streamed_citations: "tuple[object, ...]") -> "lis def _stream_builder_model_map_cost(response: ModelResponse) -> float | None: - model_name: Final = getattr(response, "model", None) + model_name: Final = response.model usage: Final = getattr(response, "usage", None) - if not isinstance(model_name, str) or not model_name or not isinstance(usage, Usage): + if not model_name or not isinstance(usage, Usage): return None try: prompt_cost, completion_tokens_cost = litellm.cost_per_token(model=model_name, usage_object=usage) diff --git a/litellm/proxy/db/routing_prisma_wrapper.py b/litellm/proxy/db/routing_prisma_wrapper.py index 1929e7d3fc8..be515392a17 100644 --- a/litellm/proxy/db/routing_prisma_wrapper.py +++ b/litellm/proxy/db/routing_prisma_wrapper.py @@ -61,6 +61,26 @@ class _RoutedActions: return getattr(self._writer_actions, name) +class WriterPinnedClient: + """PrismaClient-shaped view whose `.db` resolves to the writer while it is available. + + Read-after-write paths (e.g. the model reconcile a /model/new triggers to + verify its own just-committed row) must not read through a lagging read + replica: the row is not replayed there yet, so the reconcile concludes the + write is missing and fails the request even though it is durable (#38556). + + While the writer is degraded (`writer_unavailable`), the pin yields to the + routed wrapper so reconcile reads keep working from the replica: a proxy + that starts during a primary outage must still load DB-backed models, and + no read-after-write hazard exists then because writes are failing anyway. + """ + + __slots__ = ("db",) + + def __init__(self, db: "PrismaWrapper | RoutingPrismaWrapper") -> None: + self.db: Final = db.writer if isinstance(db, RoutingPrismaWrapper) and not db.writer_unavailable else db + + class RoutingPrismaWrapper: """ Routes Prisma operations between a writer and a reader Prisma client. diff --git a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py index 7791adeb41e..bcaffa8e91c 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py @@ -301,12 +301,12 @@ class LakeraAIGuardrail(CustomGuardrail): explicit sync below a hot reload that changes mode would pass validation but keep dispatching on the stale event_hook. """ - new_event_hook: Final = getattr(litellm_params, "mode", None) or self.event_hook - prospective_payload: Final = getattr(litellm_params, "payload", None) - prospective_breakdown: Final = getattr(litellm_params, "breakdown", None) + new_event_hook: Final = litellm_params.mode or self.event_hook + prospective_payload: Final = litellm_params.payload + prospective_breakdown: Final = litellm_params.breakdown self._validate_advisory_config( - on_flagged=getattr(litellm_params, "on_flagged", None) or self.on_flagged, - advisory_system_message=getattr(litellm_params, "advisory_system_message", None), + on_flagged=litellm_params.on_flagged or self.on_flagged, + advisory_system_message=litellm_params.advisory_system_message, payload=self.payload if prospective_payload is None else prospective_payload, breakdown=self.breakdown if prospective_breakdown is None else prospective_breakdown, ) diff --git a/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py b/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py index daeb91eb2bd..f834426d619 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py +++ b/litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py @@ -121,7 +121,7 @@ class QualifireGuardrail(CustomGuardrail): the live instance untouched instead of raising after it's already been corrupted. Mirrors LakeraAIGuardrail's own override of this same method. """ - prospective_on_flagged: Final = getattr(litellm_params, "on_flagged", None) or self.on_flagged + prospective_on_flagged: Final = litellm_params.on_flagged or self.on_flagged self._validate_on_flagged(prospective_on_flagged) super().update_in_memory_litellm_params(litellm_params=litellm_params) diff --git a/litellm/proxy/guardrails/guardrail_registry.py b/litellm/proxy/guardrails/guardrail_registry.py index 90d5f6f4970..dc13c09dd38 100644 --- a/litellm/proxy/guardrails/guardrail_registry.py +++ b/litellm/proxy/guardrails/guardrail_registry.py @@ -413,14 +413,15 @@ class GuardrailRegistry: raise Exception(f"Error getting guardrail from DB: {e}") -def _apply_configured_bool_override(instance: CustomGuardrail, litellm_params: LitellmParams, param_name: str) -> None: - """Override ``instance.`` only when ``litellm_params`` explicitly - sets it, preserving whatever default the guardrail's own constructor chose +def _apply_configured_bool_overrides(instance: CustomGuardrail, litellm_params: LitellmParams) -> None: + """Override the parallel/raw-scan flags only when ``litellm_params`` explicitly + sets them, preserving whatever default the guardrail's own constructor chose otherwise (its constructor default may be True, so blindly copying an absent/None config value would silently clobber it back to False).""" - configured: Final = getattr(litellm_params, param_name, None) - if configured is not None: - setattr(instance, param_name, bool(configured)) + if litellm_params.run_in_parallel is not None: + instance.run_in_parallel = bool(litellm_params.run_in_parallel) + if litellm_params.scan_raw_request is not None: + instance.scan_raw_request = bool(litellm_params.scan_raw_request) class InMemoryGuardrailHandler: @@ -544,8 +545,7 @@ class InMemoryGuardrailHandler: "skip_tool_message_in_guardrail are enabled together, which excludes every message from " "scanning, so no request content would ever be scanned. Remove one of the two." ) - for override_param in ("run_in_parallel", "scan_raw_request"): - _apply_configured_bool_override(custom_guardrail_callback, litellm_params, override_param) + _apply_configured_bool_overrides(custom_guardrail_callback, litellm_params) parsed_guardrail: Final = Guardrail( guardrail_id=guardrail.get("guardrail_id"), @@ -803,7 +803,6 @@ class InMemoryGuardrailHandler: previous_guardrail: Final = self.IN_MEMORY_GUARDRAILS.get(guardrail_id) previous_source: Final = self._sources.get(guardrail_id, source) - # Remove from memory if exists (also removes from callbacks) if guardrail_id in self.IN_MEMORY_GUARDRAILS: self.delete_in_memory_guardrail(guardrail_id) diff --git a/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py b/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py index 23cfef6576c..567d8375737 100644 --- a/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py +++ b/litellm/proxy/pass_through_endpoints/managed_id_rewriter.py @@ -49,6 +49,7 @@ from litellm._logging import verbose_proxy_logger from litellm.llms.base_llm.managed_resources.isolation import ( build_owner_filter, can_access_resource, + resolve_resource_owner_id, ) from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.batches_endpoints.common_utils import validate_batch_list_limit @@ -686,7 +687,7 @@ async def _mint_or_reuse_object( "file_object": json.dumps(body_snapshot), "model_object_id": namespaced_model_object_id, "file_purpose": file_purpose, - "created_by": user_api_key_dict.user_id, + "created_by": resolve_resource_owner_id(user_api_key_dict), "team_id": user_api_key_dict.team_id, "updated_by": user_api_key_dict.user_id, }, diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index a5619821197..50cb813c6fa 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -185,7 +185,7 @@ class PipelineExecutor: # snapshot instead of `data` (which earlier pass_data steps in # this same pipeline may have already rewritten), same reason # the normal sequential/parallel guardrail loops do this. - scans_raw_request: Final = getattr(callback, "scan_raw_request", False) + scans_raw_request: Final = callback.scan_raw_request hook_input: Final[dict] = ( # mutable-ok: same request-payload shape as data independent_snapshot(raw_request_snapshot) if scans_raw_request and raw_request_snapshot is not None diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 4bb2feff426..781fa21ed78 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -6844,9 +6844,18 @@ class ProxyConfig: - list: the rows (may be empty if no models exist) - None: signals a DB fetch *failure* — callers must not treat this as "all models deleted" and must not evict existing router deployments. + + Pinned to the writer DB: this read reconciles the router against the rows a + model write just committed, and reading it through a lagging read replica + makes the write-triggered reload report its own durable write as missing + (#38556). It also keeps a stale replica snapshot from evicting a deployment + another pod just added. While the writer is degraded the pin yields to the + replica so reader-only mode keeps loading DB-backed models. """ try: - new_models: Final[Sequence[_ProxyModelRow]] = await ModelRepository(prisma_client).table.find_many() + new_models: Final[Sequence[_ProxyModelRow]] = await ModelRepository( + WriterPinnedClient(prisma_client.db) + ).table.find_many() return new_models except Exception as e: verbose_proxy_logger.exception( @@ -12093,6 +12102,7 @@ async def run_thread( # ) # async def get_available_routes(user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth)): from litellm.llms.base_llm.base_utils import BaseTokenCounter +from litellm.proxy.db.routing_prisma_wrapper import WriterPinnedClient from litellm.repositories.config_repository import ConfigRepository from litellm.repositories.model_repository import ModelRepository from litellm.repositories.table_repositories import ( diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 1d16fa63607..eab56c31c39 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1416,7 +1416,7 @@ class ProxyLogging: mutation is discarded and a warning is logged so the misconfiguration is visible instead of silently forwarding unredacted content. """ - scans_raw_request: Final = getattr(callback, "scan_raw_request", False) + scans_raw_request: Final = callback.scan_raw_request should_use_raw_snapshot: Final = scans_raw_request and raw_request_snapshot is not None input_data: Final = ( # mutable-ok: same request-payload shape as data independent_snapshot(raw_request_snapshot) if should_use_raw_snapshot else data @@ -1453,7 +1453,7 @@ class ProxyLogging: "scan_raw_request is for block-only guardrails and this mutation is being " "discarded. Remove scan_raw_request from this guardrail's config if it needs " "to mask/rewrite content.", - getattr(callback, "guardrail_name", None) or callback.__class__.__name__, + callback.guardrail_name or callback.__class__.__name__, ) if scans_raw_request: if result is not None: @@ -1778,7 +1778,7 @@ class ProxyLogging: # guarantee must hold even under litellm.safe_memory_mode, which # otherwise makes deep copies return the original object. needs_raw_request_snapshot: Final = any( - isinstance(cb, CustomGuardrail) and getattr(cb, "scan_raw_request", False) + isinstance(cb, CustomGuardrail) and cb.scan_raw_request for cb in ProxyLogging._callback_capabilities().resolved_callbacks ) raw_request_snapshot: Final[dict | None] = ( # mutable-ok: same request-payload shape as data @@ -1938,7 +1938,7 @@ class ProxyLogging: """ def _input_for(callback: CustomGuardrail) -> dict: # mutable-ok: same request-payload shape as data - if not getattr(callback, "scan_raw_request", False) or raw_request_snapshot is None: + if not callback.scan_raw_request or raw_request_snapshot is None: return data return independent_snapshot(raw_request_snapshot) @@ -1962,11 +1962,7 @@ class ProxyLogging: # deployment-level guardrail sharing this name would see no marker # via _pre_call_hook_already_ran and re-run it a second time on # live kwargs. - if ( - getattr(callback, "scan_raw_request", False) - and not isinstance(result, BaseException) - and result is not None - ): + if callback.scan_raw_request and not isinstance(result, BaseException) and result is not None: callback.mark_pre_call_hook_ran(data) raised: Final = tuple(result for result in results if isinstance(result, BaseException)) blocking: Final = next((exc for exc in raised if not _exception_changes_request_flow(exc)), None) diff --git a/tests/test_litellm/enterprise/proxy/test_managed_files_access_check.py b/tests/test_litellm/enterprise/proxy/test_managed_files_access_check.py index c75c8099ea1..ad46798b788 100644 --- a/tests/test_litellm/enterprise/proxy/test_managed_files_access_check.py +++ b/tests/test_litellm/enterprise/proxy/test_managed_files_access_check.py @@ -10,11 +10,14 @@ with deployment credentials, bypassing the managed files access-control hooks. import base64 import pytest +from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch from fastapi import HTTPException -from litellm.proxy._types import UserAPIKeyAuth +from litellm.caching.dual_cache import DualCache +from litellm.proxy._types import CallTypes, UserAPIKeyAuth +from litellm.types.utils import LiteLLMBatch def _make_user_api_key_dict(user_id: str) -> UserAPIKeyAuth: @@ -161,6 +164,108 @@ async def test_service_account_blocked_from_other_team_file(): assert exc_info.value.status_code == 403 +# --- Keyless key must not be locked out of the batch it created --- + + +def _make_unified_batch_id() -> str: + raw = "litellm_proxy;model_id:my-model-id;llm_batch_id:batch_raw_123" + return base64.urlsafe_b64encode(raw.encode()).decode().rstrip("=") + + +def _make_managed_files_instance_with_object_store(): + """Managed-files hook backed by an in-memory stand-in for the managed + object table, so create and retrieve exercise the same stored row.""" + from litellm_enterprise.proxy.hooks.managed_files import ( + _PROXY_LiteLLMManagedFiles, + ) + + store = {} + + async def upsert(where, data): + store[where["unified_object_id"]] = SimpleNamespace(**data["create"]) + + async def find_first(where): + return store.get(where["unified_object_id"]) + + mock_prisma = MagicMock() + mock_prisma.db.litellm_managedobjecttable.upsert = AsyncMock(side_effect=upsert) + mock_prisma.db.litellm_managedobjecttable.find_first = AsyncMock( + side_effect=find_first + ) + + return ( + _PROXY_LiteLLMManagedFiles( + internal_usage_cache=DualCache(), + prisma_client=mock_prisma, + ), + store, + ) + + +async def _store_batch(managed_files, unified_batch_id: str, creator: UserAPIKeyAuth): + await managed_files.store_unified_object_id( + unified_object_id=unified_batch_id, + file_object=LiteLLMBatch( + id="batch_raw_123", + completion_window="24h", + created_at=0, + endpoint="/v1/chat/completions", + input_file_id="file-1", + object="batch", + status="validating", + ), + litellm_parent_otel_span=None, + model_object_id="batch_raw_123", + file_purpose="batch", + user_api_key_dict=creator, + ) + + +@pytest.mark.asyncio +async def test_keyless_key_can_retrieve_the_batch_it_created(): + """Regression: a key with no user_id and no team_id (what `/key/generate` + by a proxy admin and service-account keys produce) stamped + `created_by=None` and was then denied its own managed batch with + "User None does not have access".""" + unified_batch_id = _make_unified_batch_id() + managed_files, store = _make_managed_files_instance_with_object_store() + keyless = UserAPIKeyAuth(api_key="sk-keyless", parent_otel_span=None) + + await _store_batch(managed_files, unified_batch_id, keyless) + assert store[unified_batch_id].created_by == f"key:{keyless.token}" + + data = {"batch_id": unified_batch_id} + await managed_files.async_pre_call_hook( + user_api_key_dict=keyless, + cache=DualCache(), + data=data, + call_type=CallTypes.aretrieve_batch.value, + ) + assert data["batch_id"] == "batch_raw_123" + assert data["model"] == "my-model-id" + + +@pytest.mark.asyncio +async def test_other_keyless_key_still_denied_the_batch(): + unified_batch_id = _make_unified_batch_id() + managed_files, _ = _make_managed_files_instance_with_object_store() + + await _store_batch( + managed_files, + unified_batch_id, + UserAPIKeyAuth(api_key="sk-creator", parent_otel_span=None), + ) + + with pytest.raises(HTTPException) as exc_info: + await managed_files.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="sk-other", parent_otel_span=None), + cache=DualCache(), + data={"batch_id": unified_batch_id}, + call_type=CallTypes.aretrieve_batch.value, + ) + assert exc_info.value.status_code == 403 + + # --- Option C fix test: check_batch_cost bypasses managed files hook --- diff --git a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py b/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py index eddfc4fbd34..f3ad8a8592e 100644 --- a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py +++ b/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py @@ -527,13 +527,33 @@ async def test_afile_list_orders_newest_first_and_breaks_ties_on_the_cursor_colu @pytest.mark.asyncio -async def test_afile_list_denies_a_caller_without_a_user_or_team(): +async def test_afile_list_scopes_a_keyless_key_to_its_own_hashed_token(): + caller = UserAPIKeyAuth(api_key="sk-test", parent_otel_span=None) + managed_files, table = _make_managed_files_over_rows( + [ + _make_managed_file_row("unified-mine", created_by=f"key:{caller.token}"), + _make_managed_file_row("unified-theirs", created_by="other-user"), + ] + ) + + response = await managed_files.afile_list( + purpose=None, + litellm_parent_otel_span=None, + user_api_key_dict=caller, + ) + + assert [file.id for file in response.data] == ["unified-mine"] + assert table.find_many_calls[0]["where"] == {"created_by": f"key:{caller.token}"} + + +@pytest.mark.asyncio +async def test_afile_list_denies_a_caller_with_no_identity_at_all(): managed_files, table = _make_managed_files_over_rows([_make_managed_file_row("unified-mine")]) response = await managed_files.afile_list( purpose=None, litellm_parent_otel_span=None, - user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", parent_otel_span=None), + user_api_key_dict=UserAPIKeyAuth(parent_otel_span=None), ) assert response.data == [] diff --git a/tests/test_litellm/llms/base_llm/test_managed_resource_isolation.py b/tests/test_litellm/llms/base_llm/test_managed_resource_isolation.py index b5fcd9d8219..1746926c689 100644 --- a/tests/test_litellm/llms/base_llm/test_managed_resource_isolation.py +++ b/tests/test_litellm/llms/base_llm/test_managed_resource_isolation.py @@ -7,6 +7,7 @@ import pytest from litellm.llms.base_llm.managed_resources.isolation import ( build_owner_filter, can_access_resource, + resolve_resource_owner_id, ) from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth @@ -154,3 +155,46 @@ def test_access_identity_less_caller_always_denied(created_by, resource_team_id) ) is False ) + + +# --------------------------------------------------------------------------- +# keyless keys (no user_id, no team_id) own their resources by hashed token +# --------------------------------------------------------------------------- + + +def test_owner_id_prefers_user_id_then_falls_back_to_token(): + assert resolve_resource_owner_id(UserAPIKeyAuth(user_id="alice")) == "alice" + assert resolve_resource_owner_id(UserAPIKeyAuth(team_id="team-eng")) is None + assert resolve_resource_owner_id(UserAPIKeyAuth()) is None + + keyless = UserAPIKeyAuth(api_key="sk-keyless") + assert resolve_resource_owner_id(keyless) == f"key:{keyless.token}" + + +def test_keyless_key_can_access_its_own_resource(): + """Regression for the self-lockout: a key generated by a proxy admin (or a + service-account key) has no user_id and no team_id, so it used to stamp + `created_by=None` and then be denied its own batches and files.""" + keyless = UserAPIKeyAuth(api_key="sk-keyless") + owner_id = resolve_resource_owner_id(keyless) + + assert build_owner_filter(keyless) == {"created_by": owner_id} + assert ( + can_access_resource(keyless, created_by=owner_id, resource_team_id=None) is True + ) + + +def test_keyless_key_denied_another_keyless_keys_resource(): + """The #27004 isolation invariant: two distinct keyless keys must not see + each other's resources.""" + creator = UserAPIKeyAuth(api_key="sk-creator") + other = UserAPIKeyAuth(api_key="sk-other") + + assert ( + can_access_resource( + other, + created_by=resolve_resource_owner_id(creator), + resource_team_id=None, + ) + is False + ) diff --git a/tests/test_litellm/proxy/db/test_routing_prisma_wrapper.py b/tests/test_litellm/proxy/db/test_routing_prisma_wrapper.py index dcc0036ff04..966a638f6a4 100644 --- a/tests/test_litellm/proxy/db/test_routing_prisma_wrapper.py +++ b/tests/test_litellm/proxy/db/test_routing_prisma_wrapper.py @@ -101,6 +101,49 @@ def test_per_model_reads_route_to_reader_writes_to_writer(): assert actions.delete_many is writer_inner.litellm_usertable.delete_many +def test_writer_pinned_client_bypasses_reader_routing(): + """Regression for #38556: read-after-write reconciles must see the writer's + just-committed rows, so WriterPinnedClient must resolve reads to the writer + even when a read replica is configured.""" + from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper, WriterPinnedClient + + writer, writer_inner, reader, reader_inner = _make_wrappers() + writer_inner.litellm_proxymodeltable = _model_actions_mock("writer_models") + reader_inner.litellm_proxymodeltable = _model_actions_mock("reader_models") + routing = RoutingPrismaWrapper(writer=writer, reader=reader) + + pinned = WriterPinnedClient(routing) + + assert pinned.db is writer + assert pinned.db.litellm_proxymodeltable.find_many is writer_inner.litellm_proxymodeltable.find_many + + +def test_writer_pinned_client_passes_through_single_db(): + from litellm.proxy.db.routing_prisma_wrapper import WriterPinnedClient + + writer, _, _, _ = _make_wrappers() + + assert WriterPinnedClient(writer).db is writer + + +def test_writer_pinned_client_yields_to_routed_reads_when_writer_down(): + """The pin must not break reader-only degraded mode: a proxy that starts + during a primary outage still loads DB-backed models from the replica, so + while the writer is degraded the pin resolves to the routed wrapper.""" + from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper, WriterPinnedClient + + writer, writer_inner, reader, reader_inner = _make_wrappers() + writer_inner.litellm_proxymodeltable = _model_actions_mock("writer_models") + reader_inner.litellm_proxymodeltable = _model_actions_mock("reader_models") + routing = RoutingPrismaWrapper(writer=writer, reader=reader) + routing._writer_unavailable = True + + pinned = WriterPinnedClient(routing) + + assert pinned.db is routing + assert pinned.db.litellm_proxymodeltable.find_many is reader_inner.litellm_proxymodeltable.find_many + + @pytest.mark.asyncio async def test_connect_invokes_both_clients(): from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 70736949628..7e9145947f4 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -9566,6 +9566,76 @@ class TestDeleteDeploymentSync: assert result is None, f"Expected None on DB failure to signal fetch error, got {result!r}" + @pytest.mark.asyncio + async def test_get_models_from_db_reads_from_writer_not_replica(self): + """ + Regression for #38556: with DATABASE_URL_READ_REPLICA configured, the model + reconcile after /model/new used to read via the replica, so a lagging replica + made the reload miss the just-committed row and fail the request with a 500. + The reconcile read must be pinned to the writer. + """ + from unittest.mock import AsyncMock, MagicMock + + from litellm.proxy.db.prisma_client import PrismaWrapper + from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper + from litellm.proxy.proxy_server import ProxyConfig + + writer_inner = MagicMock(name="writer_prisma") + reader_inner = MagicMock(name="reader_prisma") + committed_row = MagicMock(name="just_committed_model_row") + writer_inner.litellm_proxymodeltable.find_many = AsyncMock(return_value=[committed_row]) + reader_inner.litellm_proxymodeltable.find_many = AsyncMock(return_value=[]) + + mock_prisma = MagicMock() + mock_prisma.db = RoutingPrismaWrapper( + writer=PrismaWrapper(original_prisma=writer_inner, iam_token_db_auth=False), + reader=PrismaWrapper(original_prisma=reader_inner, iam_token_db_auth=False), + ) + + result = await ProxyConfig()._get_models_from_db(prisma_client=mock_prisma) + + assert result == [committed_row], f"Expected the writer's just-committed row, got {result!r}" + reader_inner.litellm_proxymodeltable.find_many.assert_not_awaited() + + @pytest.mark.asyncio + async def test_get_models_from_db_falls_back_to_replica_when_writer_down(self): + """ + The writer pin must not break reader-only degraded mode: a proxy that + starts during a primary outage (writer connect failed, replica healthy) + must still load DB-backed models through the replica instead of sending + the reconcile read to the unavailable writer. + """ + from types import SimpleNamespace + from unittest.mock import AsyncMock, MagicMock + + from litellm.proxy.db.prisma_client import PrismaWrapper + from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper + from litellm.proxy.proxy_server import ProxyConfig + + writer_inner = MagicMock(name="writer_prisma") + reader_inner = MagicMock(name="reader_prisma") + replica_row = MagicMock(name="replica_model_row") + writer_inner.litellm_proxymodeltable = SimpleNamespace( + find_many=AsyncMock(side_effect=RuntimeError("writer unreachable")), + create=MagicMock(name="writer_create"), + ) + reader_inner.litellm_proxymodeltable = SimpleNamespace( + find_many=AsyncMock(return_value=[replica_row]), + create=MagicMock(name="reader_create"), + ) + + mock_prisma = MagicMock() + mock_prisma.db = RoutingPrismaWrapper( + writer=PrismaWrapper(original_prisma=writer_inner, iam_token_db_auth=False), + reader=PrismaWrapper(original_prisma=reader_inner, iam_token_db_auth=False), + ) + mock_prisma.db._writer_unavailable = True + + result = await ProxyConfig()._get_models_from_db(prisma_client=mock_prisma) + + assert result == [replica_row], f"Expected the replica's rows in degraded mode, got {result!r}" + writer_inner.litellm_proxymodeltable.find_many.assert_not_awaited() + def test_get_config_list_includes_cancel_on_disconnect(monkeypatch): """Follow-up to #30223: the flag must be discoverable via /config/list, diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py b/tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py index 0971ce09d79..9d2a27ce9d3 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py @@ -9,11 +9,14 @@ import pytest from fastapi import HTTPException import litellm +from litellm.caching.caching import DualCache from litellm.exceptions import RejectedRequestError from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger +from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.utils import ProxyLogging from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.utils import CallTypesLiteral def _load(module: str, name: str): @@ -473,7 +476,13 @@ class _RedactingGuardrail(CustomGuardrail): kwargs.setdefault("event_hook", GuardrailEventHooks.pre_call) super().__init__(guardrail_name="redactor", **kwargs) - async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): # type: ignore[override] + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: dict, + call_type: CallTypesLiteral, + ) -> dict | None: for msg in data.get("messages", []): if "SECRET" in msg.get("content", ""): msg["content"] = msg["content"].replace("SECRET", "[REDACTED]") @@ -488,7 +497,13 @@ class _BlockOnSecretGuardrail(CustomGuardrail): kwargs.setdefault("event_hook", GuardrailEventHooks.pre_call) super().__init__(guardrail_name="blocker", **kwargs) - async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): # type: ignore[override] + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: dict, + call_type: CallTypesLiteral, + ) -> dict | None: if any("SECRET" in msg.get("content", "") for msg in data.get("messages", [])): raise HTTPException(status_code=400, detail="blocked: SECRET detected") return None @@ -560,7 +575,13 @@ async def test_scan_raw_request_guardrail_does_not_undo_later_masking( separate marker (PII_TOKEN) that only the redactor reacts to.""" class _PiiRedactor(_RedactingGuardrail): - async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): # type: ignore[override] + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: dict, + call_type: CallTypesLiteral, + ) -> dict | None: for msg in data.get("messages", []): if "PII_TOKEN" in msg.get("content", ""): msg["content"] = msg["content"].replace("PII_TOKEN", "[REDACTED]") @@ -692,7 +713,13 @@ async def test_scan_raw_request_warns_when_guardrail_mutation_discarded( super().__init__(**kwargs) self.scan_raw_request = True - async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): # type: ignore[override] + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: dict, + call_type: CallTypesLiteral, + ) -> dict | None: for msg in data.get("messages", []): msg["content"] = msg["content"].replace("SECRET", "[REDACTED]") return data