mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-11 03:38:38 +00:00
fix(proxy): reject key project ownership before permission and budget writes
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yucheng
parent
1d24f7fdbc
commit
a27bb1088e
5 changed files with 462 additions and 242 deletions
|
|
@ -790,8 +790,16 @@ async def update_project(
|
|||
data, existing_project, _router_access_group_names(llm_router)
|
||||
)
|
||||
|
||||
object_permission_data: Final = (
|
||||
data.object_permission.model_dump(exclude_none=True) if data.object_permission is not None else None
|
||||
)
|
||||
object_permission_payload: Final = (
|
||||
_OBJECT_PERMISSION_PAYLOAD.validate_python(object_permission_data) if object_permission_data else None
|
||||
)
|
||||
|
||||
# Prepare update data
|
||||
update_data = _jsonified(prisma_client, data.model_dump(exclude_none=True, exclude={"project_id"}))
|
||||
update_data.pop("object_permission", None)
|
||||
update_data["updated_by"] = user_api_key_dict.user_id or litellm_proxy_admin_name
|
||||
|
||||
# Handle budget updates
|
||||
|
|
@ -801,48 +809,6 @@ async def update_project(
|
|||
**({"max_budget": None} if "max_budget" in data.model_fields_set and data.max_budget is None else {}),
|
||||
}
|
||||
|
||||
if budget_updates and existing_project.budget_id:
|
||||
# Update existing budget
|
||||
await _budget_table(prisma_client).update(
|
||||
where={"budget_id": existing_project.budget_id},
|
||||
data={
|
||||
**budget_updates,
|
||||
"updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name,
|
||||
},
|
||||
)
|
||||
# Remove budget fields from project update
|
||||
for field in budget_updates.keys():
|
||||
update_data.pop(field, None)
|
||||
|
||||
# Handle object permissions
|
||||
if "object_permission" in update_data:
|
||||
object_permission_data = update_data.pop("object_permission")
|
||||
if object_permission_data:
|
||||
object_permission_payload: Final = _OBJECT_PERMISSION_PAYLOAD.validate_python(object_permission_data)
|
||||
if existing_project.object_permission_id:
|
||||
# Update existing permission
|
||||
await _object_permission_table(prisma_client).update(
|
||||
where={"object_permission_id": existing_project.object_permission_id},
|
||||
data=object_permission_payload,
|
||||
)
|
||||
else:
|
||||
# Create new permission
|
||||
created_permission: Final = await _object_permission_table(prisma_client).create(
|
||||
data=object_permission_payload,
|
||||
)
|
||||
update_data["object_permission_id"] = created_permission.object_permission_id
|
||||
|
||||
# Handle metadata fields
|
||||
for field in LiteLLM_ManagementEndpoint_MetadataFields:
|
||||
if field in update_data:
|
||||
existing_metadata = update_data.get("metadata")
|
||||
metadata_dict: dict[str, object] = existing_metadata if isinstance(existing_metadata, dict) else {}
|
||||
metadata_dict[field] = update_data.pop(field)
|
||||
update_data["metadata"] = metadata_dict
|
||||
|
||||
# Remove budget fields (following organization_endpoints.py pattern)
|
||||
update_data = _remove_budget_fields_from_project_data(update_data)
|
||||
|
||||
if data.team_id is not None:
|
||||
current_project_record: Final = await _writer_project_table(prisma_client).find_unique(
|
||||
where={"project_id": data.project_id}
|
||||
|
|
@ -870,6 +836,42 @@ async def update_project(
|
|||
},
|
||||
)
|
||||
|
||||
if budget_updates and existing_project.budget_id:
|
||||
# Update existing budget
|
||||
await _budget_table(prisma_client).update(
|
||||
where={"budget_id": existing_project.budget_id},
|
||||
data={
|
||||
**budget_updates,
|
||||
"updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name,
|
||||
},
|
||||
)
|
||||
# Remove budget fields from project update
|
||||
for field in budget_updates.keys():
|
||||
update_data.pop(field, None)
|
||||
|
||||
if object_permission_payload is not None:
|
||||
if existing_project.object_permission_id:
|
||||
await _object_permission_table(prisma_client).update(
|
||||
where={"object_permission_id": existing_project.object_permission_id},
|
||||
data=object_permission_payload,
|
||||
)
|
||||
else:
|
||||
created_permission: Final = await _object_permission_table(prisma_client).create(
|
||||
data=object_permission_payload,
|
||||
)
|
||||
update_data["object_permission_id"] = created_permission.object_permission_id
|
||||
|
||||
# Handle metadata fields
|
||||
for field in LiteLLM_ManagementEndpoint_MetadataFields:
|
||||
if field in update_data:
|
||||
existing_metadata = update_data.get("metadata")
|
||||
metadata_dict: dict[str, object] = existing_metadata if isinstance(existing_metadata, dict) else {}
|
||||
metadata_dict[field] = update_data.pop(field)
|
||||
update_data["metadata"] = metadata_dict
|
||||
|
||||
# Remove budget fields (following organization_endpoints.py pattern)
|
||||
update_data = _remove_budget_fields_from_project_data(update_data)
|
||||
|
||||
# Update project
|
||||
updated_project: Final = await _project_table(prisma_client).update(
|
||||
where={"project_id": data.project_id},
|
||||
|
|
|
|||
|
|
@ -113,10 +113,11 @@ from litellm.proxy.management_helpers.access_group_key_sync import (
|
|||
)
|
||||
from litellm.proxy.management_helpers.key_settings_audit import with_settings_updated_at
|
||||
from litellm.proxy.management_helpers.object_permission_utils import (
|
||||
ObjectPermissionUpsert,
|
||||
_set_object_permission,
|
||||
attach_object_permission_to_dict,
|
||||
handle_update_object_permission_common,
|
||||
invalidate_cached_object_permissions,
|
||||
prepare_object_permission_upsert,
|
||||
validate_key_mcp_servers_against_team,
|
||||
validate_key_search_tools_against_team,
|
||||
validate_key_vector_stores_against_team,
|
||||
|
|
@ -139,6 +140,7 @@ from litellm.repositories.budget_repository import BudgetRepository
|
|||
from litellm.repositories.config_repository import ConfigParam, ConfigRepository
|
||||
from litellm.repositories.credentials_repository import CredentialsRepository
|
||||
from litellm.repositories.model_repository import ModelRepository
|
||||
from litellm.repositories.object_permission_repository import ObjectPermissionRepository
|
||||
from litellm.repositories.prisma_protocols import TableActions
|
||||
from litellm.repositories.table_repositories import (
|
||||
DeletedVerificationTokenRepository,
|
||||
|
|
@ -2502,11 +2504,6 @@ async def _update_key_row_with_soft_budget(
|
|||
existing_key_row=existing_key_row,
|
||||
changed_by=changed_by,
|
||||
)
|
||||
await _check_key_project_team_on_mutation(
|
||||
data=data,
|
||||
existing_key_row=existing_key_row,
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
include_object_permission: Final[prisma.types.LiteLLM_VerificationTokenInclude] = {"object_permission": True}
|
||||
updated_row: Final = await tx.litellm_verificationtoken.update(
|
||||
where=key_where,
|
||||
|
|
@ -2522,19 +2519,12 @@ async def _update_key_row_with_soft_budget(
|
|||
return result
|
||||
|
||||
|
||||
async def _update_key_row_with_project_team_check(
|
||||
async def _update_key_row(
|
||||
prisma_client: PrismaClient,
|
||||
key: str,
|
||||
data: UpdateKeyRequest,
|
||||
update_values: Mapping[str, object],
|
||||
existing_key_row: LiteLLM_VerificationToken,
|
||||
) -> _KeyUpdateResult | None:
|
||||
key_update_data: Final = MappingProxyType({**update_values, "token": key})
|
||||
await _check_key_project_team_on_mutation(
|
||||
data=data,
|
||||
existing_key_row=existing_key_row,
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
response: Final = await prisma_client.update_data(token=key, data=key_update_data)
|
||||
if response is None:
|
||||
return None
|
||||
|
|
@ -2647,27 +2637,42 @@ async def prepare_key_update_data(
|
|||
return non_default_values
|
||||
|
||||
|
||||
async def _handle_update_object_permission(
|
||||
data_json: dict,
|
||||
async def _prepare_key_update_object_permission(
|
||||
object_permission_data: object,
|
||||
existing_key_row: LiteLLM_VerificationToken,
|
||||
prisma_client: PrismaClient,
|
||||
) -> dict:
|
||||
"""Persist the requested object permission row and swap it for its id, only after the key policy allowed the write."""
|
||||
if "object_permission" not in data_json:
|
||||
return data_json
|
||||
) -> ObjectPermissionUpsert | None:
|
||||
if object_permission_data is None:
|
||||
return None
|
||||
|
||||
object_permission_id: Final = await handle_update_object_permission_common(
|
||||
data_json=data_json,
|
||||
parsed_object_permission: Final[object] = (
|
||||
json.loads(object_permission_data) if isinstance(object_permission_data, str) else object_permission_data
|
||||
)
|
||||
permission_data: Final[dict[str, object]] = (
|
||||
TypeAdapter(dict[str, object]).validate_python(parsed_object_permission)
|
||||
if isinstance(parsed_object_permission, dict)
|
||||
else {}
|
||||
)
|
||||
return await prepare_object_permission_upsert(
|
||||
new_object_permission=permission_data,
|
||||
existing_object_permission_id=existing_key_row.object_permission_id,
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
|
||||
# Add the object_permission_id to data_json if one was created/updated
|
||||
if object_permission_id is not None:
|
||||
data_json["object_permission_id"] = object_permission_id
|
||||
verbose_proxy_logger.debug("updated object_permission_id: %s", object_permission_id)
|
||||
|
||||
return data_json
|
||||
async def _write_prepared_key_update_object_permission(
|
||||
data_json: Mapping[str, object],
|
||||
upsert: ObjectPermissionUpsert | None,
|
||||
prisma_client: PrismaClient,
|
||||
) -> Mapping[str, object]:
|
||||
if upsert is None:
|
||||
return data_json
|
||||
|
||||
await ObjectPermissionRepository(prisma_client).table.upsert(
|
||||
where={"object_permission_id": upsert.object_permission_id},
|
||||
data={"create": upsert.record, "update": upsert.record},
|
||||
)
|
||||
return MappingProxyType({**data_json, "object_permission_id": upsert.object_permission_id})
|
||||
|
||||
|
||||
def is_different_team(data: UpdateKeyRequest, existing_key_row: LiteLLM_VerificationToken) -> bool:
|
||||
|
|
@ -2949,17 +2954,25 @@ async def _process_single_key_update(
|
|||
detail={"error": "Database not connected"},
|
||||
)
|
||||
|
||||
update_values: Final = await _handle_update_object_permission(
|
||||
data_json=non_default_values,
|
||||
object_permission_upsert: Final = await _prepare_key_update_object_permission(
|
||||
object_permission_data=non_default_values.get("object_permission"),
|
||||
existing_key_row=existing_key_row,
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
_data: Final = {**update_values, "token": key_request.key}
|
||||
key_update_values: Final = MappingProxyType(
|
||||
{field: value for field, value in non_default_values.items() if field != "object_permission"}
|
||||
)
|
||||
await _check_key_project_team_on_mutation(
|
||||
data=key_request,
|
||||
existing_key_row=existing_key_row,
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
update_values: Final = await _write_prepared_key_update_object_permission(
|
||||
data_json=key_update_values,
|
||||
upsert=object_permission_upsert,
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
_data: Final = {**update_values, "token": key_request.key}
|
||||
response: Final[Mapping[str, object] | None] = cast( # cast-ok: every update_data branch returns a str-keyed dict
|
||||
"Mapping[str, object] | None",
|
||||
await prisma_client.update_data(token=key_request.key, data=_data),
|
||||
|
|
@ -3627,11 +3640,24 @@ async def update_key_fn(
|
|||
if prisma_client is None:
|
||||
raise Exception("Not connected to DB!")
|
||||
|
||||
update_values: Final = await _handle_update_object_permission(
|
||||
data_json=non_default_values,
|
||||
object_permission_upsert: Final = await _prepare_key_update_object_permission(
|
||||
object_permission_data=non_default_values.get("object_permission"),
|
||||
existing_key_row=existing_key_row,
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
key_update_values: Final = MappingProxyType(
|
||||
{field: value for field, value in non_default_values.items() if field != "object_permission"}
|
||||
)
|
||||
await _check_key_project_team_on_mutation(
|
||||
data=data,
|
||||
existing_key_row=existing_key_row,
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
update_values: Final = await _write_prepared_key_update_object_permission(
|
||||
data_json=key_update_values,
|
||||
upsert=object_permission_upsert,
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
changed_by: Final = user_api_key_dict.user_id or litellm_proxy_admin_name
|
||||
response: Final = (
|
||||
await _update_key_row_with_soft_budget(
|
||||
|
|
@ -3643,12 +3669,10 @@ async def update_key_fn(
|
|||
changed_by=changed_by,
|
||||
)
|
||||
if "soft_budget" in data.model_fields_set
|
||||
else await _update_key_row_with_project_team_check(
|
||||
else await _update_key_row(
|
||||
prisma_client=prisma_client,
|
||||
key=key,
|
||||
data=data,
|
||||
update_values=update_values,
|
||||
existing_key_row=existing_key_row,
|
||||
)
|
||||
)
|
||||
|
||||
|
|
@ -3657,7 +3681,7 @@ async def update_key_fn(
|
|||
await invalidate_cached_object_permissions(
|
||||
object_permission_ids=(
|
||||
existing_key_row.object_permission_id,
|
||||
non_default_values.get("object_permission_id"),
|
||||
update_values.get("object_permission_id"),
|
||||
),
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
)
|
||||
|
|
@ -5710,7 +5734,6 @@ async def _execute_virtual_key_regeneration(
|
|||
new_token: Final = await get_new_token(data=data)
|
||||
new_token_hash: Final = hash_token(new_token)
|
||||
new_token_key_name: Final = abbreviate_api_key(api_key=new_token)
|
||||
update_data = {"token": new_token_hash, "key_name": new_token_key_name}
|
||||
|
||||
non_default_values = {}
|
||||
if data is not None:
|
||||
|
|
@ -5736,18 +5759,28 @@ async def _execute_virtual_key_regeneration(
|
|||
request=data if data is not None else RegenerateKeyRequest(),
|
||||
),
|
||||
)
|
||||
update_values: Final = await _handle_update_object_permission(
|
||||
data_json=non_default_values,
|
||||
object_permission_upsert: Final = await _prepare_key_update_object_permission(
|
||||
object_permission_data=non_default_values.get("object_permission"),
|
||||
existing_key_row=key_in_db,
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
update_data.update(update_values)
|
||||
key_update_values: Final = MappingProxyType(
|
||||
{field: value for field, value in non_default_values.items() if field != "object_permission"}
|
||||
)
|
||||
if data is not None:
|
||||
await _check_key_project_team_on_mutation(
|
||||
data=data,
|
||||
existing_key_row=key_in_db,
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
update_values: Final = await _write_prepared_key_update_object_permission(
|
||||
data_json=key_update_values,
|
||||
upsert=object_permission_upsert,
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
update_data: Final = MappingProxyType(
|
||||
{"token": new_token_hash, "key_name": new_token_key_name, **update_values}
|
||||
)
|
||||
|
||||
jsonified_update_data: Final[Mapping[str, object]] = prisma_client.jsonify_object(data=update_data)
|
||||
|
||||
|
|
|
|||
|
|
@ -33,6 +33,49 @@ def _key_rows(key: str) -> list[dict[str, JsonValue]]:
|
|||
)
|
||||
|
||||
|
||||
def _key_permission_and_budget_ids(key: str) -> list[dict[str, JsonValue]]:
|
||||
return read_rows(
|
||||
'SELECT object_permission_id, budget_id FROM "LiteLLM_VerificationToken" WHERE token = %s',
|
||||
(sha256(key.encode()).hexdigest(),),
|
||||
)
|
||||
|
||||
|
||||
def _object_permission_rows(permission_id: str) -> list[dict[str, JsonValue]]:
|
||||
return read_rows(
|
||||
'SELECT * FROM "LiteLLM_ObjectPermissionTable" WHERE object_permission_id = %s',
|
||||
(permission_id,),
|
||||
)
|
||||
|
||||
|
||||
def _budget_rows(budget_id: str) -> list[dict[str, JsonValue]]:
|
||||
return read_rows(
|
||||
'SELECT to_jsonb(b) AS row FROM "LiteLLM_BudgetTable" AS b WHERE budget_id = %s',
|
||||
(budget_id,),
|
||||
)
|
||||
|
||||
|
||||
def _clear_key_object_permission(key: str, permission_id: str) -> None:
|
||||
write_rows(
|
||||
'UPDATE "LiteLLM_VerificationToken" SET object_permission_id = NULL WHERE token = %s',
|
||||
(sha256(key.encode()).hexdigest(),),
|
||||
)
|
||||
write_rows(
|
||||
'DELETE FROM "LiteLLM_ObjectPermissionTable" WHERE object_permission_id = %s',
|
||||
(permission_id,),
|
||||
)
|
||||
|
||||
|
||||
def _clear_project_object_permission(project_id: str, permission_id: str) -> None:
|
||||
write_rows(
|
||||
'UPDATE "LiteLLM_ProjectTable" SET object_permission_id = NULL WHERE project_id = %s',
|
||||
(project_id,),
|
||||
)
|
||||
write_rows(
|
||||
'DELETE FROM "LiteLLM_ObjectPermissionTable" WHERE object_permission_id = %s',
|
||||
(permission_id,),
|
||||
)
|
||||
|
||||
|
||||
def _deleted_key_rows(key: str) -> list[dict[str, JsonValue]]:
|
||||
return read_rows(
|
||||
'SELECT token FROM "LiteLLM_DeletedVerificationToken" WHERE token = %s',
|
||||
|
|
@ -253,7 +296,22 @@ def test_key_update_rejects_team_change_and_allows_unchanged_values(ownership_ga
|
|||
team_a: Final = scenario.team(models=[model])
|
||||
team_b: Final = scenario.team(models=[model])
|
||||
project_a: Final = scenario.project(team_a, models=[model])
|
||||
key: Final = scenario.key(team_id=team_a, project_id=project_a, models=[model])
|
||||
key: Final = scenario.key(team_id=team_a, project_id=project_a, models=[model], soft_budget=3.0)
|
||||
permission_seed: Final = ownership_gateway.request(
|
||||
"POST",
|
||||
"/key/update",
|
||||
{"key": key, "object_permission": {"vector_stores": ["existing-store"]}},
|
||||
)
|
||||
assert permission_seed.status_code == 200, permission_seed.text
|
||||
key_permission_state: Final = _key_permission_and_budget_ids(key)
|
||||
assert len(key_permission_state) == 1
|
||||
permission_id: Final = string_value(key_permission_state[0]["object_permission_id"])
|
||||
budget_id: Final = string_value(key_permission_state[0]["budget_id"])
|
||||
scenario.cleanups.callback(_clear_key_object_permission, key, permission_id)
|
||||
permission_before: Final = _object_permission_rows(permission_id)
|
||||
budget_before: Final = _budget_rows(budget_id)
|
||||
assert len(permission_before) == 1
|
||||
assert len(budget_before) == 1
|
||||
aliased: Final = ownership_gateway.request("POST", "/key/update", {"key": key, "key_alias": "updated"})
|
||||
assert aliased.status_code == 200, aliased.text
|
||||
after_alias: Final = _key_rows(key)
|
||||
|
|
@ -267,9 +325,20 @@ def test_key_update_rejects_team_change_and_allows_unchanged_values(ownership_ga
|
|||
assert after_unchanged_team == after_alias
|
||||
before_reassignment: Final = _key_rows(key)
|
||||
assert len(before_reassignment) == 1
|
||||
reassigned: Final = ownership_gateway.request("POST", "/key/update", {"key": key, "team_id": team_b})
|
||||
reassigned: Final = ownership_gateway.request(
|
||||
"POST",
|
||||
"/key/update",
|
||||
{
|
||||
"key": key,
|
||||
"team_id": team_b,
|
||||
"soft_budget": 5.0,
|
||||
"object_permission": {"vector_stores": ["replacement-store"]},
|
||||
},
|
||||
)
|
||||
assert reassigned.status_code == 400, reassigned.text
|
||||
assert _key_rows(key) == before_reassignment
|
||||
assert _object_permission_rows(permission_id) == permission_before
|
||||
assert _budget_rows(budget_id) == budget_before
|
||||
|
||||
|
||||
def test_key_update_can_detach_project_and_change_team(ownership_gateway: Gateway) -> None:
|
||||
|
|
@ -588,16 +657,41 @@ def test_project_update_rejects_moving_project_with_attached_key(ownership_gatew
|
|||
model: Final = scenario.model()
|
||||
team_a: Final = scenario.team(models=[model])
|
||||
team_b: Final = scenario.team(models=[model])
|
||||
attached_project: Final = scenario.project(team_a, models=[model])
|
||||
budget: Final = scenario.budget(max_budget=3)
|
||||
attached_project: Final = scenario.project(team_a, budget_id=budget, models=[model])
|
||||
key: Final = scenario.key(team_id=team_a, project_id=attached_project, models=[model])
|
||||
permission_seed: Final = ownership_gateway.request(
|
||||
"POST",
|
||||
"/project/update",
|
||||
{"project_id": attached_project, "object_permission": {"vector_stores": ["existing-store"]}},
|
||||
)
|
||||
assert permission_seed.status_code == 200, permission_seed.text
|
||||
project_permission: Final = read_rows(
|
||||
'SELECT object_permission_id FROM "LiteLLM_ProjectTable" WHERE project_id = %s',
|
||||
(attached_project,),
|
||||
)
|
||||
assert len(project_permission) == 1
|
||||
permission_id: Final = string_value(project_permission[0]["object_permission_id"])
|
||||
scenario.cleanups.callback(_clear_project_object_permission, attached_project, permission_id)
|
||||
project_before: Final = _project_rows(attached_project)
|
||||
key_before: Final = _key_rows(key)
|
||||
permission_before: Final = _object_permission_rows(permission_id)
|
||||
budget_before: Final = _budget_rows(budget)
|
||||
moved_with_key: Final = ownership_gateway.request(
|
||||
"POST", "/project/update", {"project_id": attached_project, "team_id": team_b}
|
||||
"POST",
|
||||
"/project/update",
|
||||
{
|
||||
"project_id": attached_project,
|
||||
"team_id": team_b,
|
||||
"max_budget": 11,
|
||||
"object_permission": {"vector_stores": ["replacement-store"]},
|
||||
},
|
||||
)
|
||||
assert moved_with_key.status_code == 400, moved_with_key.text
|
||||
assert _project_rows(attached_project) == project_before
|
||||
assert _key_rows(key) == key_before
|
||||
assert _object_permission_rows(permission_id) == permission_before
|
||||
assert _budget_rows(budget) == budget_before
|
||||
|
||||
|
||||
def test_project_update_allows_moving_project_without_keys(ownership_gateway: Gateway) -> None:
|
||||
|
|
|
|||
|
|
@ -1277,6 +1277,15 @@ async def test_update_project_rejects_move_when_attached_teamless_key_exists(
|
|||
destination_team_id: Final = "team-b"
|
||||
mock_prisma: Final = _project_update_mocks(monkeypatch, {})
|
||||
mock_prisma.db.litellm_projecttable.find_unique.return_value.team_id = "team-a"
|
||||
mock_prisma.db.litellm_projecttable.find_unique.return_value.budget_id = "budget-project"
|
||||
mock_prisma.db.litellm_projecttable.find_unique.return_value.object_permission_id = "permission-project"
|
||||
budget_table: Final = mock.MagicMock()
|
||||
budget_table.update = mock.AsyncMock()
|
||||
permission_table: Final = mock.MagicMock()
|
||||
permission_table.update = mock.AsyncMock()
|
||||
permission_table.create = mock.AsyncMock()
|
||||
mock_prisma.db.litellm_budgettable = budget_table
|
||||
mock_prisma.db.litellm_objectpermissiontable = permission_table
|
||||
mock_prisma.db.litellm_teamtable.find_unique = mock.AsyncMock(
|
||||
return_value=LiteLLM_TeamTable(team_id=destination_team_id)
|
||||
)
|
||||
|
|
@ -1293,7 +1302,12 @@ async def test_update_project_rejects_move_when_attached_teamless_key_exists(
|
|||
mock_prisma.writer_db.litellm_verificationtoken.count = mock.AsyncMock(side_effect=count_teamless_keys)
|
||||
|
||||
with pytest.raises(ProxyException) as error:
|
||||
await _run_project_update(project_id, team_id=destination_team_id)
|
||||
await _run_project_update(
|
||||
project_id,
|
||||
team_id=destination_team_id,
|
||||
max_budget=50,
|
||||
object_permission={"vector_stores": ["replacement-store"]},
|
||||
)
|
||||
|
||||
expected_detail: Final = {
|
||||
"error": (
|
||||
|
|
@ -1305,6 +1319,9 @@ async def test_update_project_rejects_move_when_attached_teamless_key_exists(
|
|||
assert expected_detail["error"] in error.value.message
|
||||
mock_prisma.writer_db.litellm_verificationtoken.count.assert_awaited_once()
|
||||
mock_prisma.db.litellm_verificationtoken.count.assert_not_awaited()
|
||||
budget_table.update.assert_not_awaited()
|
||||
permission_table.update.assert_not_awaited()
|
||||
permission_table.create.assert_not_awaited()
|
||||
mock_prisma.db.litellm_projecttable.update.assert_not_awaited()
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -1076,208 +1076,146 @@ async def test_key_generation_with_mcp_tool_permissions(monkeypatch):
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_key_update_object_permissions_existing_permission():
|
||||
"""
|
||||
Test updating object permissions when a key already has an existing object_permission_id.
|
||||
|
||||
This test verifies that when updating vector stores for a key that already has an
|
||||
object_permission_id, the existing LiteLLM_ObjectPermissionTable record is updated
|
||||
with the new permissions and the object_permission_id remains the same.
|
||||
"""
|
||||
async def test_key_update_prepares_existing_object_permission_without_writing():
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.proxy._types import (
|
||||
LiteLLM_ObjectPermissionBase,
|
||||
LiteLLM_VerificationToken,
|
||||
)
|
||||
from litellm.proxy._types import LiteLLM_ObjectPermissionBase, LiteLLM_VerificationToken
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
_handle_update_object_permission,
|
||||
_prepare_key_update_object_permission,
|
||||
_write_prepared_key_update_object_permission,
|
||||
)
|
||||
|
||||
mock_prisma_client = AsyncMock()
|
||||
|
||||
# Mock existing key with object_permission_id
|
||||
existing_key_row = LiteLLM_VerificationToken(
|
||||
token="test_token_hash",
|
||||
object_permission_id="existing_perm_id_123",
|
||||
user_id="user123",
|
||||
team_id=None,
|
||||
)
|
||||
|
||||
# Mock existing object permission record
|
||||
existing_object_permission = MagicMock()
|
||||
existing_object_permission.model_dump.return_value = {
|
||||
existing_permission = MagicMock()
|
||||
existing_permission.model_dump.return_value = {
|
||||
"object_permission_id": "existing_perm_id_123",
|
||||
"vector_stores": ["old_store_1", "old_store_2"],
|
||||
"vector_stores": ["old_store"],
|
||||
}
|
||||
mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock(return_value=existing_permission)
|
||||
permission_upsert = AsyncMock()
|
||||
mock_prisma_client.db.litellm_objectpermissiontable.upsert = permission_upsert
|
||||
|
||||
mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock(
|
||||
return_value=existing_object_permission
|
||||
object_permission_data = LiteLLM_ObjectPermissionBase(vector_stores=["new_store"]).model_dump(
|
||||
exclude_unset=True, exclude_none=True
|
||||
)
|
||||
|
||||
# Mock upsert operation
|
||||
updated_permission = MagicMock()
|
||||
updated_permission.object_permission_id = "existing_perm_id_123"
|
||||
mock_prisma_client.db.litellm_objectpermissiontable.upsert = AsyncMock(
|
||||
return_value=updated_permission
|
||||
)
|
||||
|
||||
# Test data with new object permission
|
||||
data_json = {
|
||||
"object_permission": LiteLLM_ObjectPermissionBase(
|
||||
vector_stores=["new_store_1", "new_store_2", "new_store_3"]
|
||||
).model_dump(exclude_unset=True, exclude_none=True),
|
||||
"user_id": "user123",
|
||||
}
|
||||
|
||||
# Call the function
|
||||
result = await _handle_update_object_permission(
|
||||
data_json=data_json,
|
||||
upsert = await _prepare_key_update_object_permission(
|
||||
object_permission_data=object_permission_data,
|
||||
existing_key_row=existing_key_row,
|
||||
prisma_client=mock_prisma_client,
|
||||
)
|
||||
|
||||
# Verify the object_permission was removed from data_json and object_permission_id was set
|
||||
assert "object_permission" not in result
|
||||
assert result["object_permission_id"] == "existing_perm_id_123"
|
||||
|
||||
# Verify database operations were called correctly
|
||||
mock_prisma_client.db.litellm_objectpermissiontable.find_unique.assert_called_once_with(
|
||||
assert upsert is not None
|
||||
assert upsert.object_permission_id == "existing_perm_id_123"
|
||||
assert upsert.record["vector_stores"] == ["new_store"]
|
||||
permission_upsert.assert_not_awaited()
|
||||
mock_prisma_client.db.litellm_objectpermissiontable.find_unique.assert_awaited_once_with(
|
||||
where={"object_permission_id": "existing_perm_id_123"}
|
||||
)
|
||||
mock_prisma_client.db.litellm_objectpermissiontable.upsert.assert_called_once()
|
||||
|
||||
result = await _write_prepared_key_update_object_permission(
|
||||
data_json={"user_id": "user123"},
|
||||
upsert=upsert,
|
||||
prisma_client=mock_prisma_client,
|
||||
)
|
||||
|
||||
assert result == {"user_id": "user123", "object_permission_id": "existing_perm_id_123"}
|
||||
permission_upsert.assert_awaited_once_with(
|
||||
where={"object_permission_id": "existing_perm_id_123"},
|
||||
data={"create": upsert.record, "update": upsert.record},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_key_update_object_permissions_no_existing_permission():
|
||||
"""
|
||||
Test creating object permissions when a key has no existing object_permission_id.
|
||||
async def test_key_update_prepares_json_object_permission_and_upserts_new_row():
|
||||
import json
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
This test verifies that when updating object permissions for a key that has
|
||||
object_permission_id set to None, a new entry is created in the
|
||||
LiteLLM_ObjectPermissionTable and the key is updated with the new object_permission_id.
|
||||
"""
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.proxy._types import (
|
||||
LiteLLM_ObjectPermissionBase,
|
||||
LiteLLM_VerificationToken,
|
||||
)
|
||||
from litellm.proxy._types import LiteLLM_VerificationToken
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
_handle_update_object_permission,
|
||||
_prepare_key_update_object_permission,
|
||||
_write_prepared_key_update_object_permission,
|
||||
)
|
||||
|
||||
mock_prisma_client = AsyncMock()
|
||||
|
||||
existing_key_row_no_perm = LiteLLM_VerificationToken(
|
||||
existing_key_row = LiteLLM_VerificationToken(
|
||||
token="test_token_hash_2",
|
||||
object_permission_id=None,
|
||||
user_id="user456",
|
||||
team_id=None,
|
||||
)
|
||||
mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock(return_value=None)
|
||||
permission_upsert = AsyncMock()
|
||||
mock_prisma_client.db.litellm_objectpermissiontable.upsert = permission_upsert
|
||||
|
||||
# Mock find_unique to return None (no existing permission)
|
||||
mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock(
|
||||
return_value=None
|
||||
)
|
||||
|
||||
# Mock upsert to create new record
|
||||
new_permission = MagicMock()
|
||||
new_permission.object_permission_id = "new_perm_id_456"
|
||||
mock_prisma_client.db.litellm_objectpermissiontable.upsert = AsyncMock(
|
||||
return_value=new_permission
|
||||
)
|
||||
|
||||
data_json = {
|
||||
"object_permission": LiteLLM_ObjectPermissionBase(
|
||||
vector_stores=["brand_new_store"]
|
||||
).model_dump(exclude_unset=True, exclude_none=True),
|
||||
"user_id": "user456",
|
||||
}
|
||||
|
||||
result = await _handle_update_object_permission(
|
||||
data_json=data_json,
|
||||
existing_key_row=existing_key_row_no_perm,
|
||||
upsert = await _prepare_key_update_object_permission(
|
||||
object_permission_data=json.dumps({"vector_stores": ["brand_new_store"]}),
|
||||
existing_key_row=existing_key_row,
|
||||
prisma_client=mock_prisma_client,
|
||||
)
|
||||
|
||||
# Verify new object_permission_id was set
|
||||
assert "object_permission" not in result
|
||||
assert result["object_permission_id"] == "new_perm_id_456"
|
||||
# Verify upsert was called to create new record
|
||||
mock_prisma_client.db.litellm_objectpermissiontable.upsert.assert_called_once()
|
||||
assert upsert is not None
|
||||
assert upsert.record["vector_stores"] == ["brand_new_store"]
|
||||
permission_upsert.assert_not_awaited()
|
||||
result = await _write_prepared_key_update_object_permission(
|
||||
data_json={},
|
||||
upsert=upsert,
|
||||
prisma_client=mock_prisma_client,
|
||||
)
|
||||
|
||||
assert result["object_permission_id"] == upsert.object_permission_id
|
||||
permission_upsert.assert_awaited_once_with(
|
||||
where={"object_permission_id": upsert.object_permission_id},
|
||||
data={"create": upsert.record, "update": upsert.record},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_key_update_object_permissions_missing_permission_record():
|
||||
"""
|
||||
Test creating object permissions when existing object_permission_id record is not found.
|
||||
async def test_key_update_recreates_missing_object_permission_with_existing_id():
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
This test verifies that when updating object permissions for a key that has an
|
||||
object_permission_id but the corresponding record cannot be found in the database,
|
||||
a new entry is created in the LiteLLM_ObjectPermissionTable with the new permissions.
|
||||
"""
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.proxy._types import (
|
||||
LiteLLM_ObjectPermissionBase,
|
||||
LiteLLM_VerificationToken,
|
||||
)
|
||||
from litellm.proxy._types import LiteLLM_ObjectPermissionBase, LiteLLM_VerificationToken
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
_handle_update_object_permission,
|
||||
_prepare_key_update_object_permission,
|
||||
_write_prepared_key_update_object_permission,
|
||||
)
|
||||
|
||||
mock_prisma_client = AsyncMock()
|
||||
|
||||
existing_key_row_missing_perm = LiteLLM_VerificationToken(
|
||||
existing_key_row = LiteLLM_VerificationToken(
|
||||
token="test_token_hash_3",
|
||||
object_permission_id="missing_perm_id_789",
|
||||
user_id="user789",
|
||||
team_id=None,
|
||||
)
|
||||
mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock(return_value=None)
|
||||
permission_upsert = AsyncMock()
|
||||
mock_prisma_client.db.litellm_objectpermissiontable.upsert = permission_upsert
|
||||
|
||||
# Mock find_unique to return None (permission record not found)
|
||||
mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock(
|
||||
return_value=None
|
||||
)
|
||||
|
||||
# Mock upsert to create new record
|
||||
new_permission = MagicMock()
|
||||
new_permission.object_permission_id = "recreated_perm_id_789"
|
||||
mock_prisma_client.db.litellm_objectpermissiontable.upsert = AsyncMock(
|
||||
return_value=new_permission
|
||||
)
|
||||
|
||||
data_json = {
|
||||
"object_permission": LiteLLM_ObjectPermissionBase(
|
||||
vector_stores=["recreated_store"]
|
||||
).model_dump(exclude_unset=True, exclude_none=True),
|
||||
"user_id": "user789",
|
||||
}
|
||||
|
||||
result = await _handle_update_object_permission(
|
||||
data_json=data_json,
|
||||
existing_key_row=existing_key_row_missing_perm,
|
||||
upsert = await _prepare_key_update_object_permission(
|
||||
object_permission_data=LiteLLM_ObjectPermissionBase(vector_stores=["recreated_store"]).model_dump(
|
||||
exclude_unset=True, exclude_none=True
|
||||
),
|
||||
existing_key_row=existing_key_row,
|
||||
prisma_client=mock_prisma_client,
|
||||
)
|
||||
|
||||
# Verify new object_permission_id was set
|
||||
assert "object_permission" not in result
|
||||
assert result["object_permission_id"] == "recreated_perm_id_789"
|
||||
|
||||
# Verify find_unique was called with the missing permission ID
|
||||
mock_prisma_client.db.litellm_objectpermissiontable.find_unique.assert_called_once_with(
|
||||
where={"object_permission_id": "missing_perm_id_789"}
|
||||
assert upsert is not None
|
||||
assert upsert.object_permission_id == "missing_perm_id_789"
|
||||
permission_upsert.assert_not_awaited()
|
||||
await _write_prepared_key_update_object_permission(
|
||||
data_json={},
|
||||
upsert=upsert,
|
||||
prisma_client=mock_prisma_client,
|
||||
)
|
||||
permission_upsert.assert_awaited_once_with(
|
||||
where={"object_permission_id": "missing_perm_id_789"},
|
||||
data={"create": upsert.record, "update": upsert.record},
|
||||
)
|
||||
|
||||
# Verify upsert was called to create new record
|
||||
mock_prisma_client.db.litellm_objectpermissiontable.upsert.assert_called_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -7291,7 +7229,7 @@ async def test_bulk_update_keys_object_permission_is_granted_not_dropped(monkeyp
|
|||
upserted = prisma.db.litellm_objectpermissiontable.upsert.call_args.kwargs["data"]["create"]
|
||||
assert upserted["vector_stores"] == ["vs-1"]
|
||||
written = _written_key_row(prisma)
|
||||
assert written["object_permission_id"] == "objperm-bulk"
|
||||
assert written["object_permission_id"] == upserted["object_permission_id"]
|
||||
assert not {"max_budget", "team_id", "budget_id"} & written.keys()
|
||||
|
||||
|
||||
|
|
@ -13170,6 +13108,7 @@ async def test_execute_virtual_key_regeneration_hides_the_untouched_modal_expiry
|
|||
_POLICY_DENIAL_MESSAGE = "key duration must be 7d or less"
|
||||
_POLICY_HASHED_TOKEN = "0d62f396c1317066f55a96086517047c737087c61eb2bf016b72e6298927b15b"
|
||||
_POLICY_GENERATED_KEY = {"key": "sk-test-key", "expires": None, "user_id": "test-user", "team_id": None}
|
||||
_OBJECT_PERMISSION_ID_AFTER_POLICY = "perm-after-policy"
|
||||
|
||||
|
||||
def _seven_day_policy(received: list[CustomKeyPolicyRequest]):
|
||||
|
|
@ -13311,7 +13250,11 @@ async def test_regenerate_without_changes_still_runs_custom_key_policy(data):
|
|||
|
||||
def _policy_existing_team_key() -> LiteLLM_VerificationToken:
|
||||
return LiteLLM_VerificationToken(
|
||||
token=_POLICY_HASHED_TOKEN, user_id="test-user", team_id="team-a", max_budget=200.0
|
||||
token=_POLICY_HASHED_TOKEN,
|
||||
user_id="test-user",
|
||||
team_id="team-a",
|
||||
max_budget=200.0,
|
||||
object_permission_id=_OBJECT_PERMISSION_ID_AFTER_POLICY,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -13457,9 +13400,6 @@ async def test_process_single_key_update_rejects_when_custom_key_policy_denies()
|
|||
assert [policy_request.operation for policy_request in received] == ["update"]
|
||||
|
||||
|
||||
_OBJECT_PERMISSION_ID_AFTER_POLICY = "perm-after-policy"
|
||||
|
||||
|
||||
def _record_object_permission_writes(mock_prisma_client: AsyncMock, events: list[str]) -> None:
|
||||
mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock(return_value=None)
|
||||
|
||||
|
|
@ -13579,9 +13519,12 @@ async def test_regenerate_writes_the_object_permission_row_only_after_the_policy
|
|||
events: list[str] = []
|
||||
_record_object_permission_writes(mock_prisma_client, events)
|
||||
data = RegenerateKeyRequest(max_budget=50.0, object_permission=LiteLLM_ObjectPermissionBase(vector_stores=["vs-1"]))
|
||||
existing_key: Final = _make_regenerate_existing_key().model_copy(
|
||||
update={"object_permission_id": _OBJECT_PERMISSION_ID_AFTER_POLICY}
|
||||
)
|
||||
|
||||
with _regenerate_policy_mocks(_recording_policy(events, allowed=True), AsyncMock(), AsyncMock()):
|
||||
await _regenerate_under_policy(mock_prisma_client, _make_regenerate_existing_key(), data)
|
||||
await _regenerate_under_policy(mock_prisma_client, existing_key, data)
|
||||
|
||||
_assert_permission_row_written_after_policy(
|
||||
events, mock_prisma_client.db.litellm_verificationtoken.update.await_args.kwargs["data"]
|
||||
|
|
@ -20960,11 +20903,25 @@ async def test_key_update_rejects_team_change_for_project_bound_key(monkeypatch:
|
|||
token="hashed-key",
|
||||
team_id=_OWNERSHIP_KEY_TEAM,
|
||||
project_id=_OWNED_PROJECT,
|
||||
object_permission_id="permission-update",
|
||||
budget_id="budget-update",
|
||||
)
|
||||
mock_prisma_client: Final = _wire_update_key_fn(monkeypatch, existing_key_row)
|
||||
existing_permission: Final = LiteLLM_ObjectPermissionTable(
|
||||
object_permission_id="permission-update",
|
||||
vector_stores=["existing-store"],
|
||||
)
|
||||
mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock(return_value=existing_permission)
|
||||
permission_upsert: Final = AsyncMock()
|
||||
mock_prisma_client.db.litellm_objectpermissiontable.upsert = permission_upsert
|
||||
permission_before: Final = existing_permission.model_dump()
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints.get_team_object",
|
||||
AsyncMock(return_value=LiteLLM_TeamTable(team_id=_OWNERSHIP_DESTINATION_TEAM, team_members=[])),
|
||||
AsyncMock(
|
||||
return_value=LiteLLM_TeamTable(team_id=_OWNERSHIP_DESTINATION_TEAM).model_copy(
|
||||
update={"team_members": []}
|
||||
)
|
||||
),
|
||||
)
|
||||
mock_prisma_client.writer_db.litellm_projecttable.find_unique = AsyncMock(
|
||||
return_value=LiteLLM_ProjectTable(project_id=_OWNED_PROJECT, team_id=_OWNERSHIP_PROJECT_TEAM)
|
||||
|
|
@ -20976,7 +20933,12 @@ async def test_key_update_rejects_team_change_for_project_bound_key(monkeypatch:
|
|||
with pytest.raises((HTTPException, ProxyException)) as error:
|
||||
await update_key_fn(
|
||||
request=mock_request,
|
||||
data=UpdateKeyRequest(key="sk-key", team_id=_OWNERSHIP_DESTINATION_TEAM),
|
||||
data=UpdateKeyRequest(
|
||||
key="sk-key",
|
||||
team_id=_OWNERSHIP_DESTINATION_TEAM,
|
||||
soft_budget=5.0,
|
||||
object_permission=LiteLLM_ObjectPermissionBase(vector_stores=["replacement-store"]),
|
||||
),
|
||||
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin"),
|
||||
litellm_changed_by=None,
|
||||
)
|
||||
|
|
@ -20987,6 +20949,80 @@ async def test_key_update_rejects_team_change_for_project_bound_key(monkeypatch:
|
|||
f"but the key belongs to {_OWNERSHIP_DESTINATION_TEAM}"
|
||||
)
|
||||
assert expected_detail in str(getattr(error.value, "detail", None) or getattr(error.value, "message", None))
|
||||
assert existing_permission.model_dump() == permission_before
|
||||
permission_upsert.assert_not_awaited()
|
||||
mock_prisma_client.tx.assert_not_called()
|
||||
mock_prisma_client.update_data.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_key_update_ambiguous_permission_error_precedes_project_ownership(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
existing_key_row: Final = LiteLLM_VerificationToken(
|
||||
token="hashed-key",
|
||||
team_id=_OWNERSHIP_KEY_TEAM,
|
||||
project_id=_OWNED_PROJECT,
|
||||
object_permission_id="permission-ambiguous",
|
||||
)
|
||||
mock_prisma_client: Final = _wire_update_key_fn(monkeypatch, existing_key_row)
|
||||
mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock(
|
||||
return_value=LiteLLM_ObjectPermissionTable(
|
||||
object_permission_id="permission-ambiguous",
|
||||
mcp_tool_permissions={},
|
||||
)
|
||||
)
|
||||
permission_upsert: Final = AsyncMock()
|
||||
mock_prisma_client.db.litellm_objectpermissiontable.upsert = permission_upsert
|
||||
mock_prisma_client.db.litellm_mcpservertable.find_many = AsyncMock(
|
||||
return_value=[
|
||||
MagicMock(server_id="wiki-a-id", alias="wiki", server_name="wiki-a"),
|
||||
MagicMock(server_id="wiki-b-id", alias="wiki", server_name="wiki-b"),
|
||||
]
|
||||
)
|
||||
mock_prisma_client.writer_db.litellm_projecttable.find_unique = AsyncMock(
|
||||
return_value=LiteLLM_ProjectTable(project_id=_OWNED_PROJECT, team_id=_OWNERSHIP_PROJECT_TEAM)
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.management_endpoints.key_management_endpoints.get_team_object",
|
||||
AsyncMock(
|
||||
return_value=LiteLLM_TeamTable(
|
||||
team_id=_OWNERSHIP_DESTINATION_TEAM,
|
||||
object_permission=LiteLLM_ObjectPermissionTable(
|
||||
object_permission_id="permission-team",
|
||||
mcp_servers=["wiki-a-id", "wiki-b-id"],
|
||||
),
|
||||
)
|
||||
),
|
||||
)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", MagicMock())
|
||||
|
||||
with pytest.raises((HTTPException, ProxyException)) as error:
|
||||
await update_key_fn(
|
||||
request=MagicMock(query_params={}),
|
||||
data=UpdateKeyRequest(
|
||||
key="sk-key",
|
||||
team_id=_OWNERSHIP_DESTINATION_TEAM,
|
||||
object_permission=LiteLLM_ObjectPermissionBase(
|
||||
mcp_tool_permissions={"wiki": ["read_wiki_structure"]}
|
||||
),
|
||||
),
|
||||
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin"),
|
||||
litellm_changed_by=None,
|
||||
)
|
||||
|
||||
assert str(getattr(error.value, "status_code", None) or getattr(error.value, "code", None)) == "400"
|
||||
expected_detail: Final = {
|
||||
"error": (
|
||||
"Ambiguous mcp_tool_permissions key: 'wiki' matches MCP servers ['wiki-a-id', 'wiki-b-id']. "
|
||||
"Key tool permissions by server_id when servers share a name or alias."
|
||||
)
|
||||
}
|
||||
assert getattr(error.value, "detail", None) == expected_detail or getattr(error.value, "message", None) == str(
|
||||
expected_detail
|
||||
)
|
||||
permission_upsert.assert_not_awaited()
|
||||
mock_prisma_client.writer_db.litellm_projecttable.find_unique.assert_not_awaited()
|
||||
mock_prisma_client.update_data.assert_not_awaited()
|
||||
|
||||
|
||||
|
|
@ -21070,11 +21106,23 @@ async def test_bulk_key_update_rejects_project_team_change_and_allows_other_fiel
|
|||
models=[],
|
||||
team_id=_OWNERSHIP_PROJECT_TEAM,
|
||||
project_id=_OWNED_PROJECT,
|
||||
object_permission_id="permission-bulk",
|
||||
budget_id="budget-bulk",
|
||||
)
|
||||
mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=existing_key_row)
|
||||
mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(
|
||||
return_value=LiteLLM_TeamTable(team_id=_OWNERSHIP_DESTINATION_TEAM)
|
||||
)
|
||||
existing_permission: Final = LiteLLM_ObjectPermissionTable(
|
||||
object_permission_id="permission-bulk",
|
||||
vector_stores=["existing-store"],
|
||||
)
|
||||
permission_before: Final = existing_permission.model_dump()
|
||||
mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock(return_value=existing_permission)
|
||||
permission_upsert: Final = AsyncMock()
|
||||
mock_prisma_client.db.litellm_objectpermissiontable.upsert = permission_upsert
|
||||
budget_update: Final = AsyncMock()
|
||||
mock_prisma_client.db.litellm_budgettable.update = budget_update
|
||||
mock_prisma_client.get_data = AsyncMock(return_value=existing_key_row)
|
||||
updated_key: Final = MagicMock()
|
||||
updated_key.model_dump.return_value = {
|
||||
|
|
@ -21089,7 +21137,14 @@ async def test_bulk_key_update_rejects_project_team_change_and_allows_other_fiel
|
|||
|
||||
request: Final = BulkUpdateKeyRequest(
|
||||
keys=[
|
||||
BulkUpdateKeyRequestItem(key="sk-bulk-key", team_id=_OWNERSHIP_DESTINATION_TEAM),
|
||||
BulkUpdateKeyRequestItem.model_validate(
|
||||
{
|
||||
"key": "sk-bulk-key",
|
||||
"team_id": _OWNERSHIP_DESTINATION_TEAM,
|
||||
"soft_budget": 5.0,
|
||||
"object_permission": LiteLLM_ObjectPermissionBase(vector_stores=["replacement-store"]),
|
||||
}
|
||||
),
|
||||
BulkUpdateKeyRequestItem(key="sk-bulk-key", max_budget=10.0, tags=["bulk-update"]),
|
||||
]
|
||||
)
|
||||
|
|
@ -21136,6 +21191,10 @@ async def test_bulk_key_update_rejects_project_team_change_and_allows_other_fiel
|
|||
assert response.successful_updates[0].key == "sk-bulk-key"
|
||||
assert response.successful_updates[0].key_info["max_budget"] == 10.0
|
||||
assert response.successful_updates[0].key_info["tags"] == ["bulk-update"]
|
||||
assert existing_permission.model_dump() == permission_before
|
||||
permission_upsert.assert_not_awaited()
|
||||
budget_update.assert_not_awaited()
|
||||
mock_prisma_client.tx.assert_not_called()
|
||||
mock_prisma_client.update_data.assert_awaited_once()
|
||||
|
||||
|
||||
|
|
@ -21199,7 +21258,11 @@ async def test_regenerate_checks_project_team_ownership(
|
|||
expected_status: int | None,
|
||||
expected_updates: int,
|
||||
) -> None:
|
||||
existing_key: Final = LiteLLM_VerificationToken(token="abc123", team_id=key_team_id)
|
||||
existing_key: Final = LiteLLM_VerificationToken(
|
||||
token="abc123",
|
||||
team_id=key_team_id,
|
||||
object_permission_id="permission-regenerate",
|
||||
)
|
||||
mock_prisma_client: Final = _make_regenerate_mock_prisma()
|
||||
mock_prisma_client.writer_db = MagicMock()
|
||||
mock_prisma_client.writer_db.litellm_projecttable = MagicMock()
|
||||
|
|
@ -21212,6 +21275,14 @@ async def test_regenerate_checks_project_team_ownership(
|
|||
deprecated_key_table: Final = MagicMock()
|
||||
deprecated_key_table.upsert = AsyncMock()
|
||||
mock_prisma_client.db.litellm_deprecatedverificationtoken = deprecated_key_table
|
||||
mock_prisma_client.db.litellm_objectpermissiontable.find_unique = AsyncMock(
|
||||
return_value=LiteLLM_ObjectPermissionTable(
|
||||
object_permission_id="permission-regenerate",
|
||||
vector_stores=["existing-store"],
|
||||
)
|
||||
)
|
||||
permission_upsert: Final = AsyncMock()
|
||||
mock_prisma_client.db.litellm_objectpermissiontable.upsert = permission_upsert
|
||||
user_api_key_cache: Final = await _cache_with_project(_OWNED_PROJECT, [], team_id="team-b")
|
||||
|
||||
async def regenerate() -> None:
|
||||
|
|
@ -21234,6 +21305,7 @@ async def test_regenerate_checks_project_team_ownership(
|
|||
data=RegenerateKeyRequest(
|
||||
project_id=_OWNED_PROJECT,
|
||||
grace_period="1h" if expected_status is not None else None,
|
||||
object_permission=LiteLLM_ObjectPermissionBase(vector_stores=["replacement-store"]),
|
||||
),
|
||||
user_api_key_dict=_make_regenerate_user_api_key_dict(),
|
||||
litellm_changed_by=None,
|
||||
|
|
@ -21248,8 +21320,10 @@ async def test_regenerate_checks_project_team_ownership(
|
|||
assert "belongs to team team-b, but the key belongs to team-a" in str(error.value.detail)
|
||||
deleted_history_table.create_many.assert_not_awaited()
|
||||
deprecated_key_table.upsert.assert_not_awaited()
|
||||
permission_upsert.assert_not_awaited()
|
||||
else:
|
||||
await regenerate()
|
||||
permission_upsert.assert_awaited_once()
|
||||
|
||||
assert mock_prisma_client.db.litellm_verificationtoken.update.await_count == expected_updates
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue