From 217cb7da653e025f7c6d10eaae8e291b826fd76f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 4 Sep 2026 20:30:33 -0700 Subject: [PATCH] fix(managed_files): resolve the creator org through the cached team lookup Batch creation snapshotted the team's organization with a direct litellm_teamtable query on every create. Go through get_team_object instead, which serves the team auth already cached and only falls back to the database when the team was never cached. --- .../proxy/hooks/managed_files.py | 13 +++- ..._batch_update_db_managed_output_file_id.py | 68 ++++++++++++++----- 2 files changed, 61 insertions(+), 20 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 850e5aadf0e..68782c5516d 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -288,11 +288,18 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): return user_api_key_dict.org_id if not user_api_key_dict.team_id: return None + from litellm.proxy.auth.auth_checks import get_team_object + from litellm.proxy.proxy_server import proxy_logging_obj, user_api_key_cache + try: - team_row = await self.prisma_client.db.litellm_teamtable.find_unique( - where={"team_id": user_api_key_dict.team_id} + team: Final = await get_team_object( + team_id=user_api_key_dict.team_id, + prisma_client=self.prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=user_api_key_dict.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, ) - return getattr(team_row, "organization_id", None) if team_row is not None else None + return team.organization_id except Exception as e: verbose_logger.warning(f"could not resolve org for managed object attribution: {e}") return None diff --git a/tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py b/tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py index b5865ab4a13..d3e668b8987 100644 --- a/tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py +++ b/tests/test_litellm/enterprise/proxy/test_batch_update_db_managed_output_file_id.py @@ -413,29 +413,63 @@ async def test_store_unified_object_id_persists_key_and_tags_on_create(): @pytest.mark.asyncio -async def test_store_unified_object_id_resolves_org_through_the_team(): +async def test_store_unified_object_id_resolves_org_through_the_cached_team(): """Most keys belong to an org only through their team, so the auth object carries no - org_id. The create resolves the team's organization so org spend is snapshotted at - submission time instead of never being billed.""" - from types import SimpleNamespace + org_id. The create reads the team that auth already cached, so org spend is snapshotted + at submission time without a database query in the request path.""" + from litellm.models.team import LiteLLM_TeamTableCachedObj + from litellm.proxy.proxy_server import user_api_key_cache + + instance, store = _in_memory_managed_files() + creator = UserAPIKeyAuth(user_id="alice", team_id="team-cached", api_key="hash-alice") + await user_api_key_cache.async_set_cache( + key="team_id:team-cached", + value=LiteLLM_TeamTableCachedObj(team_id="team-cached", organization_id="org-via-team"), + model_type=LiteLLM_TeamTableCachedObj, + ) + try: + await instance.store_unified_object_id( + unified_object_id="unified-b", + file_object=_build_batch_response(batch_id="b", status="validating"), + litellm_parent_otel_span=None, + model_object_id="b", + file_purpose="batch", + user_api_key_dict=creator, + persist_attribution=True, + ) + finally: + user_api_key_cache.delete_cache(key="team_id:team-cached") + + assert store["unified-b"]["org_id"] == "org-via-team" + instance.prisma_client.db.litellm_teamtable.find_unique.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_store_unified_object_id_resolves_org_from_the_db_when_the_team_is_not_cached(): + """A team no request has run under yet is absent from the auth cache; its organization + still comes back from the table so the org is billed rather than dropped.""" + from litellm.models.team import LiteLLM_TeamTable + from litellm.proxy.proxy_server import user_api_key_cache instance, store = _in_memory_managed_files() instance.prisma_client.db.litellm_teamtable.find_unique = AsyncMock( - return_value=SimpleNamespace(organization_id="org-via-team") + return_value=LiteLLM_TeamTable(team_id="team-uncached", organization_id="org-via-db") ) - creator = UserAPIKeyAuth(user_id="alice", team_id="team-alpha", api_key="hash-alice") + creator = UserAPIKeyAuth(user_id="alice", team_id="team-uncached", api_key="hash-alice") + try: + await instance.store_unified_object_id( + unified_object_id="unified-b", + file_object=_build_batch_response(batch_id="b", status="validating"), + litellm_parent_otel_span=None, + model_object_id="b", + file_purpose="batch", + user_api_key_dict=creator, + persist_attribution=True, + ) + finally: + user_api_key_cache.delete_cache(key="team_id:team-uncached") - await instance.store_unified_object_id( - unified_object_id="unified-b", - file_object=_build_batch_response(batch_id="b", status="validating"), - litellm_parent_otel_span=None, - model_object_id="b", - file_purpose="batch", - user_api_key_dict=creator, - persist_attribution=True, - ) - - assert store["unified-b"]["org_id"] == "org-via-team" + assert store["unified-b"]["org_id"] == "org-via-db" @pytest.mark.asyncio