From 846dbecbf2bcf3e11c60f56971c16c21f13f40fd Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 30 Jun 2026 19:02:00 -0700 Subject: [PATCH] feat(proxy): support object_permission in default_key_generate_params (#31776) * feat(proxy): support object_permission in default_key_generate_params default_key_generate_params filled in a fixed whitelist of scalar fields plus a full-replace for models/metadata, but never touched object_permission, so admins had no way to set a default (e.g. mcp_tool_search_enabled, vector_stores) applied to every new key. Merge object_permission field-by-field instead of replacing it wholesale, so a caller-supplied field (e.g. mcp_servers) is preserved alongside defaulted fields the caller left unset. * ci: retrigger proxy_pass_through_endpoint_tests (suspected flake, unrelated to this PR's diff) * fix(proxy): apply default object_permission after team-scope validation Injecting the default before validate_key_vector_stores_against_team / validate_key_search_tools_against_team ran meant a default containing a team-scoped field (e.g. vector_stores) looked like a caller-requested permission, turning ordinary non-admin personal key creation into a 403. Merge the default into data_json after those checks instead, and guard against a non-dict default value. --- .../key_management_endpoints.py | 17 ++ .../test_key_management_endpoints.py | 258 ++++++++++++++++++ 2 files changed, 275 insertions(+) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 4106eae606c..523cb9b74e5 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -965,6 +965,23 @@ async def _common_key_generation_helper( is_proxy_admin=_is_proxy_admin_caller, ) + # Merge default_key_generate_params.object_permission in *after* the team-scope + # checks above, so an admin-configured default (e.g. vector_stores, search_tools) + # is never mistaken for a caller-requested permission and rejected by those + # non-admin/no-team checks. Only fields the caller left unset are filled in. + _default_object_permission = ( + litellm.default_key_generate_params.get("object_permission") + if litellm.default_key_generate_params is not None + else None + ) + if isinstance(_default_object_permission, dict): + _caller_object_permission = data_json.get("object_permission") + if _caller_object_permission is None: + data_json["object_permission"] = dict(_default_object_permission) + elif isinstance(_caller_object_permission, dict): + for _op_field, _op_default_value in _default_object_permission.items(): + _caller_object_permission.setdefault(_op_field, _op_default_value) + data_json = await _set_object_permission( data_json=data_json, prisma_client=prisma_client, diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 04048020e18..ae62be8799f 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -7781,6 +7781,264 @@ async def test_default_key_generate_params_duration(monkeypatch): litellm.default_key_generate_params = original_value +async def test_default_key_generate_params_object_permission_applied_when_absent( + monkeypatch, +): + """ + default_key_generate_params.object_permission is applied to a key that + doesn't specify object_permission at all. + """ + import litellm + + mock_prisma_client = AsyncMock() + mock_insert_data = AsyncMock( + return_value=MagicMock( + token="hashed_token_123", litellm_budget_table=None, object_permission=None + ) + ) + mock_prisma_client.insert_data = mock_insert_data + mock_prisma_client.db = MagicMock() + mock_prisma_client.db.litellm_verificationtoken = MagicMock() + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( + return_value=None + ) + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock( + return_value=[] + ) + mock_prisma_client.db.litellm_verificationtoken.count = AsyncMock(return_value=0) + mock_prisma_client.db.litellm_verificationtoken.update = AsyncMock( + return_value=MagicMock( + token="hashed_token_123", litellm_budget_table=None, object_permission=None + ) + ) + mock_prisma_client.db.litellm_objectpermissiontable = MagicMock() + mock_prisma_client.db.litellm_objectpermissiontable.create = AsyncMock( + return_value=MagicMock(object_permission_id="objperm-1") + ) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + + original_value = litellm.default_key_generate_params + litellm.default_key_generate_params = { + "object_permission": {"vector_stores": ["default-vs"]} + } + + try: + request = GenerateKeyRequest() # No object_permission specified + await _common_key_generation_helper( + data=request, + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-1234", + user_id="1234", + ), + litellm_changed_by=None, + team_table=None, + ) + + created_data = mock_prisma_client.db.litellm_objectpermissiontable.create.call_args.kwargs["data"] + assert created_data["vector_stores"] == ["default-vs"] + finally: + litellm.default_key_generate_params = original_value + + +async def test_default_key_generate_params_object_permission_merges_partial( + monkeypatch, +): + """ + default_key_generate_params.object_permission fills only the fields the + caller left unset - an explicitly supplied field (agents here) is + preserved alongside the defaulted field (vector_stores). + """ + import litellm + from litellm.proxy._types import LiteLLM_ObjectPermissionBase + + mock_prisma_client = AsyncMock() + mock_insert_data = AsyncMock( + return_value=MagicMock( + token="hashed_token_123", litellm_budget_table=None, object_permission=None + ) + ) + mock_prisma_client.insert_data = mock_insert_data + mock_prisma_client.db = MagicMock() + mock_prisma_client.db.litellm_verificationtoken = MagicMock() + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( + return_value=None + ) + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock( + return_value=[] + ) + mock_prisma_client.db.litellm_verificationtoken.count = AsyncMock(return_value=0) + mock_prisma_client.db.litellm_verificationtoken.update = AsyncMock( + return_value=MagicMock( + token="hashed_token_123", litellm_budget_table=None, object_permission=None + ) + ) + mock_prisma_client.db.litellm_objectpermissiontable = MagicMock() + mock_prisma_client.db.litellm_objectpermissiontable.create = AsyncMock( + return_value=MagicMock(object_permission_id="objperm-2") + ) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + + original_value = litellm.default_key_generate_params + litellm.default_key_generate_params = { + "object_permission": {"vector_stores": ["default-vs"]} + } + + try: + request = GenerateKeyRequest( + object_permission=LiteLLM_ObjectPermissionBase(agents=["agent-1"]) + ) + await _common_key_generation_helper( + data=request, + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-1234", + user_id="1234", + ), + litellm_changed_by=None, + team_table=None, + ) + + created_data = mock_prisma_client.db.litellm_objectpermissiontable.create.call_args.kwargs["data"] + assert created_data["agents"] == ["agent-1"] + assert created_data["vector_stores"] == ["default-vs"] + finally: + litellm.default_key_generate_params = original_value + + +async def test_default_key_generate_params_object_permission_does_not_override_explicit( + monkeypatch, +): + """ + A field the caller explicitly set on object_permission must win over the + same field in default_key_generate_params. + """ + import litellm + from litellm.proxy._types import LiteLLM_ObjectPermissionBase + + mock_prisma_client = AsyncMock() + mock_insert_data = AsyncMock( + return_value=MagicMock( + token="hashed_token_123", litellm_budget_table=None, object_permission=None + ) + ) + mock_prisma_client.insert_data = mock_insert_data + mock_prisma_client.db = MagicMock() + mock_prisma_client.db.litellm_verificationtoken = MagicMock() + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( + return_value=None + ) + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock( + return_value=[] + ) + mock_prisma_client.db.litellm_verificationtoken.count = AsyncMock(return_value=0) + mock_prisma_client.db.litellm_verificationtoken.update = AsyncMock( + return_value=MagicMock( + token="hashed_token_123", litellm_budget_table=None, object_permission=None + ) + ) + mock_prisma_client.db.litellm_objectpermissiontable = MagicMock() + mock_prisma_client.db.litellm_objectpermissiontable.create = AsyncMock( + return_value=MagicMock(object_permission_id="objperm-3") + ) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + + original_value = litellm.default_key_generate_params + litellm.default_key_generate_params = { + "object_permission": {"vector_stores": ["default-vs"]} + } + + try: + request = GenerateKeyRequest( + object_permission=LiteLLM_ObjectPermissionBase( + vector_stores=["explicit-vs"] + ) + ) + await _common_key_generation_helper( + data=request, + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-1234", + user_id="1234", + ), + litellm_changed_by=None, + team_table=None, + ) + + created_data = mock_prisma_client.db.litellm_objectpermissiontable.create.call_args.kwargs["data"] + assert created_data["vector_stores"] == ["explicit-vs"] + finally: + litellm.default_key_generate_params = original_value + + +async def test_default_key_generate_params_object_permission_not_rejected_for_non_admin_personal_key( + monkeypatch, +): + """ + Regression test: a default_key_generate_params.object_permission containing + a team-scoped field (vector_stores) must not turn ordinary non-admin + personal key creation into a 403. The default is merged in *after* the + caller-scope validation, so it is never mistaken for a caller-requested + permission. + """ + import litellm + + mock_prisma_client = AsyncMock() + mock_insert_data = AsyncMock( + return_value=MagicMock( + token="hashed_token_123", litellm_budget_table=None, object_permission=None + ) + ) + mock_prisma_client.insert_data = mock_insert_data + mock_prisma_client.db = MagicMock() + mock_prisma_client.db.litellm_verificationtoken = MagicMock() + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( + return_value=None + ) + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock( + return_value=[] + ) + mock_prisma_client.db.litellm_verificationtoken.count = AsyncMock(return_value=0) + mock_prisma_client.db.litellm_verificationtoken.update = AsyncMock( + return_value=MagicMock( + token="hashed_token_123", litellm_budget_table=None, object_permission=None + ) + ) + mock_prisma_client.db.litellm_objectpermissiontable = MagicMock() + mock_prisma_client.db.litellm_objectpermissiontable.create = AsyncMock( + return_value=MagicMock(object_permission_id="objperm-4") + ) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + + original_value = litellm.default_key_generate_params + litellm.default_key_generate_params = { + "object_permission": {"vector_stores": ["default-vs"]} + } + + try: + request = GenerateKeyRequest(user_id="alice") # No object_permission specified + response = await _common_key_generation_helper( + data=request, + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-alice", + user_id="alice", + ), + litellm_changed_by=None, + team_table=None, + ) + + assert response is not None + created_data = mock_prisma_client.db.litellm_objectpermissiontable.create.call_args.kwargs["data"] + assert created_data["vector_stores"] == ["default-vs"] + finally: + litellm.default_key_generate_params = original_value + + @pytest.mark.asyncio async def test_build_key_filter_member_team_service_accounts(): """