diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index fe0bda5ddb9..cb7a18cd107 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -18914,7 +18914,7 @@ } } }, - "description": "\nUnified rate-limit error.\n\nEvery rate-limit condition surfaced by litellm \u2014 whether it originated from\nan upstream LLM provider, a vendor batch endpoint, or one of litellm's own\nproxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\nmax-iterations, etc.) \u2014 is raised as an instance of this class.\n\nThe :attr:`category` attribute lets callers distinguish the source. See\n:class:`RateLimitErrorCategory` for the available values.\n" + "description": "\n Unified rate-limit error.\n\n Every rate-limit condition surfaced by litellm \u2014 whether it originated from\n an upstream LLM provider, a vendor batch endpoint, or one of litellm's own\n proxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\n max-iterations, etc.) \u2014 is raised as an instance of this class.\n\n The :attr:`category` attribute lets callers distinguish the source. See\n :class:`RateLimitErrorCategory` for the available values.\n " }, "500": { "content": { diff --git a/litellm/proxy/hooks/key_management_event_hooks.py b/litellm/proxy/hooks/key_management_event_hooks.py index 8c86f354295..5cfef11df8d 100644 --- a/litellm/proxy/hooks/key_management_event_hooks.py +++ b/litellm/proxy/hooks/key_management_event_hooks.py @@ -3,6 +3,8 @@ import json from datetime import datetime, timezone from typing import Final +from pydantic import TypeAdapter + import litellm from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid @@ -108,33 +110,32 @@ class KeyManagementEventHooks: from litellm.proxy.proxy_server import litellm_proxy_admin_name if is_audit_logging_enabled(): - updated_fields: Final = data.model_dump(exclude_none=True) - if "project_id" in data.model_fields_set: - updated_fields["project_id"] = data.project_id - _updated_values: Final = json.dumps(updated_fields, default=str) - - _before_value = existing_key_row.json(exclude_none=True) - _before_value = json.dumps(_before_value, default=str) - - asyncio.create_task( - create_audit_log_for_update( - request_data=LiteLLM_AuditLogs( - id=str(uuid.uuid4()), - updated_at=datetime.now(timezone.utc), - changed_by=get_audit_log_changed_by( - litellm_changed_by=litellm_changed_by, - user_api_key_dict=user_api_key_dict, - litellm_proxy_admin_name=litellm_proxy_admin_name, - ), - changed_by_api_key=user_api_key_dict.api_key, - table_name=LitellmTableNames.KEY_TABLE_NAME, - object_id=_hash_token_if_needed(data.key), - action="updated", - updated_values=_updated_values, - before_value=_before_value, - ) - ) + updated_fields: Final = { + **data.model_dump(exclude_none=True), + **({"project_id": data.project_id} if "project_id" in data.model_fields_set else {}), + } + audit_log: Final = LiteLLM_AuditLogs( + id=str(uuid.uuid4()), + updated_at=datetime.now(timezone.utc), + changed_by=get_audit_log_changed_by( + litellm_changed_by=litellm_changed_by, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=litellm_proxy_admin_name, + ), + changed_by_api_key=user_api_key_dict.api_key, + table_name=LitellmTableNames.KEY_TABLE_NAME, + object_id=_hash_token_if_needed(data.key), + action="updated", + updated_values=json.dumps(updated_fields, default=str), + before_value=json.dumps(existing_key_row.json(exclude_none=True), default=str), ) + masked_values: Final = TypeAdapter(dict[str, object]).validate_json(str(audit_log.updated_values)) + request_data: Final = ( + audit_log.model_copy(update={"updated_values": json.dumps({**masked_values, "project_id": None})}) + if "project_id" in data.model_fields_set and data.project_id is None + else audit_log + ) + asyncio.create_task(create_audit_log_for_update(request_data=request_data)) @staticmethod async def async_key_rotated_hook( diff --git a/tests/e2e/management/test_key_management_e2e.py b/tests/e2e/management/test_key_management_e2e.py index 585fcd676b2..353b0f7cf09 100644 --- a/tests/e2e/management/test_key_management_e2e.py +++ b/tests/e2e/management/test_key_management_e2e.py @@ -12,7 +12,7 @@ asserting once. from __future__ import annotations import time -from collections.abc import Callable +from collections.abc import Callable, Iterator from typing import Final, Literal import pytest @@ -21,8 +21,11 @@ from e2e_config import unique_marker from e2e_http import NoBody, StreamingResponse, unwrap from lifecycle import ResourceManager from management_client import ManagementClient -from models import CLEAR, ChatResponse, KeyDeleteBody, KeyGenerateBody, KeyInfo, KeyUpdateBody, LiteLLMParamsBody, OrgNewBody, TeamNewBody -from pydantic import BaseModel +from models import ( + CLEAR, ChatResponse, KeyDeleteBody, KeyGenerateBody, KeyInfo, KeyUpdateBody, + LiteLLMParamsBody, OrgNewBody, TeamNewBody, +) +from pydantic import BaseModel, RootModel pytestmark = pytest.mark.e2e @@ -145,11 +148,23 @@ class ProjectBlockBody(ProjectIdentity): blocked: bool +class ProjectDeleteBody(BaseModel): + project_ids: list[str] + + +@pytest.fixture +def project_resources(client: ManagementClient) -> Iterator[ResourceManager]: + manager: Final = ResourceManager(client=client.proxy, strict_cleanup=True) + yield manager + manager.teardown() + + class TestKeyManagementRoutes: @pytest.mark.covers("mgmt.key.update.persists") def test_project_detachment_preserves_key_scope_and_refreshes_auth( - self, client: ManagementClient, resources: ResourceManager + self, client: ManagementClient, project_resources: ResourceManager ) -> None: + resources: Final = project_resources name: Final = f"e2e-detach-{unique_marker()}" model_id: Final = client.proxy.create_model( name, LiteLLMParamsBody(model="openai/synthetic-detachment", api_key="synthetic", mock_response="orbit") @@ -164,8 +179,9 @@ class TestKeyManagementRoutes: json=ProjectCreateBody(team_id=team_id, project_alias=name, models=[name]), response_type=ProjectIdentity, )) - resources.defer(lambda: unwrap(client.proxy.transport.post( - "/project/delete", headers=client.proxy.transport.master, json=project, response_type=NoBody, + resources.defer(lambda: unwrap(client.proxy.transport.delete( + "/project/delete", headers=client.proxy.transport.master, + json=ProjectDeleteBody(project_ids=[project.project_id]), response_type=RootModel[list[ProjectIdentity]], ))) key: Final = _generate_key(client, resources, KeyGenerateBody( key_alias=name, team_id=team_id, organization_id=org_id, project_id=project.project_id, diff --git a/tests/e2e/test_proxy_client.py b/tests/e2e/test_proxy_client.py index 3b84a47e3cc..1b0133f12cb 100644 --- a/tests/e2e/test_proxy_client.py +++ b/tests/e2e/test_proxy_client.py @@ -245,6 +245,7 @@ class TestReplicasFor: replica_urls=("http://gateway-1", "http://gateway-2"), ) assert set(client.replicas_for("/key/info")) == {"http://backend"} + assert set(client.replicas_for("/project/info")) == {"http://backend"} assert set(client.replicas_for("/v1/models")) == {"http://gateway-1", "http://gateway-2"} def test_monolith_reads_management_routes_back_from_every_replica(self) -> None: diff --git a/tests/e2e/transport.py b/tests/e2e/transport.py index 44fdbaa3e41..e8caa801467 100644 --- a/tests/e2e/transport.py +++ b/tests/e2e/transport.py @@ -295,6 +295,7 @@ CONTROL_PLANE_PREFIXES: tuple[str, ...] = ( "/user", "/team", "/organization", + "/project", "/customer", "/end_user", "/tag", diff --git a/tests/test_litellm/proxy/hooks/test_key_management_event_hooks.py b/tests/test_litellm/proxy/hooks/test_key_management_event_hooks.py index bdadf1657e9..1aa9382f3fe 100644 --- a/tests/test_litellm/proxy/hooks/test_key_management_event_hooks.py +++ b/tests/test_litellm/proxy/hooks/test_key_management_event_hooks.py @@ -6,6 +6,7 @@ Validates that email and secret manager operations are independent and non-block import asyncio import json +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -497,9 +498,9 @@ class TestKeyUpdatedAuditLogObjectId: project_id="project-orbit", ) - data = UpdateKeyRequest(key=request_key, max_budget=2000.0) - if detach_project: - data.project_id = None + data: Final = UpdateKeyRequest( + key=request_key, max_budget=2000.0, **({"project_id": None} if detach_project else {}) + ) with ( patch("litellm.store_audit_logs", True), @@ -544,12 +545,14 @@ class TestKeyUpdatedAuditLogObjectId: hashed_key = hash_token("sk-raw-test-key-31620") - audit_row = await self._run_updated_hook_and_capture_audit_log(request_key=hashed_key, detach_project=detach_project) + audit_row: Final = await self._run_updated_hook_and_capture_audit_log( + request_key=hashed_key, detach_project=detach_project, + ) assert audit_row.object_id == hashed_key - updated_values = json.loads(audit_row.updated_values) + updated_values: Final = json.loads(audit_row.updated_values) assert ("project_id" in updated_values) is detach_project if detach_project: - assert updated_values["project_id"] == "None" + assert updated_values["project_id"] is None assert json.loads(audit_row.before_value)["project_id"] == "project-orbit" assert updated_values["max_budget"] == 2000.0 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 0d3434942c0..646ae43f37a 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 @@ -18167,7 +18167,8 @@ async def test_project_detachment_uses_effective_project_for_validation(project_ None, False, MagicMock(), cache, ) assert exc.value.status_code == 400 - assert ("not in project's allowed models" if project_id == "project-orbit" else "reassignment") in str(exc.value.detail) + expected: Final = "not in project's allowed models" if project_id == "project-orbit" else "reassignment" + assert expected in str(exc.value.detail) @pytest.mark.asyncio