diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index a566d491597..f5a312cef23 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1241,6 +1241,27 @@ class UpdateKeyRequest(KeyRequestBase): return self +class PatchKeyRequest(UpdateKeyRequest): + """ + Body of PATCH /key/{key}. + + Differs from UpdateKeyRequest in two ways. `key` is optional, because PATCH takes it + from the path; a `key` in the body is still accepted when it matches. Unknown fields + are rejected rather than ignored, because on a merge patch the set of fields present + *is* the request, so a misspelled field has to fail loudly instead of silently + no-op'ing. + """ + + model_config = ConfigDict(extra="forbid") + + key: str | None = None + + @model_validator(mode="after") + def validate_key_identifier(self) -> "PatchKeyRequest": + """The path supplies the identifier, so the body needs neither key nor key_alias.""" + return self + + class RegenerateKeyRequest(GenerateKeyRequest): # This needs to be different from UpdateKeyRequest, because "key" is optional for this key: str | None = None diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index ca2607653a1..004cefb2f72 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -20,11 +20,12 @@ import secrets import traceback from collections.abc import Awaitable, Callable, Mapping, Sequence from datetime import datetime, timedelta, timezone -from typing import Any, Final, Literal, Optional, Protocol, TypeVar, cast +from typing import Annotated, Any, Final, Literal, Optional, Protocol, TypeVar, cast import fastapi import yaml from fastapi import APIRouter, Depends, Header, HTTPException, Query, Request, status +from pydantic import JsonValue import litellm from litellm._logging import verbose_proxy_logger @@ -68,6 +69,7 @@ from litellm.proxy.common_utils.config_sync_pubsub import ( coordination_redis_cache, publish_config_change, ) +from litellm.proxy.common_utils.json_merge_patch import apply_json_merge_patch from litellm.proxy.common_utils.rbac_utils import check_org_admin_can_generate_keys from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache @@ -2903,6 +2905,79 @@ async def update_key_fn( ) +async def _merge_key_metadata( + key: str, + prisma_client: PrismaClient | None, + metadata_patch: JsonValue, +) -> JsonValue: + """Deep-merge a metadata patch onto the key's stored metadata, per RFC 7386.""" + existing_key_row: Final = await _get_and_validate_existing_key(token=key, prisma_client=prisma_client) + existing_metadata: Final = existing_key_row.metadata if isinstance(existing_key_row.metadata, dict) else {} + return apply_json_merge_patch(existing_metadata, metadata_patch) + + +@router.patch( + "/key/{key:path}", + tags=["key management"], + dependencies=[Depends(user_api_key_auth)], + include_in_schema=False, +) +async def patch_key( + key: str, + data: PatchKeyRequest, + request: Request, + user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)], + litellm_changed_by: Annotated[ + str | None, + Header( + description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability", + ), + ] = None, +): + """ + Partially update a key using RFC 7386 JSON Merge Patch semantics. + + `key` is taken from the path; a `key` in the body is accepted only when it matches. + `metadata` is merged with the key's stored metadata rather than replacing it: an + omitted entry is preserved, `entry: null` deletes it, and any other value overwrites + (recursing into nested objects). Every other field behaves exactly like + `POST /key/update` (omitted preserves, `null` clears, a value overwrites). An unknown + field is rejected with a 422 rather than silently ignored. Returns the updated key. + + ``` + curl --location --request PATCH 'http://0.0.0.0:4000/key/sk-1234' \ + --header 'Authorization: Bearer sk-1234' \ + --header 'Content-Type: application/json' \ + --data-raw '{ + "metadata": {"cost_center": "1234", "deprecated_entry": null} + }' + ``` + """ + from litellm.proxy.proxy_server import prisma_client + + if data.key is not None and data.key != key: + raise ProxyException( + message="key in body does not match key in path", + type=ProxyErrorTypes.bad_request_error, + param="key", + code=status.HTTP_400_BAD_REQUEST, + ) + + patch_fields: Final = data.model_dump(exclude_unset=True, exclude={"key"}) + merged_fields: Final = ( + {**patch_fields, "metadata": await _merge_key_metadata(key, prisma_client, patch_fields["metadata"])} + if "metadata" in patch_fields + else patch_fields + ) + + return await update_key_fn( + request=request, + data=UpdateKeyRequest.model_validate({"key": key, **merged_fields}), + user_api_key_dict=user_api_key_dict, + litellm_changed_by=litellm_changed_by, + ) + + @router.post( "/key/bulk_update", tags=["key management"], 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 bdf09a95e4b..7fd752b3756 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 @@ -16444,3 +16444,197 @@ async def test_regenerate_key_repoints_live_membership_not_the_key_row_it_read( assert await _authorized_models_for_key( access_groups, new_token_hash, ["ag-revoked-since", "ag-attached-since"] ) == ["attached-model"] + + +# --------------------------------------------------------------------------- +# PATCH /key/{key} - RFC 7386 JSON merge patch. +# Drives POST update_key_fn and PATCH patch_key against the same mocked key row +# and asserts on the exact dict handed to prisma_client.update_data. +# --------------------------------------------------------------------------- + +_PATCH_KEY_TOKEN = "a1b2c3d4" * 8 + + +async def _drive_key_write( + kind, + monkeypatch, + *, + existing_metadata=None, + existing_kwargs=None, + payload=None, +): + """Drive POST ``update_key_fn`` or PATCH ``patch_key`` against a mocked key. + + Returns the dict handed to ``prisma_client.update_data``. + """ + from litellm.proxy._types import PatchKeyRequest + from litellm.proxy.management_endpoints.key_management_endpoints import ( + patch_key, + update_key_fn, + ) + + key_in_db = LiteLLM_VerificationToken( + token=_PATCH_KEY_TOKEN, + user_id="test-user", + **({"metadata": existing_metadata} if existing_metadata is not None else {}), + **(existing_kwargs or {}), + ) + + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock( + return_value=key_in_db + ) + mock_prisma_client.update_data = AsyncMock(return_value={"data": {}}) + _setup_update_key_mocks(monkeypatch, mock_prisma_client) + + auth = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user" + ) + + with patch( + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", + new=AsyncMock(), + ): + if kind == "post": + await update_key_fn( + request=MagicMock(), + data=UpdateKeyRequest(key=_PATCH_KEY_TOKEN, **(payload or {})), + user_api_key_dict=auth, + litellm_changed_by=None, + ) + else: + await patch_key( + key=_PATCH_KEY_TOKEN, + data=PatchKeyRequest.model_validate(dict(payload or {})), + request=MagicMock(), + user_api_key_dict=auth, + litellm_changed_by=None, + ) + + return mock_prisma_client.update_data.call_args.kwargs["data"] + + +# (label, existing_metadata, merge_patch_body, expected_POST_metadata, expected_PATCH_metadata) +_KEY_METADATA_MAPPING = [ + ( + "sibling entries survive a merge patch but not a POST", + {"cost_center": "cc-1", "owner": "data-eng"}, + {"cost_center": "cc-2"}, + {"cost_center": "cc-2"}, + {"cost_center": "cc-2", "owner": "data-eng"}, + ), + ( + "a nested object recurses instead of being replaced", + {"nested": {"a": 1, "b": 2}, "owner": "data-eng"}, + {"nested": {"b": 99}}, + {"nested": {"b": 99}}, + {"nested": {"a": 1, "b": 99}, "owner": "data-eng"}, + ), + ( + "null deletes only its own entry", + {"cost_center": "cc-1", "owner": "data-eng"}, + {"cost_center": None}, + {"cost_center": None}, + {"owner": "data-eng"}, + ), +] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "label,existing,body,expected_post,expected_patch", + _KEY_METADATA_MAPPING, + ids=[row[0] for row in _KEY_METADATA_MAPPING], +) +async def test_patch_key_merges_metadata_where_post_replaces_it( + monkeypatch, label, existing, body, expected_post, expected_patch +): + """POST /key/update writes the submitted metadata verbatim, dropping every entry + the caller did not resend. PATCH merges it onto what is already stored.""" + written_post = await _drive_key_write( + "post", monkeypatch, existing_metadata=existing, payload={"metadata": body} + ) + assert written_post["metadata"] == expected_post + + written_patch = await _drive_key_write( + "patch", monkeypatch, existing_metadata=existing, payload={"metadata": body} + ) + assert written_patch["metadata"] == expected_patch + + +@pytest.mark.asyncio +async def test_patch_key_null_clears_and_omission_preserves(monkeypatch): + """Both directions in one test: a route that cleared everything would pass a + clear-only assertion, and a route that cleared nothing would pass a preserve-only one.""" + cleared = await _drive_key_write( + "patch", monkeypatch, existing_kwargs={"tpm_limit": 500}, payload={"tpm_limit": None} + ) + assert "tpm_limit" in cleared + assert cleared["tpm_limit"] is None + + preserved = await _drive_key_write( + "patch", monkeypatch, existing_kwargs={"tpm_limit": 500}, payload={"rpm_limit": 9} + ) + assert "tpm_limit" not in preserved + assert preserved["rpm_limit"] == 9 + + +@pytest.mark.asyncio +async def test_patch_key_does_not_slide_the_budget_window(monkeypatch): + """A merge patch is idempotent, so a patch that never mentions budget_duration + must leave budget_reset_at alone rather than postponing the key's reset.""" + written = await _drive_key_write( + "patch", + monkeypatch, + existing_kwargs={"budget_duration": "30d"}, + payload={"rpm_limit": 5}, + ) + assert written["rpm_limit"] == 5 + assert "budget_reset_at" not in written + assert "budget_duration" not in written + + +@pytest.mark.asyncio +async def test_patch_key_rejects_a_body_key_that_disagrees_with_the_path(monkeypatch): + """The path is authoritative, and the rejection must not echo either secret back.""" + from litellm.proxy._types import PatchKeyRequest + from litellm.proxy.management_endpoints.key_management_endpoints import patch_key + + mock_prisma_client = AsyncMock() + _setup_update_key_mocks(monkeypatch, mock_prisma_client) + + with pytest.raises(ProxyException) as exc_info: + await patch_key( + key=_PATCH_KEY_TOKEN, + data=PatchKeyRequest(key="sk-a-different-key", tpm_limit=1), + request=MagicMock(), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user" + ), + litellm_changed_by=None, + ) + + assert str(exc_info.value.code) == "400" + assert "sk-a-different-key" not in str(exc_info.value.message) + assert _PATCH_KEY_TOKEN not in str(exc_info.value.message) + mock_prisma_client.update_data.assert_not_called() + + +def test_patch_key_request_rejects_unknown_fields(): + """On a merge patch the set of fields present IS the request, so a misspelled + field has to fail loudly instead of silently no-op'ing the way POST does.""" + from pydantic import ValidationError + + from litellm.proxy._types import PatchKeyRequest + + with pytest.raises(ValidationError): + PatchKeyRequest.model_validate({"tpm_limitt": 5}) + + assert UpdateKeyRequest(key="sk-x", **{"tpm_limitt": 5}).tpm_limit is None + + +def test_patch_key_request_makes_key_optional(): + """PATCH takes the identifier from the path, so the body needs neither key nor key_alias.""" + from litellm.proxy._types import PatchKeyRequest + + assert PatchKeyRequest.model_validate({"tpm_limit": 5}).key is None diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 095afc6110f..8fca7eaad4d 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -7371,6 +7371,39 @@ export interface paths { patch?: never; trace?: never; }; + "/key/{key}": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + /** + * Patch Key + * @description Partially update a key using RFC 7386 JSON Merge Patch semantics. + * + * `key` is taken from the path; a `key` in the body is accepted only when it matches. + * `metadata` is merged with the key's stored metadata rather than replacing it: an + * omitted entry is preserved, `entry: null` deletes it, and any other value overwrites + * (recursing into nested objects). Every other field behaves exactly like + * `POST /key/update` (omitted preserves, `null` clears, a value overwrites). An unknown + * field is rejected with a 422 rather than silently ignored. Returns the updated key. + * + * ``` + * curl --location --request PATCH 'http://0.0.0.0:4000/key/sk-1234' --header 'Authorization: Bearer sk-1234' --header 'Content-Type: application/json' --data-raw '{ + * "metadata": {"cost_center": "1234", "deprecated_entry": null} + * }' + * ``` + */ + patch: operations["patch_key_key__key__patch"]; + trace?: never; + }; "/key/{key}/regenerate": { parameters: { query?: never; @@ -30307,6 +30340,159 @@ export interface components { guardrail_name?: string | null; litellm_params?: components["schemas"]["BaseLitellmParams-Input"] | null; }; + /** + * PatchKeyRequest + * @description Body of PATCH /key/{key}. + * + * Differs from UpdateKeyRequest in two ways. `key` is optional, because PATCH takes it + * from the path; a `key` in the body is still accepted when it matches. Unknown fields + * are rejected rather than ignored, because on a merge patch the set of fields present + * *is* the request, so a misspelled field has to fail loudly instead of silently + * no-op'ing. + */ + PatchKeyRequest: { + /** Access Group Ids */ + access_group_ids?: string[] | null; + /** Agent Id */ + agent_id?: string | null; + /** + * Aliases + * @default {} + */ + aliases: { + [key: string]: unknown; + } | null; + /** + * Allowed Cache Controls + * @default [] + */ + allowed_cache_controls: unknown[] | null; + /** Allowed Passthrough Routes */ + allowed_passthrough_routes?: unknown[] | null; + /** + * Allowed Routes + * @default [] + */ + allowed_routes: unknown[] | null; + /** Allowed Vector Store Indexes */ + allowed_vector_store_indexes?: components["schemas"]["AllowedVectorStoreIndexItem"][] | null; + /** Auto Rotate */ + auto_rotate?: boolean | null; + /** Blocked */ + blocked?: boolean | null; + /** Budget Duration */ + budget_duration?: string | null; + /** Budget Fallbacks */ + budget_fallbacks?: { + [key: string]: string[]; + } | null; + /** Budget Id */ + budget_id?: string | null; + /** Budget Limits */ + budget_limits?: components["schemas"]["BudgetLimitEntry"][] | null; + /** + * Config + * @default {} + */ + config: { + [key: string]: unknown; + } | null; + /** Default Estimated Output Tokens */ + default_estimated_output_tokens?: number | null; + /** Default Estimated Output Tokens Per Model */ + default_estimated_output_tokens_per_model?: { + [key: string]: number; + } | null; + /** Disable Global Guardrails */ + disable_global_guardrails?: boolean | null; + /** Duration */ + duration?: string | null; + /** Enable Prompt Caching */ + enable_prompt_caching?: boolean | null; + /** Enforced Params */ + enforced_params?: string[] | null; + /** Guardrails */ + guardrails?: string[] | null; + /** Key */ + key?: string | null; + /** Key Alias */ + key_alias?: string | null; + /** Max Budget */ + max_budget?: number | null; + /** Max Parallel Requests */ + max_parallel_requests?: number | null; + /** Mcp Rpm Limit */ + mcp_rpm_limit?: { + [key: string]: number; + } | null; + /** Metadata */ + metadata?: { + [key: string]: unknown; + } | null; + /** + * Model Max Budget + * @default {} + */ + model_max_budget: { + [key: string]: unknown; + } | null; + /** Model Rpm Limit */ + model_rpm_limit?: { + [key: string]: unknown; + } | null; + /** Model Tpm Limit */ + model_tpm_limit?: { + [key: string]: unknown; + } | null; + /** + * Models + * @default [] + */ + models: unknown[] | null; + object_permission?: components["schemas"]["LiteLLM_ObjectPermissionBase"] | null; + /** Organization Id */ + organization_id?: string | null; + /** + * Permissions + * @default {} + */ + permissions: { + [key: string]: unknown; + } | null; + /** Policies */ + policies?: string[] | null; + /** Prompts */ + prompts?: string[] | null; + /** Rotation Interval */ + rotation_interval?: string | null; + router_settings?: components["schemas"]["UpdateRouterConfig"] | null; + /** Rpm Limit */ + rpm_limit?: number | null; + /** Rpm Limit Type */ + rpm_limit_type?: ("guaranteed_throughput" | "best_effort_throughput" | "dynamic") | null; + /** Spend */ + spend?: number | null; + /** Tag Rpm Limit */ + tag_rpm_limit?: { + [key: string]: number; + } | null; + /** Tags */ + tags?: string[] | null; + /** Team Id */ + team_id?: string | null; + /** Temp Budget Expiry */ + temp_budget_expiry?: string | null; + /** Temp Budget Increase */ + temp_budget_increase?: number | null; + /** Throttle On Budget Exceeded */ + throttle_on_budget_exceeded?: boolean | null; + /** Tpm Limit */ + tpm_limit?: number | null; + /** Tpm Limit Type */ + tpm_limit_type?: ("guaranteed_throughput" | "best_effort_throughput" | "dynamic") | null; + /** User Id */ + user_id?: string | null; + }; /** PatchPromptRequest */ PatchPromptRequest: { litellm_params?: components["schemas"]["PromptLiteLLMParams"] | null; @@ -45575,6 +45761,44 @@ export interface operations { }; }; }; + patch_key_key__key__patch: { + parameters: { + query?: never; + header?: { + /** @description The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability */ + "litellm-changed-by"?: string | null; + }; + path: { + key: string; + }; + cookie?: never; + }; + requestBody: { + content: { + "application/json": components["schemas"]["PatchKeyRequest"]; + }; + }; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; regenerate_key_fn_key__key__regenerate_post: { parameters: { query?: never;