scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 89

Merge remote-tracking branch 'origin/litellm_window_spend_schema' into litellm_window_spend_writer

Browse source
This commit is contained in:
ryan-crabbe-berri 2026-08-29 12:27:58 -07:00
commit 2fac72392a
19 changed files with 418 additions and 52 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
basedpyright-code-budget.json
View file

@ -1,6 +1,6 @@
{
"reportAny": {
"limit": 17271
"limit": 17270
},
"reportArgumentType": {
"limit": 2539

7
enterprise/litellm_enterprise/proxy/hooks/managed_files.py
View file

@ -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,

5
litellm/llms/base_llm/managed_resources/base_managed_resource.py
View file

@ -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,
}

49
litellm/llms/base_llm/managed_resources/isolation.py
View file

@ -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": <user_id>}`` for user-keyed callers.
- ``{"created_by": <owner_id>}`` for user-keyed callers, and for keys
with no user_id and no team_id (owner id is their hashed token).
- ``{"team_id": <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

4
litellm/main.py
View file

@ -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)

20
litellm/proxy/db/routing_prisma_wrapper.py
View file

@ -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.

10
litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py
View file

@ -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,
)

2
litellm/proxy/guardrails/guardrail_hooks/qualifire/qualifire.py
View file

@ -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)

17
litellm/proxy/guardrails/guardrail_registry.py
View file

@ -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.<param_name>`` 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)

3
litellm/proxy/pass_through_endpoints/managed_id_rewriter.py
View file

@ -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,
},

2
litellm/proxy/policy_engine/pipeline_executor.py
View file

@ -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

12
litellm/proxy/proxy_server.py
View file

@ -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 (

14
litellm/proxy/utils.py
View file

@ -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)

107
tests/test_litellm/enterprise/proxy/test_managed_files_access_check.py
View file

@ -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 ---

24
tests/test_litellm/enterprise/proxy/test_managed_files_hook.py
View file

@ -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 == []

44
tests/test_litellm/llms/base_llm/test_managed_resource_isolation.py
View file

@ -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
)

43
tests/test_litellm/proxy/db/test_routing_prisma_wrapper.py
View file

@ -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

70
tests/test_litellm/proxy/test_proxy_server.py
View file

@ -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,

35
tests/test_litellm/proxy/utils/proxy_logging/test_pre_call_hook.py
View file

@ -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