scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-06 02:48:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

feat(proxy): support object_permission in default_key_generate_params (#31776)

Browse source 
* 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.
This commit is contained in:
Krrish Dholakia 2026-06-30 19:02:00 -07:00 • committed by GitHub
parent 6c21029cb7
commit 846dbecbf2
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 275 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

17
litellm/proxy/management_endpoints/key_management_endpoints.py
View file

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

258
tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py
View file

@ -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():
"""