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/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/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 + )