mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
fix(auth): add missing await in get_org_object() cache lookup
`get_org_object()` at auth_checks.py:2440 was calling `user_api_key_cache.async_get_cache()` without `await`, returning a coroutine object instead of the cached value. Since a coroutine is always truthy, `cached_org_obj is not None` was always True, but the subsequent `isinstance` checks against dict/LiteLLM_OrganizationTable always failed, causing every org lookup to fall through to the database. This is the only call site missing `await` — all 14 other calls to `async_get_cache` in the same file correctly use `await`. Added a regression test that verifies cache hits return the org object directly without hitting the database. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This commit is contained in:
parent
d1df4e838b
commit
46390b5084
2 changed files with 39 additions and 1 deletions
|
|
@ -2437,7 +2437,7 @@ async def get_org_object(
|
|||
cache_key = "org_id:{}:with_budget".format(org_id)
|
||||
|
||||
# check if in cache
|
||||
cached_org_obj = user_api_key_cache.async_get_cache(key=cache_key)
|
||||
cached_org_obj = await user_api_key_cache.async_get_cache(key=cache_key)
|
||||
if cached_org_obj is not None:
|
||||
if isinstance(cached_org_obj, dict):
|
||||
return LiteLLM_OrganizationTable(**cached_org_obj)
|
||||
|
|
|
|||
|
|
@ -38,6 +38,7 @@ from litellm.proxy.auth.auth_checks import (
|
|||
_virtual_key_max_budget_check,
|
||||
_virtual_key_soft_budget_check,
|
||||
get_key_object,
|
||||
get_org_object,
|
||||
get_user_object,
|
||||
vector_store_access_check,
|
||||
)
|
||||
|
|
@ -1780,3 +1781,40 @@ async def test_team_member_budget_check_reads_from_spend_counter():
|
|||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
assert exc_info.value.current_cost == 1.5
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_org_object_returns_cached_org():
|
||||
"""
|
||||
Verify that get_org_object() correctly returns a cached organization object
|
||||
instead of falling through to a DB query. Regression test for a missing
|
||||
`await` on async_get_cache that caused the cache to be bypassed on every call.
|
||||
"""
|
||||
from litellm.proxy._types import LiteLLM_OrganizationTable
|
||||
|
||||
cached_org = LiteLLM_OrganizationTable(
|
||||
organization_id="test-org-123",
|
||||
organization_alias="Test Org",
|
||||
budget_id="budget-1",
|
||||
models=["gpt-4"],
|
||||
created_by="admin",
|
||||
updated_by="admin",
|
||||
)
|
||||
|
||||
mock_cache = AsyncMock()
|
||||
mock_cache.async_get_cache = AsyncMock(return_value=cached_org)
|
||||
|
||||
mock_prisma = MagicMock()
|
||||
# If cache works, find_unique should never be called
|
||||
mock_prisma.db.litellm_organizationtable.find_unique = AsyncMock()
|
||||
|
||||
result = await get_org_object(
|
||||
org_id="test-org-123",
|
||||
prisma_client=mock_prisma,
|
||||
user_api_key_cache=mock_cache,
|
||||
)
|
||||
|
||||
assert result is not None
|
||||
assert result.organization_id == "test-org-123"
|
||||
mock_cache.async_get_cache.assert_awaited_once()
|
||||
mock_prisma.db.litellm_organizationtable.find_unique.assert_not_awaited()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue