From a5dd01d8aca49a94433d4a430958fe369ac18a2f Mon Sep 17 00:00:00 2001 From: kothamah <104782493+kothamah@users.noreply.github.com> Date: Thu, 19 Mar 2026 13:52:23 -0400 Subject: [PATCH 001/420] added bedrock guardrail API exception --- .../guardrail_hooks/bedrock_guardrails.py | 37 ++++++++++++------- 1 file changed, 23 insertions(+), 14 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 8ef188bb23c..02d87a4962d 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -82,7 +82,12 @@ def _redact_pii_matches(response_json: dict) -> dict: redacted_response = copy.deepcopy(response_json) # Get assessments from the response - assessments = redacted_response.get("assessments", []) + # NOTE: We use `.get("key") or []` instead of `.get("key", [])` because + # the Bedrock API can return explicit `null` for list fields (e.g. "regexes": null). + # In Python, dict.get("key", []) returns None (not []) when the key exists + # with a None/null value. The `or []` ensures we always get an iterable, + # preventing "TypeError: 'NoneType' object is not iterable". + assessments = redacted_response.get("assessments") or [] if not assessments: return redacted_response @@ -90,13 +95,13 @@ def _redact_pii_matches(response_json: dict) -> dict: # Redact PII entities in sensitive information policy sensitive_info_policy = assessment.get("sensitiveInformationPolicy") if sensitive_info_policy: - pii_entities = sensitive_info_policy.get("piiEntities", []) + pii_entities = sensitive_info_policy.get("piiEntities") or [] for pii_entity in pii_entities: if "match" in pii_entity: pii_entity["match"] = "[REDACTED]" # Redact regex matches - regexes = sensitive_info_policy.get("regexes", []) + regexes = sensitive_info_policy.get("regexes") or [] for regex_match in regexes: if "match" in regex_match: regex_match["match"] = "[REDACTED]" @@ -104,12 +109,12 @@ def _redact_pii_matches(response_json: dict) -> dict: # Redact custom word matches in word policy word_policy = assessment.get("wordPolicy") if word_policy: - custom_words = word_policy.get("customWords", []) + custom_words = word_policy.get("customWords") or [] for custom_word in custom_words: if "match" in custom_word: custom_word["match"] = "[REDACTED]" - managed_words = word_policy.get("managedWordLists", []) + managed_words = word_policy.get("managedWordLists") or [] for managed_word in managed_words: if "match" in managed_word: managed_word["match"] = "[REDACTED]" @@ -682,7 +687,9 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): return False # Check assessments to determine if any actions were BLOCKED (vs ANONYMIZED) - assessments = response.get("assessments", []) + # NOTE: Use `or []` instead of default param to handle explicit null from Bedrock API. + # See _redact_pii_matches() for detailed explanation of the null safety pattern. + assessments = response.get("assessments") or [] if not assessments: return False @@ -690,7 +697,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): # Check topic policy topic_policy = assessment.get("topicPolicy") if topic_policy: - topics = topic_policy.get("topics", []) + topics = topic_policy.get("topics") or [] for topic in topics: if topic.get("action") == "BLOCKED": return True @@ -698,7 +705,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): # Check content policy content_policy = assessment.get("contentPolicy") if content_policy: - filters = content_policy.get("filters", []) + filters = content_policy.get("filters") or [] for filter_item in filters: if filter_item.get("action") == "BLOCKED": return True @@ -706,11 +713,11 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): # Check word policy word_policy = assessment.get("wordPolicy") if word_policy: - custom_words = word_policy.get("customWords", []) + custom_words = word_policy.get("customWords") or [] for custom_word in custom_words: if custom_word.get("action") == "BLOCKED": return True - managed_words = word_policy.get("managedWordLists", []) + managed_words = word_policy.get("managedWordLists") or [] for managed_word in managed_words: if managed_word.get("action") == "BLOCKED": return True @@ -718,12 +725,12 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): # Check sensitive information policy sensitive_info_policy = assessment.get("sensitiveInformationPolicy") if sensitive_info_policy: - pii_entities = sensitive_info_policy.get("piiEntities", []) + pii_entities = sensitive_info_policy.get("piiEntities") or [] if pii_entities: for pii_entity in pii_entities: if pii_entity.get("action") == "BLOCKED": return True - regexes = sensitive_info_policy.get("regexes", []) + regexes = sensitive_info_policy.get("regexes") or [] if regexes: for regex in regexes: if regex.get("action") == "BLOCKED": @@ -732,7 +739,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): # Check contextual grounding policy contextual_grounding_policy = assessment.get("contextualGroundingPolicy") if contextual_grounding_policy: - grounding_filters = contextual_grounding_policy.get("filters", []) + grounding_filters = contextual_grounding_policy.get("filters") or [] for grounding_filter in grounding_filters: if grounding_filter.get("action") == "BLOCKED": return True @@ -1391,7 +1398,9 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): Raises: Exception: If content is blocked by Bedrock guardrail """ - texts = inputs.get("texts", []) + # NOTE: Use `or []` to handle case where inputs["texts"] is explicitly None. + # dict.get("texts", []) would return None if the key exists with a None value. + texts = inputs.get("texts") or [] try: verbose_proxy_logger.debug( f"Bedrock Guardrail: Applying guardrail to {len(texts)} text(s)" From ead822b6985443827d3accb7c0d6cfb7fce40df7 Mon Sep 17 00:00:00 2001 From: kothamah <104782493+kothamah@users.noreply.github.com> Date: Thu, 19 Mar 2026 13:55:36 -0400 Subject: [PATCH 002/420] Added test cases for the null type handling --- .../test_bedrock_guardrails.py | 2972 ++++++++--------- 1 file changed, 1472 insertions(+), 1500 deletions(-) diff --git a/tests/guardrails_tests/test_bedrock_guardrails.py b/tests/guardrails_tests/test_bedrock_guardrails.py index cb594c221c4..03fe63b307a 100644 --- a/tests/guardrails_tests/test_bedrock_guardrails.py +++ b/tests/guardrails_tests/test_bedrock_guardrails.py @@ -1,1090 +1,213 @@ -import sys +""" +Unit tests for Bedrock Guardrails +""" +import json import os -import io, asyncio +import sys +from unittest.mock import AsyncMock, MagicMock, patch + import pytest +from fastapi import HTTPException + +sys.path.insert(0, os.path.abspath("../../../../../..")) -sys.path.insert(0, os.path.abspath("../..")) -import litellm -from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import BedrockGuardrail from litellm.proxy._types import UserAPIKeyAuth -from litellm.caching import DualCache -from unittest.mock import MagicMock, AsyncMock, patch +from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( + BedrockGuardrail, + _redact_pii_matches, +) @pytest.mark.asyncio -async def test_bedrock_guardrails_pii_masking(): - # Create proper mock objects - mock_user_api_key_dict = UserAPIKeyAuth() +async def test__redact_pii_matches_function(): + """Test the _redact_pii_matches function directly""" - guardrail = BedrockGuardrail( - guardrailIdentifier="wf0hkdb5x07f", - guardrailVersion="DRAFT", - ) - - request_data = { - "model": "gpt-4o", - "messages": [ - {"role": "user", "content": "Hello, my phone number is +1 412 555 1212"}, - {"role": "assistant", "content": "Hello, how can I help you today?"}, - {"role": "user", "content": "I need to cancel my order"}, - { - "role": "user", - "content": "ok, my credit card number is 1234-5678-9012-3456", - }, - ], - } - - response = await guardrail.async_moderation_hook( - data=request_data, - user_api_key_dict=mock_user_api_key_dict, - call_type="completion", - ) - print("response after moderation hook", response) - - if response: # Only assert if response is not None - assert response["messages"][0]["content"] == "Hello, my phone number is {PHONE}" - assert response["messages"][1]["content"] == "Hello, how can I help you today?" - assert response["messages"][2]["content"] == "I need to cancel my order" - assert ( - response["messages"][3]["content"] - == "ok, my credit card number is {CREDIT_DEBIT_CARD_NUMBER}" - ) - - -@pytest.mark.asyncio -async def test_bedrock_guardrails_pii_masking_content_list(): - # Create proper mock objects - mock_user_api_key_dict = UserAPIKeyAuth() - - guardrail = BedrockGuardrail( - guardrailIdentifier="wf0hkdb5x07f", - guardrailVersion="DRAFT", - ) - - request_data = { - "model": "gpt-4o", - "messages": [ - { - "role": "user", - "content": [ - { - "type": "text", - "text": "Hello, my phone number is +1 412 555 1212", - }, - {"type": "text", "text": "what time is it?"}, - ], - }, - {"role": "assistant", "content": "Hello, how can I help you today?"}, - {"role": "user", "content": "who is the president of the united states?"}, - ], - } - - response = await guardrail.async_moderation_hook( - data=request_data, - user_api_key_dict=mock_user_api_key_dict, - call_type="completion", - ) - print(response) - - if response: # Only assert if response is not None - # Verify that the list content is properly masked - assert isinstance(response["messages"][0]["content"], list) - assert ( - response["messages"][0]["content"][0]["text"] - == "Hello, my phone number is {PHONE}" - ) - assert response["messages"][0]["content"][1]["text"] == "what time is it?" - assert response["messages"][1]["content"] == "Hello, how can I help you today?" - assert ( - response["messages"][2]["content"] - == "who is the president of the united states?" - ) - - -@pytest.mark.asyncio -async def test_bedrock_guardrails_block_messages_api(): - """ - Test that guardrails block messages API requests containing 'coffee' and raise the expected exception. - """ - from fastapi import HTTPException - - # Create proper mock objects - mock_user_api_key_dict = UserAPIKeyAuth() - - guardrail = BedrockGuardrail( - guardrailIdentifier="ff6ujrregl1q", - guardrailVersion="DRAFT", - ) - - request_data = { - "model": "claude-3-5-sonnet-20240620", - "messages": [ - { - "role": "user", - "content": [ - { - "type": "text", - "text": "Hello, my phone number is +1 412 555 1212", - }, - {"type": "text", "text": "what time is it?"}, - ], - }, - {"role": "user", "content": "tell me about coffee"}, - ], - } - - with pytest.raises(HTTPException) as exc_info: - await guardrail.async_pre_call_hook( - data=request_data, - user_api_key_dict=mock_user_api_key_dict, - call_type="anthropic_messages", - cache=MagicMock(spec=DualCache), - ) - - exception = exc_info.value - assert exception.status_code == 400 - detail = exception.detail - assert isinstance(detail, dict) - assert detail["error"] == "Violated guardrail policy" - assert ( - detail["bedrock_guardrail_response"] - == "Sorry, the model cannot answer this question. coffee guardrail applied " - ) - - -@pytest.mark.asyncio -async def test_bedrock_guardrails_block_responses_api(): - """ - Test that guardrails block responses API requests containing 'coffee' and raise the expected exception. - """ - from fastapi import HTTPException - - # Create proper mock objects - mock_user_api_key_dict = UserAPIKeyAuth() - - guardrail = BedrockGuardrail( - guardrailIdentifier="ff6ujrregl1q", - guardrailVersion="DRAFT", - ) - - request_data = { - "model": "gpt-4.1", - "input": "Tell me a three sentence bedtime story about a unicorn drinking coffee", - "stream": False, - } - - with pytest.raises(HTTPException) as exc_info: - await guardrail.async_pre_call_hook( - data=request_data, - user_api_key_dict=mock_user_api_key_dict, - call_type="responses", - cache=MagicMock(spec=DualCache), - ) - - exception = exc_info.value - assert exception.status_code == 400 - detail = exception.detail - assert isinstance(detail, dict) - assert detail["error"] == "Violated guardrail policy" - assert ( - detail["bedrock_guardrail_response"] - == "Sorry, the model cannot answer this question. coffee guardrail applied " - ) - - -@pytest.mark.asyncio -async def test_bedrock_guardrails_with_streaming(): - from litellm.proxy.utils import ProxyLogging - from litellm.types.guardrails import GuardrailEventHooks - - # Create proper mock objects - mock_user_api_key_cache = MagicMock(spec=DualCache) - mock_user_api_key_dict = UserAPIKeyAuth() - - with pytest.raises(Exception): # Assert that this raises an exception - proxy_logging_obj = ProxyLogging( - user_api_key_cache=mock_user_api_key_cache, - premium_user=True, - ) - - guardrail = BedrockGuardrail( - guardrailIdentifier="ff6ujrregl1q", - guardrailVersion="DRAFT", - supported_event_hooks=[GuardrailEventHooks.post_call], - guardrail_name="bedrock-post-guard", - ) - - litellm.callbacks.append(guardrail) - - request_data = { - "model": "gpt-4o", - "messages": [{"role": "user", "content": "Hi I like coffee"}], - "stream": True, - "metadata": {"guardrails": ["bedrock-post-guard"]}, - } - - response = await litellm.acompletion( - **request_data, - ) - - response = proxy_logging_obj.async_post_call_streaming_iterator_hook( - user_api_key_dict=mock_user_api_key_dict, - response=response, - request_data=request_data, - ) - - async for chunk in response: - print(chunk) - - -@pytest.mark.asyncio -async def test_bedrock_guardrails_with_streaming_no_violation(): - from litellm.proxy.utils import ProxyLogging - from litellm.types.guardrails import GuardrailEventHooks - - # Create proper mock objects - mock_user_api_key_cache = MagicMock(spec=DualCache) - mock_user_api_key_dict = UserAPIKeyAuth() - - proxy_logging_obj = ProxyLogging( - user_api_key_cache=mock_user_api_key_cache, - premium_user=True, - ) - - guardrail = BedrockGuardrail( - guardrailIdentifier="ff6ujrregl1q", - guardrailVersion="DRAFT", - supported_event_hooks=[GuardrailEventHooks.post_call], - guardrail_name="bedrock-post-guard", - ) - - litellm.callbacks.append(guardrail) - - request_data = { - "model": "gpt-4o", - "messages": [{"role": "user", "content": "hi"}], - "stream": True, - "metadata": {"guardrails": ["bedrock-post-guard"]}, - } - - response = await litellm.acompletion( - **request_data, - ) - - response = proxy_logging_obj.async_post_call_streaming_iterator_hook( - user_api_key_dict=mock_user_api_key_dict, - response=response, - request_data=request_data, - ) - - async for chunk in response: - print(chunk) - - -@pytest.mark.asyncio -async def test_bedrock_guardrails_streaming_request_body_mock(): - """Test that the exact request body sent to Bedrock matches expected format when using streaming""" - import json - from unittest.mock import AsyncMock, MagicMock, patch - from litellm.proxy._types import UserAPIKeyAuth - from litellm.caching import DualCache - from litellm.types.guardrails import GuardrailEventHooks - - # Create mock objects - mock_user_api_key_dict = UserAPIKeyAuth() - mock_cache = MagicMock(spec=DualCache) - - # Create the guardrail - guardrail = BedrockGuardrail( - guardrailIdentifier="wf0hkdb5x07f", - guardrailVersion="DRAFT", - supported_event_hooks=[GuardrailEventHooks.post_call], - guardrail_name="bedrock-post-guard", - ) - - # Mock the assembled response from streaming - mock_response = litellm.ModelResponse( - id="test-id", - choices=[ - litellm.Choices( - index=0, - message=litellm.Message( - role="assistant", content="The capital of Spain is Madrid." - ), - finish_reason="stop", - ) - ], - created=1234567890, - model="gpt-4o", - object="chat.completion", - ) - - # Mock Bedrock API response - mock_bedrock_response = MagicMock() - mock_bedrock_response.status_code = 200 - mock_bedrock_response.json.return_value = {"action": "NONE", "outputs": []} - - # Patch the async_handler.post method to capture the request body - with patch.object(guardrail, "async_handler") as mock_async_handler: - mock_async_handler.post = AsyncMock(return_value=mock_bedrock_response) - - # Test data - simulating request data and assembled response - request_data = { - "model": "gpt-4o", - "messages": [{"role": "user", "content": "what's the capital of spain?"}], - "stream": True, - "metadata": {"guardrails": ["bedrock-post-guard"]}, - } - - # Call the method that should make the Bedrock API request - await guardrail.make_bedrock_api_request( - source="OUTPUT", response=mock_response, request_data=request_data - ) - - # Verify the API call was made - mock_async_handler.post.assert_called_once() - - # Get the request data that was passed - call_args = mock_async_handler.post.call_args - - # The data should be in the 'data' parameter of the prepared request - # We need to parse the JSON from the prepared request body - prepared_request_body = call_args.kwargs.get("data") - - # Parse the JSON body - if isinstance(prepared_request_body, bytes): - actual_body = json.loads(prepared_request_body.decode("utf-8")) - else: - actual_body = json.loads(prepared_request_body) - - # Expected body based on the convert_to_bedrock_format method behavior - expected_body = { - "source": "OUTPUT", - "content": [{"text": {"text": "The capital of Spain is Madrid."}}], - } - - print("Actual Bedrock request body:", json.dumps(actual_body, indent=2)) - print("Expected Bedrock request body:", json.dumps(expected_body, indent=2)) - - # Assert the request body matches exactly - assert ( - actual_body == expected_body - ), f"Request body mismatch. Expected: {expected_body}, Got: {actual_body}" - - -@pytest.mark.asyncio -async def test_bedrock_guardrail_aws_param_persistence(): - """Test that AWS auth params set on init are used for every request and not popped out.""" - from litellm.proxy._types import UserAPIKeyAuth - from litellm.types.guardrails import GuardrailEventHooks - - guardrail = BedrockGuardrail( - guardrailIdentifier="wf0hkdb5x07f", - guardrailVersion="DRAFT", - aws_access_key_id="test-access-key", - aws_secret_access_key="test-secret-key", - aws_region_name="us-east-1", - supported_event_hooks=[GuardrailEventHooks.post_call], - guardrail_name="bedrock-post-guard", - ) - - with patch.object( - guardrail, "get_credentials", wraps=guardrail.get_credentials - ) as mock_get_creds: - for i in range(3): - request_data = { - "model": "gpt-4o", - "messages": [{"role": "user", "content": f"request {i}"}], - "stream": False, - "metadata": {"guardrails": ["bedrock-post-guard"]}, - } - with patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post: - # Configure the mock response properly - mock_response = AsyncMock() - mock_response.status_code = 200 - mock_response.json = MagicMock( - return_value={"action": "NONE", "outputs": []} - ) - mock_post.return_value = mock_response - await guardrail.make_bedrock_api_request( - source="INPUT", - messages=request_data.get("messages"), - request_data=request_data, - ) - - assert mock_get_creds.call_count == 3 - for call in mock_get_creds.call_args_list: - kwargs = call.kwargs - print("used the following kwargs to get credentials=", kwargs) - assert kwargs["aws_access_key_id"] == "test-access-key" - assert kwargs["aws_secret_access_key"] == "test-secret-key" - assert kwargs["aws_region_name"] == "us-east-1" - - -@pytest.mark.asyncio -async def test_bedrock_guardrail_blocked_vs_anonymized_actions(): - """Test that BLOCKED actions raise exceptions but ANONYMIZED actions do not""" - from unittest.mock import MagicMock - from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( - BedrockGuardrail, - ) - from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( - BedrockGuardrailResponse, - ) - - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) - - # Test 1: ANONYMIZED action should NOT raise exception - anonymized_response: BedrockGuardrailResponse = { + # Test case 1: Response with PII entities + response_with_pii = { "action": "GUARDRAIL_INTERVENED", - "outputs": [{"text": "Hello, my phone number is {PHONE}"}], "assessments": [ { "sensitiveInformationPolicy": { "piiEntities": [ + {"type": "NAME", "match": "John Smith", "action": "BLOCKED"}, { - "type": "PHONE", - "match": "+1 412 555 1212", - "action": "ANONYMIZED", - } - ] - } - } - ], - } - - should_raise = guardrail._should_raise_guardrail_blocked_exception( - anonymized_response - ) - assert should_raise is False, "ANONYMIZED actions should not raise exceptions" - - # Test 2: BLOCKED action should raise exception - blocked_response: BedrockGuardrailResponse = { - "action": "GUARDRAIL_INTERVENED", - "outputs": [{"text": "I can't provide that information."}], - "assessments": [ - { - "topicPolicy": { - "topics": [ - {"name": "Sensitive Topic", "type": "DENY", "action": "BLOCKED"} - ] - } - } - ], - } - - should_raise = guardrail._should_raise_guardrail_blocked_exception(blocked_response) - assert should_raise is True, "BLOCKED actions should raise exceptions" - - # Test 3: Mixed actions - should raise if ANY action is BLOCKED - mixed_response: BedrockGuardrailResponse = { - "action": "GUARDRAIL_INTERVENED", - "outputs": [{"text": "I can't provide that information."}], - "assessments": [ - { - "sensitiveInformationPolicy": { - "piiEntities": [ - { - "type": "PHONE", - "match": "+1 412 555 1212", - "action": "ANONYMIZED", - } - ] - }, - "topicPolicy": { - "topics": [ - {"name": "Blocked Topic", "type": "DENY", "action": "BLOCKED"} - ] - }, - } - ], - } - - should_raise = guardrail._should_raise_guardrail_blocked_exception(mixed_response) - assert ( - should_raise is True - ), "Mixed actions with any BLOCKED should raise exceptions" - - # Test 4: NONE action should not raise exception - none_response: BedrockGuardrailResponse = { - "action": "NONE", - "outputs": [], - "assessments": [], - } - - should_raise = guardrail._should_raise_guardrail_blocked_exception(none_response) - assert should_raise is False, "NONE actions should not raise exceptions" - - # Test 5: Test other policy types with BLOCKED actions - content_blocked_response: BedrockGuardrailResponse = { - "action": "GUARDRAIL_INTERVENED", - "outputs": [{"text": "I can't provide that information."}], - "assessments": [ - { - "contentPolicy": { - "filters": [ - {"type": "VIOLENCE", "confidence": "HIGH", "action": "BLOCKED"} - ] - } - } - ], - } - - should_raise = guardrail._should_raise_guardrail_blocked_exception( - content_blocked_response - ) - assert ( - should_raise is True - ), "Content policy BLOCKED actions should raise exceptions" - - -@pytest.mark.asyncio -async def test_bedrock_guardrail_masking_with_anonymized_response(): - """Test that masking works correctly when guardrail returns ANONYMIZED actions""" - from unittest.mock import AsyncMock, MagicMock, patch - from litellm.proxy._types import UserAPIKeyAuth - from litellm.caching import DualCache - - # Create proper mock objects - mock_user_api_key_dict = UserAPIKeyAuth() - - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", - guardrailVersion="DRAFT", - mask_request_content=True, - ) - - # Mock the Bedrock API response with ANONYMIZED action - mock_bedrock_response = MagicMock() - mock_bedrock_response.status_code = 200 - mock_bedrock_response.json.return_value = { - "action": "GUARDRAIL_INTERVENED", - "outputs": [{"text": "Hello, my phone number is {PHONE}"}], - "assessments": [ - { - "sensitiveInformationPolicy": { - "piiEntities": [ - { - "type": "PHONE", - "match": "+1 412 555 1212", - "action": "ANONYMIZED", - } - ] - } - } - ], - } - - request_data = { - "model": "gpt-4o", - "messages": [ - {"role": "user", "content": "Hello, my phone number is +1 412 555 1212"}, - ], - } - - # Patch the async_handler.post method - with patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post: - mock_post.return_value = mock_bedrock_response - - # This should NOT raise an exception since action is ANONYMIZED - try: - response = await guardrail.async_moderation_hook( - data=request_data, - user_api_key_dict=mock_user_api_key_dict, - call_type="completion", - ) - # Should succeed and return data with masked content - assert response is not None - assert ( - response["messages"][0]["content"] - == "Hello, my phone number is {PHONE}" - ) - except Exception as e: - pytest.fail( - f"Should not raise exception for ANONYMIZED actions, but got: {e}" - ) - - -@pytest.mark.asyncio -async def test_bedrock_guardrail_uses_masked_output_without_masking_flags(): - """Test that masked output from guardrails is used even when masking flags are not enabled""" - from unittest.mock import AsyncMock, MagicMock, patch - from litellm.proxy._types import UserAPIKeyAuth - - # Create proper mock objects - mock_user_api_key_dict = UserAPIKeyAuth() - - # Create guardrail WITHOUT masking flags enabled - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", - guardrailVersion="DRAFT", - # Note: No mask_request_content=True or mask_response_content=True - ) - - # Mock the Bedrock API response with ANONYMIZED action and masked output - mock_bedrock_response = MagicMock() - mock_bedrock_response.status_code = 200 - mock_bedrock_response.json.return_value = { - "action": "GUARDRAIL_INTERVENED", - "outputs": [{"text": "Hello, my phone number is {PHONE} and email is {EMAIL}"}], - "assessments": [ - { - "sensitiveInformationPolicy": { - "piiEntities": [ - { - "type": "PHONE", - "match": "+1 412 555 1212", - "action": "ANONYMIZED", + "type": "US_SOCIAL_SECURITY_NUMBER", + "match": "324-12-3212", + "action": "BLOCKED", }, + {"type": "PHONE", "match": "607-456-7890", "action": "BLOCKED"}, + ] + } + } + ], + "outputs": [{"text": "Input blocked by PII policy"}], + } + + # Call the redaction function + redacted_response = _redact_pii_matches(response_with_pii) + + # Verify that PII matches are redacted + pii_entities = redacted_response["assessments"][0]["sensitiveInformationPolicy"][ + "piiEntities" + ] + + assert pii_entities[0]["match"] == "[REDACTED]", "Name should be redacted" + assert pii_entities[1]["match"] == "[REDACTED]", "SSN should be redacted" + assert pii_entities[2]["match"] == "[REDACTED]", "Phone should be redacted" + + # Verify other fields remain unchanged + assert pii_entities[0]["type"] == "NAME" + assert pii_entities[1]["type"] == "US_SOCIAL_SECURITY_NUMBER" + assert pii_entities[2]["type"] == "PHONE" + assert redacted_response["action"] == "GUARDRAIL_INTERVENED" + assert redacted_response["outputs"][0]["text"] == "Input blocked by PII policy" + + print("PII redaction function test passed") + + +@pytest.mark.asyncio +async def test__redact_pii_matches_no_pii(): + """Test _redact_pii_matches with response that has no PII""" + + response_no_pii = {"action": "NONE", "assessments": [], "outputs": []} + + # Call the redaction function + redacted_response = _redact_pii_matches(response_no_pii) + + # Should return the same response unchanged + assert redacted_response == response_no_pii + print("No PII redaction test passed") + + +@pytest.mark.asyncio +async def test__redact_pii_matches_empty_assessments(): + """Test _redact_pii_matches with empty assessments""" + + response_empty_assessments = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [{"sensitiveInformationPolicy": {"piiEntities": []}}], + "outputs": [{"text": "Some output"}], + } + + # Call the redaction function + redacted_response = _redact_pii_matches(response_empty_assessments) + + # Should return the same response unchanged + assert redacted_response == response_empty_assessments + print("Empty assessments redaction test passed") + + +@pytest.mark.asyncio +async def test__redact_pii_matches_malformed_response(): + """Test _redact_pii_matches with malformed response (should not crash)""" + + # Test with completely malformed response + malformed_response = { + "action": "GUARDRAIL_INTERVENED", + "assessments": "not_a_list", # This should cause an exception + } + + # Should not crash and return original response + redacted_response = _redact_pii_matches(malformed_response) + assert redacted_response == malformed_response + + # Test with missing keys + missing_keys_response = { + "action": "GUARDRAIL_INTERVENED" + # Missing assessments key + } + + redacted_response = _redact_pii_matches(missing_keys_response) + assert redacted_response == missing_keys_response + + print("Malformed response redaction test passed") + + +@pytest.mark.asyncio +async def test__redact_pii_matches_multiple_assessments(): + """Test _redact_pii_matches with multiple assessments containing PII""" + + response_multiple_assessments = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [ + { + "sensitiveInformationPolicy": { + "piiEntities": [ { "type": "EMAIL", - "match": "user@example.com", + "match": "john@example.com", "action": "ANONYMIZED", - }, + } ] } - } - ], - } - - request_data = { - "model": "gpt-4o", - "messages": [ - { - "role": "user", - "content": "Hello, my phone number is +1 412 555 1212 and email is user@example.com", }, - ], - } - - # Patch the async_handler.post method - with patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post: - mock_post.return_value = mock_bedrock_response - - # This should use the masked output even without masking flags - response = await guardrail.async_moderation_hook( - data=request_data, - user_api_key_dict=mock_user_api_key_dict, - call_type="completion", - ) - - # Should use the masked content from guardrail output - assert response is not None - assert ( - response["messages"][0]["content"] - == "Hello, my phone number is {PHONE} and email is {EMAIL}" - ) - print("✅ Masked output was applied even without masking flags enabled") - - -@pytest.mark.asyncio -async def test_bedrock_guardrail_response_pii_masking_non_streaming(): - """Test that PII masking is applied to response content in non-streaming scenarios""" - from unittest.mock import AsyncMock, MagicMock, patch - from litellm.proxy._types import UserAPIKeyAuth - - # Create proper mock objects - mock_user_api_key_dict = UserAPIKeyAuth() - - # Create guardrail with response masking enabled - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", - guardrailVersion="DRAFT", - ) - - # Mock the Bedrock API response with ANONYMIZED PII - mock_bedrock_response = MagicMock() - mock_bedrock_response.status_code = 200 - mock_bedrock_response.json.return_value = { - "action": "GUARDRAIL_INTERVENED", - "outputs": [ - { - "text": "My credit card number is {CREDIT_DEBIT_CARD_NUMBER} and my phone is {PHONE}" - } - ], - "assessments": [ { "sensitiveInformationPolicy": { "piiEntities": [ { "type": "CREDIT_DEBIT_CARD_NUMBER", "match": "1234-5678-9012-3456", + "action": "BLOCKED", + }, + { + "type": "ADDRESS", + "match": "123 Main St, Anytown USA", "action": "ANONYMIZED", }, + ] + } + }, + ], + "outputs": [{"text": "Multiple PII detected"}], + } + + # Call the redaction function + redacted_response = _redact_pii_matches(response_multiple_assessments) + + # Verify all PII in all assessments are redacted + assessment1_pii = redacted_response["assessments"][0]["sensitiveInformationPolicy"][ + "piiEntities" + ] + assessment2_pii = redacted_response["assessments"][1]["sensitiveInformationPolicy"][ + "piiEntities" + ] + + assert assessment1_pii[0]["match"] == "[REDACTED]", "Email should be redacted" + assert assessment2_pii[0]["match"] == "[REDACTED]", "Credit card should be redacted" + assert assessment2_pii[1]["match"] == "[REDACTED]", "Address should be redacted" + + # Verify types remain unchanged + assert assessment1_pii[0]["type"] == "EMAIL" + assert assessment2_pii[0]["type"] == "CREDIT_DEBIT_CARD_NUMBER" + assert assessment2_pii[1]["type"] == "ADDRESS" + + print("Multiple assessments redaction test passed") + + +@pytest.mark.asyncio +async def test_bedrock_guardrail_logging_uses_redacted_response(): + """Test that the Bedrock guardrail uses redacted response for logging""" + + # Create proper mock objects + mock_user_api_key_dict = UserAPIKeyAuth() + + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" + ) + + # Mock the Bedrock API response with PII + mock_bedrock_response = MagicMock() + mock_bedrock_response.status_code = 200 + mock_bedrock_response.json.return_value = { + "action": "GUARDRAIL_INTERVENED", + "outputs": [{"text": "Hello, my phone number is {PHONE}"}], + "assessments": [ + { + "sensitiveInformationPolicy": { + "piiEntities": [ { "type": "PHONE", - "match": "+1 412 555 1212", - "action": "ANONYMIZED", - }, - ] - } - } - ], - } - - # Create a mock response that contains PII - mock_response = litellm.ModelResponse( - id="test-id", - choices=[ - litellm.Choices( - index=0, - message=litellm.Message( - role="assistant", - content="My credit card number is 1234-5678-9012-3456 and my phone is +1 412 555 1212", - ), - finish_reason="stop", - ) - ], - created=1234567890, - model="gpt-4o", - object="chat.completion", - ) - - request_data = { - "model": "gpt-4o", - "messages": [ - {"role": "user", "content": "What's your credit card and phone number?"}, - ], - } - - # Patch the async_handler.post method - with patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post: - mock_post.return_value = mock_bedrock_response - - # Call the post-call success hook - await guardrail.async_post_call_success_hook( - data=request_data, - user_api_key_dict=mock_user_api_key_dict, - response=mock_response, - ) - - # Verify that the response content was masked - assert ( - mock_response.choices[0].message.content - == "My credit card number is {CREDIT_DEBIT_CARD_NUMBER} and my phone is {PHONE}" - ) - print("✓ Non-streaming response PII masking test passed") - - -@pytest.mark.asyncio -async def test_bedrock_guardrail_response_pii_masking_streaming(): - """Test that PII masking is applied to response content in streaming scenarios""" - from unittest.mock import AsyncMock, MagicMock, patch - from litellm.proxy._types import UserAPIKeyAuth - from litellm.types.utils import ModelResponseStream - - # Create proper mock objects - mock_user_api_key_dict = UserAPIKeyAuth() - - # Create guardrail with response masking enabled - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", - guardrailVersion="DRAFT", - ) - - # Mock the Bedrock API response with ANONYMIZED PII - mock_bedrock_response = MagicMock() - mock_bedrock_response.status_code = 200 - mock_bedrock_response.json.return_value = { - "action": "GUARDRAIL_INTERVENED", - "outputs": [{"text": "Sure! My email is {EMAIL} and SSN is {US_SSN}"}], - "assessments": [ - { - "sensitiveInformationPolicy": { - "piiEntities": [ - { - "type": "EMAIL", - "match": "john@example.com", - "action": "ANONYMIZED", - }, - { - "type": "US_SSN", - "match": "123-45-6789", - "action": "ANONYMIZED", - }, - ] - } - } - ], - } - - # Create mock streaming chunks - async def mock_streaming_response(): - chunks = [ - ModelResponseStream( - id="test-id", - choices=[ - litellm.utils.StreamingChoices( - index=0, - delta=litellm.utils.Delta(content="Sure! My email is "), - finish_reason=None, - ) - ], - created=1234567890, - model="gpt-4o", - object="chat.completion.chunk", - ), - ModelResponseStream( - id="test-id", - choices=[ - litellm.utils.StreamingChoices( - index=0, - delta=litellm.utils.Delta( - content="john@example.com and SSN is " - ), - finish_reason=None, - ) - ], - created=1234567890, - model="gpt-4o", - object="chat.completion.chunk", - ), - ModelResponseStream( - id="test-id", - choices=[ - litellm.utils.StreamingChoices( - index=0, - delta=litellm.utils.Delta(content="123-45-6789"), - finish_reason="stop", - ) - ], - created=1234567890, - model="gpt-4o", - object="chat.completion.chunk", - ), - ] - for chunk in chunks: - yield chunk - - request_data = { - "model": "gpt-4o", - "messages": [ - {"role": "user", "content": "What's your email and SSN?"}, - ], - "stream": True, - } - - # Patch the async_handler.post method - with patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post: - mock_post.return_value = mock_bedrock_response - - # Call the streaming hook - masked_stream = guardrail.async_post_call_streaming_iterator_hook( - user_api_key_dict=mock_user_api_key_dict, - response=mock_streaming_response(), - request_data=request_data, - ) - - # Collect all chunks from the masked stream - masked_chunks = [] - async for chunk in masked_stream: - masked_chunks.append(chunk) - - # Verify that we got chunks back - assert len(masked_chunks) > 0 - - # Reconstruct the full response from chunks to verify masking - full_content = "" - for chunk in masked_chunks: - if hasattr(chunk, "choices") and chunk.choices: - if hasattr(chunk.choices[0], "delta") and chunk.choices[0].delta: - if ( - hasattr(chunk.choices[0].delta, "content") - and chunk.choices[0].delta.content - ): - full_content += chunk.choices[0].delta.content - - # Verify that the reconstructed content contains the masked PII - assert "Sure! My email is {EMAIL} and SSN is {US_SSN}" == full_content - print("✓ Streaming response PII masking test passed") - - -@pytest.mark.asyncio -async def test_convert_to_bedrock_format_input_source(): - """Test convert_to_bedrock_format with INPUT source and mock messages""" - from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( - BedrockGuardrail, - ) - from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( - BedrockRequest, - ) - from unittest.mock import patch - - # Create the guardrail instance - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) - - # Mock messages - mock_messages = [ - {"role": "user", "content": "Hello, how are you?"}, - {"role": "assistant", "content": "I'm doing well, thank you!"}, - { - "role": "user", - "content": [ - {"type": "text", "text": "What's the weather like?"}, - {"type": "text", "text": "Is it sunny today?"}, - ], - }, - ] - - # Call the method - result = guardrail.convert_to_bedrock_format(source="INPUT", messages=mock_messages) - - # Verify the result structure - assert isinstance(result, dict) - assert result.get("source") == "INPUT" - assert "content" in result - assert isinstance(result.get("content"), list) - - # Verify content items - expected_content_items = [ - {"text": {"text": "Hello, how are you?"}}, - {"text": {"text": "I'm doing well, thank you!"}}, - {"text": {"text": "What's the weather like?"}}, - {"text": {"text": "Is it sunny today?"}}, - ] - - assert result.get("content") == expected_content_items - print("✅ INPUT source test passed - result:", result) - - -@pytest.mark.asyncio -async def test_convert_to_bedrock_format_output_source(): - """Test convert_to_bedrock_format with OUTPUT source and mock ModelResponse""" - from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( - BedrockGuardrail, - ) - from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( - BedrockRequest, - ) - import litellm - from unittest.mock import patch - - # Create the guardrail instance - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) - - # Mock ModelResponse - mock_response = litellm.ModelResponse( - id="test-response-id", - choices=[ - litellm.Choices( - index=0, - message=litellm.Message( - role="assistant", content="This is a test response from the model." - ), - finish_reason="stop", - ), - litellm.Choices( - index=1, - message=litellm.Message( - role="assistant", content="This is a second choice response." - ), - finish_reason="stop", - ), - ], - created=1234567890, - model="gpt-4o", - object="chat.completion", - ) - - # Call the method - result = guardrail.convert_to_bedrock_format( - source="OUTPUT", response=mock_response - ) - - # Verify the result structure - assert isinstance(result, dict) - assert result.get("source") == "OUTPUT" - assert "content" in result - assert isinstance(result.get("content"), list) - - # Verify content items - should contain both choice contents - expected_content_items = [ - {"text": {"text": "This is a test response from the model."}}, - {"text": {"text": "This is a second choice response."}}, - ] - - assert result.get("content") == expected_content_items - print("✅ OUTPUT source test passed - result:", result) - - -@pytest.mark.asyncio -async def test_convert_to_bedrock_format_post_call_streaming_hook(): - """Test async_post_call_streaming_iterator_hook makes OUTPUT bedrock request and applies masking""" - from unittest.mock import AsyncMock, MagicMock, patch - from litellm.proxy._types import UserAPIKeyAuth - from litellm.types.utils import ModelResponseStream - import litellm - - # Create proper mock objects - mock_user_api_key_dict = UserAPIKeyAuth() - - # Create guardrail instance - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) - - # Mock streaming chunks that contain PII - async def mock_streaming_response(): - chunks = [ - ModelResponseStream( - id="test-id", - choices=[ - litellm.utils.StreamingChoices( - index=0, - delta=litellm.utils.Delta(content="My email is "), - finish_reason=None, - ) - ], - created=1234567890, - model="gpt-4o", - object="chat.completion.chunk", - ), - ModelResponseStream( - id="test-id", - choices=[ - litellm.utils.StreamingChoices( - index=0, - delta=litellm.utils.Delta(content="john@example.com"), - finish_reason="stop", - ) - ], - created=1234567890, - model="gpt-4o", - object="chat.completion.chunk", - ), - ] - for chunk in chunks: - yield chunk - - # Mock Bedrock API response with PII masking - mock_bedrock_response = MagicMock() - mock_bedrock_response.status_code = 200 - mock_bedrock_response.json.return_value = { - "action": "GUARDRAIL_INTERVENED", - "outputs": [{"text": "My email is {EMAIL}"}], - "assessments": [ - { - "sensitiveInformationPolicy": { - "piiEntities": [ - { - "type": "EMAIL", - "match": "john@example.com", + "match": "+1 412 555 1212", # This should be redacted in logs "action": "ANONYMIZED", } ] @@ -1095,99 +218,82 @@ async def test_convert_to_bedrock_format_post_call_streaming_hook(): request_data = { "model": "gpt-4o", - "messages": [{"role": "user", "content": "What's your email?"}], - "stream": True, + "messages": [ + {"role": "user", "content": "Hello, my phone number is +1 412 555 1212"}, + ], } - # Track which bedrock API calls were made - bedrock_calls = [] + # Mock AWS credentials to avoid credential loading issues in CI + mock_credentials = MagicMock() + mock_credentials.access_key = "test-access-key" + mock_credentials.secret_key = "test-secret-key" + mock_credentials.token = None - # Mock the make_bedrock_api_request method to track calls - async def mock_make_bedrock_api_request( - source, messages=None, response=None, request_data=None - ): - bedrock_calls.append( - { - "source": source, - "messages": messages, - "response": response, - "request_data": request_data, - } - ) - # Return the mock bedrock response - from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( - BedrockGuardrailResponse, - ) - - return BedrockGuardrailResponse(**mock_bedrock_response.json()) - - # Patch the bedrock API request method + # Mock AWS-related methods to ensure test runs without external dependencies with patch.object( - guardrail, "make_bedrock_api_request", side_effect=mock_make_bedrock_api_request - ): + guardrail.async_handler, "post", new_callable=AsyncMock + ) as mock_post, patch( + "litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails.verbose_proxy_logger.debug" + ) as mock_debug, patch.object( + guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") + ) as mock_load_creds, patch.object( + guardrail, "_prepare_request", return_value=MagicMock() + ) as mock_prepare_request: - # Call the streaming hook - result_generator = guardrail.async_post_call_streaming_iterator_hook( - user_api_key_dict=mock_user_api_key_dict, - response=mock_streaming_response(), + mock_post.return_value = mock_bedrock_response + + # Call the method that should log the redacted response + await guardrail.make_bedrock_api_request( + source="INPUT", + messages=request_data.get("messages"), request_data=request_data, ) - # Collect all chunks from the result - result_chunks = [] - async for chunk in result_generator: - result_chunks.append(chunk) + # Verify that debug logging was called + mock_debug.assert_called() + + # Get the logged response (second argument to debug call) + logged_calls = mock_debug.call_args_list + bedrock_response_log_call = None + + for call in logged_calls: + args, kwargs = call + if len(args) >= 2 and "Bedrock AI response" in str(args[0]): + bedrock_response_log_call = call + break - # Verify bedrock API calls were made - # Note: When event_hook is None (default), the guardrail is considered enabled for all hooks. - # In post_call, INPUT validation is skipped if pre_call/during_call is already enabled - # to avoid redundant validation. Since event_hook=None means all hooks are enabled, - # only OUTPUT validation should be performed in post_call. assert ( - len(bedrock_calls) == 1 - ), f"Expected 1 bedrock call (OUTPUT only), got {len(bedrock_calls)}" + bedrock_response_log_call is not None + ), "Should have logged Bedrock AI response" - # Verify the OUTPUT call - output_call = bedrock_calls[0] - assert output_call["source"] == "OUTPUT" - assert output_call["response"] is not None - assert output_call["messages"] is None # OUTPUT calls don't need messages + # Extract the logged response data + logged_response = bedrock_response_log_call[0][ + 1 + ] # Second argument to debug call - # Verify that the response content was masked - # The streaming chunks should now contain the masked content - full_content = "" - for chunk in result_chunks: - if hasattr(chunk, "choices") and chunk.choices: - if ( - hasattr(chunk.choices[0], "delta") - and chunk.choices[0].delta.content - ): - full_content += chunk.choices[0].delta.content - - # The content should be masked (contains {EMAIL} instead of john@example.com) + # Verify that the logged response has redacted PII assert ( - "{EMAIL}" in full_content - ), f"Expected masked content with {{EMAIL}}, got: {full_content}" - assert ( - "john@example.com" not in full_content - ), f"Original email should be masked, got: {full_content}" - - print( - "✅ Post-call streaming hook test passed - OUTPUT source used for masking" + logged_response["assessments"][0]["sensitiveInformationPolicy"][ + "piiEntities" + ][0]["match"] + == "[REDACTED]" ) - print( - f"✅ Bedrock calls made: {[call['source'] for call in bedrock_calls]} " - "(INPUT validation skipped due to event_hook=None implying pre_call/during_call enabled)" + + # Verify other fields are preserved + assert logged_response["action"] == "GUARDRAIL_INTERVENED" + assert ( + logged_response["assessments"][0]["sensitiveInformationPolicy"][ + "piiEntities" + ][0]["type"] + == "PHONE" ) - print(f"✅ Final masked content: {full_content}") + + print("Bedrock guardrail logging redaction test passed") @pytest.mark.asyncio -async def test_bedrock_guardrail_blocked_action_shows_output_text(): - """Test that BLOCKED actions raise HTTPException with the output text in the detail""" - from unittest.mock import AsyncMock, MagicMock, patch - from litellm.proxy._types import UserAPIKeyAuth - from fastapi import HTTPException +async def test_bedrock_guardrail_original_response_not_modified(): + """Test that the original response is not modified by redaction, only the logged version""" # Create proper mock objects mock_user_api_key_dict = UserAPIKeyAuth() @@ -1196,408 +302,1274 @@ async def test_bedrock_guardrail_blocked_action_shows_output_text(): guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" ) - # Mock the Bedrock API response with BLOCKED action and output text - mock_bedrock_response = MagicMock() - mock_bedrock_response.status_code = 200 - mock_bedrock_response.json.return_value = { + # Mock the Bedrock API response with PII + original_response_data = { "action": "GUARDRAIL_INTERVENED", - "outputs": [{"text": "this violates litellm corporate guardrail policy"}], + "outputs": [{"text": "Hello, my phone number is {PHONE}"}], "assessments": [ { - "topicPolicy": { - "topics": [ - {"name": "Sensitive Topic", "type": "DENY", "action": "BLOCKED"} + "sensitiveInformationPolicy": { + "piiEntities": [ + { + "type": "PHONE", + "match": "+1 412 555 1212", # This should NOT be modified in original + "action": "ANONYMIZED", + } ] } } ], } + mock_bedrock_response = MagicMock() + mock_bedrock_response.status_code = 200 + mock_bedrock_response.json.return_value = original_response_data + request_data = { "model": "gpt-4o", "messages": [ - {"role": "user", "content": "Tell me how to make explosives"}, + {"role": "user", "content": "Hello, my phone number is +1 412 555 1212"}, ], } - # Patch the async_handler.post method + # Mock AWS credentials to avoid credential loading issues in CI + mock_credentials = MagicMock() + mock_credentials.access_key = "test-access-key" + mock_credentials.secret_key = "test-secret-key" + mock_credentials.token = None + + # Mock AWS-related methods to ensure test runs without external dependencies with patch.object( guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post: + ) as mock_post, patch.object( + guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") + ) as mock_load_creds, patch.object( + guardrail, "_prepare_request", return_value=MagicMock() + ) as mock_prepare_request: + mock_post.return_value = mock_bedrock_response - # This should raise HTTPException due to BLOCKED action - with pytest.raises(HTTPException) as exc_info: - await guardrail.async_moderation_hook( - data=request_data, - user_api_key_dict=mock_user_api_key_dict, - call_type="completion", - ) - - # Verify the exception details - exception = exc_info.value - assert exception.status_code == 400 - assert "detail" in exception.__dict__ - - # Check that the detail contains the expected structure - detail = exception.detail - assert isinstance(detail, dict) - assert detail["error"] == "Violated guardrail policy" - - # Verify that the output text from both outputs is included - expected_output_text = "this violates litellm corporate guardrail policy" - assert detail["bedrock_guardrail_response"] == expected_output_text - - print( - "✅ BLOCKED action HTTPException test passed - output text properly included" + # Call the method + result = await guardrail.make_bedrock_api_request( + source="INPUT", + messages=request_data.get("messages"), + request_data=request_data, ) + # Verify that the original response data was not modified + # (The json() method should return the original data) + original_data = mock_bedrock_response.json() + assert ( + original_data["assessments"][0]["sensitiveInformationPolicy"][ + "piiEntities" + ][0]["match"] + == "+1 412 555 1212" + ) + + # Verify that the returned BedrockGuardrailResponse contains original data + assert ( + result["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0][ + "match" + ] + == "+1 412 555 1212" + ) + + print("Original response not modified test passed") + @pytest.mark.asyncio -async def test_bedrock_guardrail_blocked_action_empty_outputs(): - """Test that BLOCKED actions with empty outputs still raise HTTPException""" - from unittest.mock import AsyncMock, MagicMock, patch - from litellm.proxy._types import UserAPIKeyAuth - from fastapi import HTTPException +async def test__redact_pii_matches_preserves_non_pii_entities(): + """Test that _redact_pii_matches only affects PII-related entities and preserves other assessment data""" - # Create proper mock objects - mock_user_api_key_dict = UserAPIKeyAuth() - - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) - - # Mock the Bedrock API response with BLOCKED action but no outputs - mock_bedrock_response = MagicMock() - mock_bedrock_response.status_code = 200 - mock_bedrock_response.json.return_value = { + response_with_mixed_data = { "action": "GUARDRAIL_INTERVENED", - "outputs": [], # Empty outputs "assessments": [ { - "contentPolicy": { - "filters": [ - {"type": "VIOLENCE", "confidence": "HIGH", "action": "BLOCKED"} - ] - } - } - ], - } - - request_data = { - "model": "gpt-4o", - "messages": [ - {"role": "user", "content": "Violent content here"}, - ], - } - - # Patch the async_handler.post method - with patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post: - mock_post.return_value = mock_bedrock_response - - # This should raise HTTPException due to BLOCKED action - with pytest.raises(HTTPException) as exc_info: - await guardrail.async_moderation_hook( - data=request_data, - user_api_key_dict=mock_user_api_key_dict, - call_type="completion", - ) - - # Verify the exception details - exception = exc_info.value - assert exception.status_code == 400 - - # Check that the detail contains the expected structure with empty output text - detail = exception.detail - assert isinstance(detail, dict) - assert detail["error"] == "Violated guardrail policy" - assert detail["bedrock_guardrail_response"] == "" # Empty string for no outputs - - print("✅ BLOCKED action with empty outputs test passed") - - -@pytest.mark.asyncio -async def test_bedrock_guardrail_disable_exception_on_block_non_streaming(): - """Test that disable_exception_on_block=True prevents exceptions in non-streaming scenarios""" - from unittest.mock import AsyncMock, MagicMock, patch - from litellm.proxy._types import UserAPIKeyAuth - from fastapi import HTTPException - - # Create proper mock objects - mock_user_api_key_dict = UserAPIKeyAuth() - - # Test 1: disable_exception_on_block=False (default) - should raise exception - guardrail_default = BedrockGuardrail( - guardrailIdentifier="test-guardrail", - guardrailVersion="DRAFT", - disable_exception_on_block=False, - ) - - # Mock the Bedrock API response with BLOCKED action - mock_bedrock_response = MagicMock() - mock_bedrock_response.status_code = 200 - mock_bedrock_response.json.return_value = { - "action": "GUARDRAIL_INTERVENED", - "outputs": [{"text": "I can't provide that information."}], - "assessments": [ - { - "topicPolicy": { - "topics": [ - {"name": "Sensitive Topic", "type": "DENY", "action": "BLOCKED"} - ] - } - } - ], - } - - request_data = { - "model": "gpt-4o", - "messages": [ - {"role": "user", "content": "Tell me how to make explosives"}, - ], - } - - # Patch the async_handler.post method - with patch.object( - guardrail_default.async_handler, "post", new_callable=AsyncMock - ) as mock_post: - mock_post.return_value = mock_bedrock_response - - # Should raise HTTPException when disable_exception_on_block=False - with pytest.raises(HTTPException) as exc_info: - await guardrail_default.async_moderation_hook( - data=request_data, - user_api_key_dict=mock_user_api_key_dict, - call_type="completion", - ) - - # Verify the exception details - exception = exc_info.value - assert exception.status_code == 400 - assert "Violated guardrail policy" in str(exception.detail) - - # Test 2: disable_exception_on_block=True - should NOT raise exception - guardrail_disabled = BedrockGuardrail( - guardrailIdentifier="test-guardrail", - guardrailVersion="DRAFT", - disable_exception_on_block=True, - ) - - with patch.object( - guardrail_disabled.async_handler, "post", new_callable=AsyncMock - ) as mock_post: - mock_post.return_value = mock_bedrock_response - - # Should NOT raise exception when disable_exception_on_block=True - try: - response = await guardrail_disabled.async_moderation_hook( - data=request_data, - user_api_key_dict=mock_user_api_key_dict, - call_type="completion", - ) - # Should succeed and return data (even though content was blocked) - assert response is not None - print("✅ No exception raised when disable_exception_on_block=True") - except Exception as e: - pytest.fail( - f"Should not raise exception when disable_exception_on_block=True, but got: {e}" - ) - - -@pytest.mark.asyncio -async def test_bedrock_guardrail_disable_exception_on_block_streaming(): - """Test that disable_exception_on_block=True prevents exceptions in streaming scenarios""" - from unittest.mock import AsyncMock, MagicMock, patch - from litellm.proxy._types import UserAPIKeyAuth - from litellm.types.utils import ModelResponseStream - from fastapi import HTTPException - import litellm - - # Create proper mock objects - mock_user_api_key_dict = UserAPIKeyAuth() - - # Mock streaming chunks that would normally trigger a block - async def mock_streaming_response(): - chunks = [ - ModelResponseStream( - id="test-id", - choices=[ - litellm.utils.StreamingChoices( - index=0, - delta=litellm.utils.Delta( - content="Here's how to make explosives: " - ), - finish_reason=None, - ) - ], - created=1234567890, - model="gpt-4o", - object="chat.completion.chunk", - ), - ModelResponseStream( - id="test-id", - choices=[ - litellm.utils.StreamingChoices( - index=0, - delta=litellm.utils.Delta(content="step 1, step 2..."), - finish_reason="stop", - ) - ], - created=1234567890, - model="gpt-4o", - object="chat.completion.chunk", - ), - ] - for chunk in chunks: - yield chunk - - # Mock Bedrock API response with BLOCKED action - mock_bedrock_response = MagicMock() - mock_bedrock_response.status_code = 200 - mock_bedrock_response.json.return_value = { - "action": "GUARDRAIL_INTERVENED", - "outputs": [{"text": "I can't provide that information."}], - "assessments": [ - { - "contentPolicy": { - "filters": [ - {"type": "VIOLENCE", "confidence": "HIGH", "action": "BLOCKED"} - ] - } - } - ], - } - - request_data = { - "model": "gpt-4o", - "messages": [{"role": "user", "content": "Tell me how to make explosives"}], - "stream": True, - } - - # Test 1: disable_exception_on_block=False (default) - should raise exception - guardrail_default = BedrockGuardrail( - guardrailIdentifier="test-guardrail", - guardrailVersion="DRAFT", - disable_exception_on_block=False, - ) - - with patch.object( - guardrail_default.async_handler, "post", new_callable=AsyncMock - ) as mock_post: - mock_post.return_value = mock_bedrock_response - - # Should raise exception during streaming processing - with pytest.raises(HTTPException): - result_generator = ( - guardrail_default.async_post_call_streaming_iterator_hook( - user_api_key_dict=mock_user_api_key_dict, - response=mock_streaming_response(), - request_data=request_data, - ) - ) - - # Try to consume the generator - should raise exception - async for chunk in result_generator: - pass - - # Test 2: disable_exception_on_block=True - should NOT raise exception - guardrail_disabled = BedrockGuardrail( - guardrailIdentifier="test-guardrail", - guardrailVersion="DRAFT", - disable_exception_on_block=True, - ) - - with patch.object( - guardrail_disabled.async_handler, "post", new_callable=AsyncMock - ) as mock_post: - mock_post.return_value = mock_bedrock_response - - # Should NOT raise exception when disable_exception_on_block=True - try: - result_generator = ( - guardrail_disabled.async_post_call_streaming_iterator_hook( - user_api_key_dict=mock_user_api_key_dict, - response=mock_streaming_response(), - request_data=request_data, - ) - ) - - # Consume the generator - should succeed without exceptions - result_chunks = [] - async for chunk in result_generator: - result_chunks.append(chunk) - - # Should have received chunks back even though content was blocked - assert len(result_chunks) > 0 - print( - "✅ Streaming completed without exception when disable_exception_on_block=True" - ) - - except Exception as e: - pytest.fail( - f"Should not raise exception when disable_exception_on_block=True in streaming, but got: {e}" - ) - - -@pytest.mark.asyncio -async def test_bedrock_guardrail_post_call_success_hook_no_output_text(): - """Test that async_post_call_success_hook skips when there's no output text""" - from unittest.mock import AsyncMock, MagicMock, patch - from litellm.proxy._types import UserAPIKeyAuth - from litellm.types.utils import ModelResponseStream - import litellm - - # Create proper mock objects - mock_user_api_key_dict = UserAPIKeyAuth() - - # Create guardrail instance - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) - - # Create a ModelResponse with tool calls (no text content) - # This simulates a response where the LLM is making a tool call - mock_response = litellm.ModelResponse( - id="test-id", - choices=[ - litellm.Choices( - index=0, - message=litellm.Message( - role="assistant", - content=None, # No text content - tool_calls=[ - litellm.utils.ChatCompletionMessageToolCall( - id="tooluse_kZJMlvQmRJ6eAyJE5GIl7Q", - function=litellm.utils.Function( - name="top_song", arguments='{"sign": "WZPZ"}' - ), - type="function", - ) + "sensitiveInformationPolicy": { + "piiEntities": [ + { + "type": "EMAIL", + "match": "user@example.com", + "action": "ANONYMIZED", + "confidence": "HIGH", + } ], - ), - finish_reason="tool_calls", - ) + "regexes": [ + { + "name": "custom_pattern", + "match": "some_pattern_match", + "action": "BLOCKED", + } + ], + }, + "contentPolicy": { + "filters": [ + { + "type": "VIOLENCE", + "confidence": "MEDIUM", + "action": "BLOCKED", + } + ] + }, + "topicPolicy": { + "topics": [ + { + "name": "Restricted Topic", + "type": "DENY", + "action": "BLOCKED", + } + ] + }, + } ], - created=1234567890, - model="gpt-4o", - object="chat.completion", + "outputs": [{"text": "Content blocked"}], + } + + # Call the redaction function + redacted_response = _redact_pii_matches(response_with_mixed_data) + + # Verify that PII entity matches are redacted + pii_entities = redacted_response["assessments"][0]["sensitiveInformationPolicy"][ + "piiEntities" + ] + assert pii_entities[0]["match"] == "[REDACTED]", "PII match should be redacted" + assert pii_entities[0]["type"] == "EMAIL", "PII type should be preserved" + assert pii_entities[0]["action"] == "ANONYMIZED", "PII action should be preserved" + assert pii_entities[0]["confidence"] == "HIGH", "PII confidence should be preserved" + + # Verify that regex matches are also redacted (updated behavior) + regexes = redacted_response["assessments"][0]["sensitiveInformationPolicy"][ + "regexes" + ] + assert regexes[0]["match"] == "[REDACTED]", "Regex match should be redacted" + assert regexes[0]["name"] == "custom_pattern", "Regex name should be preserved" + assert regexes[0]["action"] == "BLOCKED", "Regex action should be preserved" + + # Verify that other policies are completely unchanged + content_policy = redacted_response["assessments"][0]["contentPolicy"] + assert content_policy["filters"][0]["type"] == "VIOLENCE" + assert content_policy["filters"][0]["confidence"] == "MEDIUM" + + topic_policy = redacted_response["assessments"][0]["topicPolicy"] + assert topic_policy["topics"][0]["name"] == "Restricted Topic" + + # Verify top-level fields are unchanged + assert redacted_response["action"] == "GUARDRAIL_INTERVENED" + assert redacted_response["outputs"][0]["text"] == "Content blocked" + + print("Preserves non-PII entities test passed") + + +@pytest.mark.asyncio +async def test_pii_redaction_matches_debug_output_format(): + """Test that demonstrates the exact behavior shown in your debug output""" + + # This matches the structure from your debug output + original_response = { + "action": "GUARDRAIL_INTERVENED", + "actionReason": "Guardrail blocked.", + "assessments": [ + { + "invocationMetrics": { + "guardrailCoverage": { + "textCharacters": {"guarded": 84, "total": 84} + }, + "guardrailProcessingLatency": 322, + "usage": { + "contentPolicyImageUnits": 0, + "contentPolicyUnits": 0, + "contextualGroundingPolicyUnits": 0, + "sensitiveInformationPolicyFreeUnits": 0, + "sensitiveInformationPolicyUnits": 1, + "topicPolicyUnits": 0, + "wordPolicyUnits": 0, + }, + }, + "sensitiveInformationPolicy": { + "piiEntities": [ + { + "action": "BLOCKED", + "detected": True, + "match": "John Smith", + "type": "NAME", + }, + { + "action": "BLOCKED", + "detected": True, + "match": "324-12-3212", + "type": "US_SOCIAL_SECURITY_NUMBER", + }, + { + "action": "BLOCKED", + "detected": True, + "match": "607-456-7890", + "type": "PHONE", + }, + ] + }, + } + ], + "blockedResponse": "Input blocked by PII policy", + "guardrailCoverage": {"textCharacters": {"guarded": 84, "total": 84}}, + "output": [{"text": "Input blocked by PII policy"}], + "outputs": [{"text": "Input blocked by PII policy"}], + "usage": { + "contentPolicyImageUnits": 0, + "contentPolicyUnits": 0, + "contextualGroundingPolicyUnits": 0, + "sensitiveInformationPolicyFreeUnits": 0, + "sensitiveInformationPolicyUnits": 1, + "topicPolicyUnits": 0, + "wordPolicyUnits": 0, + }, + } + + # Apply redaction + redacted_response = _redact_pii_matches(original_response) + + # Verify the redacted response matches your expected debug output + pii_entities = redacted_response["assessments"][0]["sensitiveInformationPolicy"][ + "piiEntities" + ] + + # All PII matches should be redacted + assert pii_entities[0]["match"] == "[REDACTED]", "NAME should be redacted" + assert pii_entities[1]["match"] == "[REDACTED]", "SSN should be redacted" + assert pii_entities[2]["match"] == "[REDACTED]", "PHONE should be redacted" + + # But all other fields should be preserved + assert pii_entities[0]["type"] == "NAME" + assert pii_entities[1]["type"] == "US_SOCIAL_SECURITY_NUMBER" + assert pii_entities[2]["type"] == "PHONE" + assert pii_entities[0]["action"] == "BLOCKED" + assert pii_entities[0]["detected"] == True + + # Verify that the original response is unchanged + original_pii_entities = original_response["assessments"][0][ + "sensitiveInformationPolicy" + ]["piiEntities"] + assert ( + original_pii_entities[0]["match"] == "John Smith" + ), "Original should be unchanged" + assert ( + original_pii_entities[1]["match"] == "324-12-3212" + ), "Original should be unchanged" + assert ( + original_pii_entities[2]["match"] == "607-456-7890" + ), "Original should be unchanged" + + # Verify all other metadata is preserved in redacted response + assert redacted_response["action"] == "GUARDRAIL_INTERVENED" + assert redacted_response["actionReason"] == "Guardrail blocked." + assert redacted_response["blockedResponse"] == "Input blocked by PII policy" + assert ( + redacted_response["assessments"][0]["invocationMetrics"][ + "guardrailProcessingLatency" + ] + == 322 ) - data = { + print("PII redaction matches debug output format test passed") + print( + f"Original PII values preserved: {[e['match'] for e in original_pii_entities]}" + ) + print(f"Redacted PII values: {[e['match'] for e in pii_entities]}") + + +@pytest.mark.asyncio +async def test__redact_pii_matches_with_regex_matches(): + """Test redaction of regex matches in sensitive information policy""" + + response_with_regex = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [ + { + "sensitiveInformationPolicy": { + "regexes": [ + { + "name": "SSN_PATTERN", + "match": "123-45-6789", + "action": "BLOCKED", + }, + { + "name": "CREDIT_CARD_PATTERN", + "match": "4111-1111-1111-1111", + "action": "ANONYMIZED", + }, + ] + } + } + ], + "outputs": [{"text": "Regex patterns detected"}], + } + + # Call the redaction function + redacted_response = _redact_pii_matches(response_with_regex) + + # Verify that regex matches are redacted + regexes = redacted_response["assessments"][0]["sensitiveInformationPolicy"][ + "regexes" + ] + + assert regexes[0]["match"] == "[REDACTED]", "SSN regex match should be redacted" + assert ( + regexes[1]["match"] == "[REDACTED]" + ), "Credit card regex match should be redacted" + + # Verify other fields are preserved + assert regexes[0]["name"] == "SSN_PATTERN", "Regex name should be preserved" + assert regexes[0]["action"] == "BLOCKED", "Regex action should be preserved" + assert regexes[1]["name"] == "CREDIT_CARD_PATTERN", "Regex name should be preserved" + assert regexes[1]["action"] == "ANONYMIZED", "Regex action should be preserved" + + # Verify original response is unchanged + original_regexes = response_with_regex["assessments"][0][ + "sensitiveInformationPolicy" + ]["regexes"] + assert original_regexes[0]["match"] == "123-45-6789", "Original should be unchanged" + assert ( + original_regexes[1]["match"] == "4111-1111-1111-1111" + ), "Original should be unchanged" + + print("Regex matches redaction test passed") + + +@pytest.mark.asyncio +async def test__redact_pii_matches_with_custom_words(): + """Test redaction of custom word matches in word policy""" + + response_with_custom_words = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [ + { + "wordPolicy": { + "customWords": [ + { + "match": "confidential_data", + "action": "BLOCKED", + }, + { + "match": "secret_information", + "action": "ANONYMIZED", + }, + ] + } + } + ], + "outputs": [{"text": "Custom words detected"}], + } + + # Call the redaction function + redacted_response = _redact_pii_matches(response_with_custom_words) + + # Verify that custom word matches are redacted + custom_words = redacted_response["assessments"][0]["wordPolicy"]["customWords"] + + assert ( + custom_words[0]["match"] == "[REDACTED]" + ), "First custom word match should be redacted" + assert ( + custom_words[1]["match"] == "[REDACTED]" + ), "Second custom word match should be redacted" + + # Verify other fields are preserved + assert ( + custom_words[0]["action"] == "BLOCKED" + ), "Custom word action should be preserved" + assert ( + custom_words[1]["action"] == "ANONYMIZED" + ), "Custom word action should be preserved" + + # Verify original response is unchanged + original_custom_words = response_with_custom_words["assessments"][0]["wordPolicy"][ + "customWords" + ] + assert ( + original_custom_words[0]["match"] == "confidential_data" + ), "Original should be unchanged" + assert ( + original_custom_words[1]["match"] == "secret_information" + ), "Original should be unchanged" + + print("Custom words redaction test passed") + + +@pytest.mark.asyncio +async def test__redact_pii_matches_with_managed_words(): + """Test redaction of managed word matches in word policy""" + + response_with_managed_words = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [ + { + "wordPolicy": { + "managedWordLists": [ + { + "match": "inappropriate_word", + "action": "BLOCKED", + "type": "PROFANITY", + }, + { + "match": "offensive_term", + "action": "ANONYMIZED", + "type": "HATE_SPEECH", + }, + ] + } + } + ], + "outputs": [{"text": "Managed words detected"}], + } + + # Call the redaction function + redacted_response = _redact_pii_matches(response_with_managed_words) + + # Verify that managed word matches are redacted + managed_words = redacted_response["assessments"][0]["wordPolicy"][ + "managedWordLists" + ] + + assert ( + managed_words[0]["match"] == "[REDACTED]" + ), "First managed word match should be redacted" + assert ( + managed_words[1]["match"] == "[REDACTED]" + ), "Second managed word match should be redacted" + + # Verify other fields are preserved + assert ( + managed_words[0]["action"] == "BLOCKED" + ), "Managed word action should be preserved" + assert ( + managed_words[0]["type"] == "PROFANITY" + ), "Managed word type should be preserved" + assert ( + managed_words[1]["action"] == "ANONYMIZED" + ), "Managed word action should be preserved" + assert ( + managed_words[1]["type"] == "HATE_SPEECH" + ), "Managed word type should be preserved" + + # Verify original response is unchanged + original_managed_words = response_with_managed_words["assessments"][0][ + "wordPolicy" + ]["managedWordLists"] + assert ( + original_managed_words[0]["match"] == "inappropriate_word" + ), "Original should be unchanged" + assert ( + original_managed_words[1]["match"] == "offensive_term" + ), "Original should be unchanged" + + print("Managed words redaction test passed") + + +@pytest.mark.asyncio +async def test__redact_pii_matches_comprehensive_coverage(): + """Test redaction across all supported policy types in a single response""" + + comprehensive_response = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [ + { + "sensitiveInformationPolicy": { + "piiEntities": [ + { + "type": "EMAIL", + "match": "user@example.com", + "action": "ANONYMIZED", + } + ], + "regexes": [ + { + "name": "PHONE_PATTERN", + "match": "555-123-4567", + "action": "BLOCKED", + } + ], + }, + "wordPolicy": { + "customWords": [ + { + "match": "confidential", + "action": "BLOCKED", + } + ], + "managedWordLists": [ + { + "match": "inappropriate", + "action": "ANONYMIZED", + "type": "PROFANITY", + } + ], + }, + } + ], + "outputs": [{"text": "Multiple policy violations detected"}], + } + + # Call the redaction function + redacted_response = _redact_pii_matches(comprehensive_response) + + # Verify all match fields are redacted + assessment = redacted_response["assessments"][0] + + # PII entities + pii_entities = assessment["sensitiveInformationPolicy"]["piiEntities"] + assert ( + pii_entities[0]["match"] == "[REDACTED]" + ), "PII entity match should be redacted" + + # Regex matches + regexes = assessment["sensitiveInformationPolicy"]["regexes"] + assert regexes[0]["match"] == "[REDACTED]", "Regex match should be redacted" + + # Custom words + custom_words = assessment["wordPolicy"]["customWords"] + assert ( + custom_words[0]["match"] == "[REDACTED]" + ), "Custom word match should be redacted" + + # Managed words + managed_words = assessment["wordPolicy"]["managedWordLists"] + assert ( + managed_words[0]["match"] == "[REDACTED]" + ), "Managed word match should be redacted" + + # Verify all other fields are preserved + assert pii_entities[0]["type"] == "EMAIL" + assert regexes[0]["name"] == "PHONE_PATTERN" + assert managed_words[0]["type"] == "PROFANITY" + + # Verify original response is unchanged + original_assessment = comprehensive_response["assessments"][0] + assert ( + original_assessment["sensitiveInformationPolicy"]["piiEntities"][0]["match"] + == "user@example.com" + ) + assert ( + original_assessment["sensitiveInformationPolicy"]["regexes"][0]["match"] + == "555-123-4567" + ) + assert ( + original_assessment["wordPolicy"]["customWords"][0]["match"] == "confidential" + ) + assert ( + original_assessment["wordPolicy"]["managedWordLists"][0]["match"] + == "inappropriate" + ) + + print("Comprehensive coverage redaction test passed") + +@pytest.mark.asyncio +async def test_bedrock_guardrail_respects_custom_runtime_endpoint(monkeypatch): + """Test that BedrockGuardrail respects aws_bedrock_runtime_endpoint when set""" + + # Clear any existing environment variable to ensure clean test + monkeypatch.delenv("AWS_BEDROCK_RUNTIME_ENDPOINT", raising=False) + + # Create guardrail with custom runtime endpoint + custom_endpoint = "https://custom-bedrock.example.com" + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT", + aws_bedrock_runtime_endpoint=custom_endpoint, + ) + + # Mock credentials + mock_credentials = MagicMock() + mock_credentials.access_key = "test-access-key" + mock_credentials.secret_key = "test-secret-key" + mock_credentials.token = None + + # Test data + data = {"source": "INPUT", "content": [{"text": {"text": "test content"}}]} + optional_params = {} + aws_region_name = "us-east-1" + + # Mock the _load_credentials method to avoid actual AWS credential loading + with patch.object( + guardrail, "_load_credentials", return_value=(mock_credentials, aws_region_name) + ): + # Call _prepare_request which internally calls get_runtime_endpoint + prepped_request = guardrail._prepare_request( + credentials=mock_credentials, + data=data, + optional_params=optional_params, + aws_region_name=aws_region_name, + ) + + # Verify that the custom endpoint is used in the URL + expected_url = f"{custom_endpoint}/guardrail/{guardrail.guardrailIdentifier}/version/{guardrail.guardrailVersion}/apply" + assert ( + prepped_request.url == expected_url + ), f"Expected URL to contain custom endpoint. Got: {prepped_request.url}" + + print(f"Custom runtime endpoint test passed. URL: {prepped_request.url}") + + +@pytest.mark.asyncio +async def test_bedrock_guardrail_respects_env_runtime_endpoint(monkeypatch): + """Test that BedrockGuardrail respects AWS_BEDROCK_RUNTIME_ENDPOINT environment variable""" + + custom_endpoint = "https://env-bedrock.example.com" + + # Set the environment variable + monkeypatch.setenv("AWS_BEDROCK_RUNTIME_ENDPOINT", custom_endpoint) + + # Create guardrail without explicit aws_bedrock_runtime_endpoint + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" + ) + + # Mock credentials + mock_credentials = MagicMock() + mock_credentials.access_key = "test-access-key" + mock_credentials.secret_key = "test-secret-key" + mock_credentials.token = None + + # Test data + data = {"source": "INPUT", "content": [{"text": {"text": "test content"}}]} + optional_params = {} + aws_region_name = "us-east-1" + + # Mock the _load_credentials method + with patch.object( + guardrail, "_load_credentials", return_value=(mock_credentials, aws_region_name) + ): + # Call _prepare_request which internally calls get_runtime_endpoint + prepped_request = guardrail._prepare_request( + credentials=mock_credentials, + data=data, + optional_params=optional_params, + aws_region_name=aws_region_name, + ) + + # Verify that the custom endpoint from environment is used in the URL + expected_url = f"{custom_endpoint}/guardrail/{guardrail.guardrailIdentifier}/version/{guardrail.guardrailVersion}/apply" + assert ( + prepped_request.url == expected_url + ), f"Expected URL to contain env endpoint. Got: {prepped_request.url}" + + print(f"Environment runtime endpoint test passed. URL: {prepped_request.url}") + + +@pytest.mark.asyncio +async def test_bedrock_guardrail_uses_default_endpoint_when_no_custom_set(monkeypatch): + """Test that BedrockGuardrail uses default endpoint when no custom endpoint is set""" + + # Ensure no environment variable is set + monkeypatch.delenv("AWS_BEDROCK_RUNTIME_ENDPOINT", raising=False) + + # Create guardrail without any custom endpoint + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" + ) + + # Mock credentials + mock_credentials = MagicMock() + mock_credentials.access_key = "test-access-key" + mock_credentials.secret_key = "test-secret-key" + mock_credentials.token = None + + # Test data + data = {"source": "INPUT", "content": [{"text": {"text": "test content"}}]} + optional_params = {} + aws_region_name = "us-west-2" + + # Mock the _load_credentials method + with patch.object( + guardrail, "_load_credentials", return_value=(mock_credentials, aws_region_name) + ): + # Call _prepare_request which internally calls get_runtime_endpoint + prepped_request = guardrail._prepare_request( + credentials=mock_credentials, + data=data, + optional_params=optional_params, + aws_region_name=aws_region_name, + ) + + # Verify that the default endpoint is used + expected_url = f"https://bedrock-runtime.{aws_region_name}.amazonaws.com/guardrail/{guardrail.guardrailIdentifier}/version/{guardrail.guardrailVersion}/apply" + assert ( + prepped_request.url == expected_url + ), f"Expected default URL. Got: {prepped_request.url}" + + print(f"Default endpoint test passed. URL: {prepped_request.url}") + + +@pytest.mark.asyncio +async def test_bedrock_guardrail_parameter_takes_precedence_over_env(monkeypatch): + """Test that aws_bedrock_runtime_endpoint parameter takes precedence over environment variable + + This test verifies the corrected behavior where the parameter should take precedence + over the environment variable, consistent with the endpoint_url logic. + """ + + param_endpoint = "https://param-bedrock.example.com" + env_endpoint = "https://env-bedrock.example.com" + + # Set environment variable + monkeypatch.setenv("AWS_BEDROCK_RUNTIME_ENDPOINT", env_endpoint) + + # Create guardrail with explicit aws_bedrock_runtime_endpoint + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT", + aws_bedrock_runtime_endpoint=param_endpoint, + ) + + # Mock credentials + mock_credentials = MagicMock() + mock_credentials.access_key = "test-access-key" + mock_credentials.secret_key = "test-secret-key" + mock_credentials.token = None + + # Test data + data = {"source": "INPUT", "content": [{"text": {"text": "test content"}}]} + optional_params = {} + aws_region_name = "us-east-1" + + # Mock the _load_credentials method + with patch.object( + guardrail, "_load_credentials", return_value=(mock_credentials, aws_region_name) + ): + # Call _prepare_request which internally calls get_runtime_endpoint + prepped_request = guardrail._prepare_request( + credentials=mock_credentials, + data=data, + optional_params=optional_params, + aws_region_name=aws_region_name, + ) + + # Verify that the parameter takes precedence over environment variable + expected_url = f"{param_endpoint}/guardrail/{guardrail.guardrailIdentifier}/version/{guardrail.guardrailVersion}/apply" + assert ( + prepped_request.url == expected_url + ), f"Expected parameter endpoint to take precedence. Got: {prepped_request.url}" + + print(f"Parameter precedence test passed. URL: {prepped_request.url}") + + +@pytest.mark.asyncio +async def test_bedrock_apply_guardrail_with_only_tool_calls_response(): + """Test that apply_guardrail handles response with tool_calls (no text content) without calling Bedrock API""" + # Create a BedrockGuardrail instance + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" + ) + + # Mock the make_bedrock_api_request method + with patch.object( + guardrail, "make_bedrock_api_request", new_callable=AsyncMock + ) as mock_api_request: + # Test the apply_guardrail method with tool_calls in response + inputs = { + "texts": [], + "tool_calls": [ + { + "id": "call_eFSCWFsyL7MclHYnzKrcQnMK", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"location":"São Paulo"}', + }, + } + ], + } + + guardrailed_inputs = await guardrail.apply_guardrail( + inputs=inputs, + request_data={}, + input_type="response", + logging_obj=None, + ) + + # Verify the result - should succeed without errors + assert guardrailed_inputs is not None + assert "tool_calls" in guardrailed_inputs + assert len(guardrailed_inputs["tool_calls"]) == 1 + assert ( + guardrailed_inputs["tool_calls"][0]["id"] + == "call_eFSCWFsyL7MclHYnzKrcQnMK" + ) + assert guardrailed_inputs["tool_calls"][0]["function"]["name"] == "get_weather" + assert ( + guardrailed_inputs["tool_calls"][0]["function"]["arguments"] + == '{"location":"São Paulo"}' + ) + # Verify that the Bedrock API was NOT called since there's no text to process + mock_api_request.assert_not_called() + print("✅ apply_guardrail with tool_calls test passed - no API call made") + + +@pytest.mark.asyncio +async def test_bedrock_guardrail_blocked_content_with_masking_enabled(): + """Test that BLOCKED content raises exception even when masking is enabled + + This test verifies the bug fix where previously mask_request_content=True or + mask_response_content=True would bypass all BLOCKED content checks. Now it + properly distinguishes between BLOCKED (raise exception) and ANONYMIZED (apply masking). + """ + + # Create guardrail with masking enabled + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT", + mask_request_content=True, # Masking enabled + mask_response_content=True, # Masking enabled + ) + + # Mock Bedrock response with BLOCKED content (hate speech) + blocked_response = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [ + { + "contentPolicy": { + "filters": [ + { + "type": "HATE", + "confidence": "HIGH", + "action": "BLOCKED", # Should raise exception + } + ] + }, + "sensitiveInformationPolicy": { + "piiEntities": [ + { + "type": "NAME", + "match": "John Doe", + "action": "ANONYMIZED", # Should be masked + } + ] + }, + } + ], + "outputs": [{"text": "Content blocked due to policy violation"}], + } + + mock_bedrock_response = MagicMock() + mock_bedrock_response.status_code = 200 + mock_bedrock_response.json.return_value = blocked_response + + # Mock credentials + mock_credentials = MagicMock() + mock_credentials.access_key = "test-access-key" + mock_credentials.secret_key = "test-secret-key" + mock_credentials.token = None + + request_data = { "model": "gpt-4o", "messages": [ - {"role": "user", "content": "Hello"}, + {"role": "user", "content": "Test message with PII and hate speech"}, ], } - mock_user_api_key_dict = UserAPIKeyAuth() + + # Mock AWS-related methods + with patch.object( + guardrail.async_handler, "post", new_callable=AsyncMock + ) as mock_post, patch.object( + guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") + ), patch.object( + guardrail, "_prepare_request", return_value=MagicMock() + ): + mock_post.return_value = mock_bedrock_response + + # Should raise HTTPException for BLOCKED content + with pytest.raises(HTTPException) as exc_info: + await guardrail.make_bedrock_api_request( + source="INPUT", + messages=request_data.get("messages"), + request_data=request_data, + ) + + # Verify exception details + assert exc_info.value.status_code == 400 + assert "Violated guardrail policy" in str(exc_info.value.detail) + + print("✅ BLOCKED content with masking enabled raises exception correctly") - result = await guardrail.async_post_call_success_hook( - data=data, - response=mock_response, - user_api_key_dict=mock_user_api_key_dict, - ) - # If no error is raised and result is None, then the test passes - assert result is None - print("✅ No output text in response test passed") + +# ────────────────────────────────────────────────────────────────────────────── +# Null-safety tests for Bedrock guardrail responses +# +# The Bedrock ApplyGuardrail API can return explicit null/None for list fields +# such as "regexes", "piiEntities", "topics", "filters", "customWords", and +# "managedWordLists" when a particular policy category is present in the +# assessment but has no matches. +# +# Python's dict.get("key", []) returns None (NOT []) when the key exists with +# a None value. The `or []` fallback ensures we always iterate over a list. +# +# Without the fix, iterating over None raises: +# TypeError: 'NoneType' object is not iterable +# which surfaces to callers as: +# openai.InternalServerError: Error code: 500 +# {'error': {'message': "Bedrock guardrail failed: 'NoneType' object is not iterable", ...}} +# ────────────────────────────────────────────────────────────────────────────── + + +class TestRedactPiiMatchesNullSafety: + """Tests for _redact_pii_matches handling of null/None list fields from Bedrock API.""" + + @pytest.mark.asyncio + async def test_should_handle_null_regexes_in_sensitive_info_policy(self): + """Bedrock can return regexes: null while piiEntities has data. + + Real-world scenario: guardrail detects PII (e.g. EMAIL) but has no + custom regex patterns configured, so the API returns regexes: null. + """ + response = { + "action": "NONE", + "actionReason": "No action.", + "assessments": [ + { + "sensitiveInformationPolicy": { + "piiEntities": [ + { + "action": "NONE", + "detected": True, + "match": "joebloggs@gmail.com", + "type": "EMAIL", + } + ], + "regexes": None, # Explicit null from Bedrock API + }, + } + ], + } + + # Should not raise TypeError: 'NoneType' object is not iterable + redacted = _redact_pii_matches(response) + + # PII match should be redacted + pii = redacted["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"] + assert pii[0]["match"] == "[REDACTED]" + assert pii[0]["type"] == "EMAIL" + + @pytest.mark.asyncio + async def test_should_handle_null_pii_entities_in_sensitive_info_policy(self): + """Bedrock can return piiEntities: null while regexes has data.""" + response = { + "action": "NONE", + "assessments": [ + { + "sensitiveInformationPolicy": { + "piiEntities": None, # null from Bedrock API + "regexes": [ + { + "name": "CUSTOM_PATTERN", + "match": "secret-abc-123", + "action": "BLOCKED", + } + ], + }, + } + ], + } + + redacted = _redact_pii_matches(response) + + regexes = redacted["assessments"][0]["sensitiveInformationPolicy"]["regexes"] + assert regexes[0]["match"] == "[REDACTED]" + + @pytest.mark.asyncio + async def test_should_handle_null_custom_words_and_managed_words(self): + """Bedrock can return null for customWords and managedWordLists in wordPolicy.""" + response = { + "action": "NONE", + "assessments": [ + { + "wordPolicy": { + "customWords": None, # null from Bedrock API + "managedWordLists": None, # null from Bedrock API + }, + } + ], + } + + # Should not raise TypeError + redacted = _redact_pii_matches(response) + + # Values should remain None (no crash) + assert redacted["assessments"][0]["wordPolicy"]["customWords"] is None + assert redacted["assessments"][0]["wordPolicy"]["managedWordLists"] is None + + @pytest.mark.asyncio + async def test_should_handle_null_assessments_list(self): + """Bedrock can return assessments: null.""" + response = { + "action": "NONE", + "assessments": None, # null from Bedrock API + } + + # Should not raise TypeError + redacted = _redact_pii_matches(response) + assert redacted["assessments"] is None + + @pytest.mark.asyncio + async def test_should_handle_all_null_policy_sub_lists_together(self): + """All sub-list fields are null at the same time — worst-case scenario.""" + response = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [ + { + "sensitiveInformationPolicy": { + "piiEntities": None, + "regexes": None, + }, + "wordPolicy": { + "customWords": None, + "managedWordLists": None, + }, + "topicPolicy": None, + "contentPolicy": None, + "contextualGroundingPolicy": None, + } + ], + } + + # Should not raise any exception + redacted = _redact_pii_matches(response) + assert redacted is not None + + +class TestShouldRaiseGuardrailBlockedExceptionNullSafety: + """Tests for _should_raise_guardrail_blocked_exception handling of null list fields.""" + + def _create_guardrail(self) -> BedrockGuardrail: + return BedrockGuardrail( + guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" + ) + + @pytest.mark.asyncio + async def test_should_handle_all_null_policy_sub_lists(self): + """All policy sub-lists are null — should not crash, should return False.""" + guardrail = self._create_guardrail() + + response = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [ + { + "topicPolicy": { + "topics": None, # null from Bedrock API + }, + "contentPolicy": { + "filters": None, # null + }, + "wordPolicy": { + "customWords": None, # null + "managedWordLists": None, # null + }, + "sensitiveInformationPolicy": { + "piiEntities": None, # null + "regexes": None, # null + }, + "contextualGroundingPolicy": { + "filters": None, # null + }, + } + ], + } + + # No BLOCKED actions found (all lists null) → should return False + result = guardrail._should_raise_guardrail_blocked_exception(response) + assert result is False + + @pytest.mark.asyncio + async def test_should_detect_blocked_despite_other_null_lists(self): + """A mix of null lists and a real BLOCKED action — should still detect it.""" + guardrail = self._create_guardrail() + + response = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [ + { + "topicPolicy": { + "topics": None, # null — should not crash + }, + "contentPolicy": { + "filters": [ + { + "type": "HATE", + "confidence": "HIGH", + "action": "BLOCKED", + } + ], + }, + "wordPolicy": { + "customWords": None, # null + "managedWordLists": None, # null + }, + "sensitiveInformationPolicy": { + "piiEntities": None, # null + "regexes": None, # null + }, + "contextualGroundingPolicy": None, # entire policy is null + } + ], + } + + # Should return True because contentPolicy has a BLOCKED filter + result = guardrail._should_raise_guardrail_blocked_exception(response) + assert result is True + + @pytest.mark.asyncio + async def test_should_handle_null_assessments_list(self): + """assessments itself is null — should return False.""" + guardrail = self._create_guardrail() + + response = { + "action": "GUARDRAIL_INTERVENED", + "assessments": None, # null from Bedrock API + } + + result = guardrail._should_raise_guardrail_blocked_exception(response) + assert result is False + + @pytest.mark.asyncio + async def test_should_handle_null_topics_with_blocked_word_policy(self): + """topics is null but wordPolicy has a BLOCKED customWord.""" + guardrail = self._create_guardrail() + + response = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [ + { + "topicPolicy": { + "topics": None, + }, + "wordPolicy": { + "customWords": [ + {"match": "badword", "action": "BLOCKED"} + ], + "managedWordLists": None, + }, + } + ], + } + + result = guardrail._should_raise_guardrail_blocked_exception(response) + assert result is True + + @pytest.mark.asyncio + async def test_should_handle_null_pii_with_blocked_regex(self): + """piiEntities is null but regexes has a BLOCKED match.""" + guardrail = self._create_guardrail() + + response = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [ + { + "sensitiveInformationPolicy": { + "piiEntities": None, + "regexes": [ + {"name": "SSN", "match": "123-45-6789", "action": "BLOCKED"} + ], + }, + } + ], + } + + result = guardrail._should_raise_guardrail_blocked_exception(response) + assert result is True + + @pytest.mark.asyncio + async def test_should_handle_null_grounding_filters(self): + """contextualGroundingPolicy.filters is null — should not crash.""" + guardrail = self._create_guardrail() + + response = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [ + { + "contextualGroundingPolicy": { + "filters": None, + }, + } + ], + } + + result = guardrail._should_raise_guardrail_blocked_exception(response) + assert result is False + + @pytest.mark.asyncio + async def test_should_not_crash_when_action_is_not_intervened(self): + """If action != GUARDRAIL_INTERVENED, null lists should never be reached.""" + guardrail = self._create_guardrail() + + response = { + "action": "NONE", + "assessments": [ + { + "sensitiveInformationPolicy": { + "piiEntities": None, + "regexes": None, + }, + } + ], + } + + result = guardrail._should_raise_guardrail_blocked_exception(response) + assert result is False + + +class TestApplyGuardrailNullSafety: + """Tests for apply_guardrail handling of null/None texts input.""" + + @pytest.mark.asyncio + async def test_should_handle_none_texts_in_inputs(self): + """inputs[\"texts\"] is explicitly None — should not crash.""" + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" + ) + + inputs = {"texts": None} # Explicit None + + mock_credentials = MagicMock() + + with patch.object( + guardrail.async_handler, "post", new_callable=AsyncMock + ) as mock_post, patch.object( + guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") + ), patch.object( + guardrail, "_prepare_request", return_value=MagicMock() + ): + # With empty texts (from None → []), no Bedrock API call should be made + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={}, + input_type="request", + ) + + # Should return empty texts without crashing + assert result.get("texts") == [] + # No Bedrock API call should be made for empty input + mock_post.assert_not_called() + + @pytest.mark.asyncio + async def test_should_handle_missing_texts_key(self): + """inputs has no \"texts\" key at all — should not crash.""" + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" + ) + + inputs = {} # No "texts" key + + mock_credentials = MagicMock() + + with patch.object( + guardrail.async_handler, "post", new_callable=AsyncMock + ) as mock_post, patch.object( + guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") + ), patch.object( + guardrail, "_prepare_request", return_value=MagicMock() + ): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={}, + input_type="request", + ) + + assert result.get("texts") == [] + mock_post.assert_not_called() From cf94f4d8b720e63beadd10a28cea8c1787830814 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?d=20=F0=9F=94=B9?= Date: Sun, 5 Apr 2026 09:23:32 +0800 Subject: [PATCH 003/420] fix(mcp): is_tool_name_prefixed validates against known server prefixes (#25085) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fixes #25081. is_tool_name_prefixed() checked for the presence of MCP_TOOL_PREFIX_SEPARATOR (default '-') anywhere in the tool name. Any non-MCP tool whose name contains a hyphen (e.g. 'text-to-speech', 'code-review') was silently misclassified as an MCP-prefixed tool. When the semantic tool filter is enabled, these tools would be routed through semantic matching and potentially dropped. Fix: accept an optional known_server_prefixes set. When supplied, the function extracts the candidate prefix (text before the first separator) and checks it against the normalised set of registered server prefixes. Only a genuine match returns True. Without the set, legacy behaviour is preserved for backward compatibility. Updated _get_mcp_server_from_tool_name() to build the prefix set from the live registry and pass it through. 9 new tests. Co-authored-by: d 🔹 <258577966+voidborne-d@users.noreply.github.com> --- .../mcp_server/mcp_server_manager.py | 7 +- .../proxy/_experimental/mcp_server/utils.py | 32 +++++-- .../mcp_server/test_is_tool_name_prefixed.py | 90 +++++++++++++++++++ 3 files changed, 123 insertions(+), 6 deletions(-) create mode 100644 tests/test_litellm/proxy/_experimental/mcp_server/test_is_tool_name_prefixed.py diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 7b87e7e7e61..5363e317ff7 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -2552,7 +2552,12 @@ class MCPServerManager: return server # If not found and tool name is prefixed, try extracting server name from prefix - if is_tool_name_prefixed(tool_name): + known_prefixes = { + normalize_server_name(get_server_prefix(s)) + for s in self.get_registry().values() + if get_server_prefix(s) + } + if is_tool_name_prefixed(tool_name, known_server_prefixes=known_prefixes): ( original_tool_name, server_name_from_prefix, diff --git a/litellm/proxy/_experimental/mcp_server/utils.py b/litellm/proxy/_experimental/mcp_server/utils.py index 8189f212bcb..79942eda54e 100644 --- a/litellm/proxy/_experimental/mcp_server/utils.py +++ b/litellm/proxy/_experimental/mcp_server/utils.py @@ -100,17 +100,39 @@ def split_server_prefix_from_name(prefixed_name: str) -> Tuple[str, str]: return prefixed_name, "" -def is_tool_name_prefixed(tool_name: str) -> bool: +def is_tool_name_prefixed( + tool_name: str, + known_server_prefixes: Optional[set] = None, +) -> bool: """ - Check if tool name has server prefix + Check if tool name has a known MCP server prefix. + + When ``known_server_prefixes`` is provided the function verifies that the + substring before the first separator is an actual registered server + prefix. Without it the check falls back to the legacy heuristic + (separator present anywhere in the name), which can produce false + positives for non-MCP tools whose names contain hyphens + (e.g. ``text-to-speech``, ``code-review``). Args: - tool_name: Tool name to check + tool_name: Tool name to check. + known_server_prefixes: Optional set of normalised server prefixes + currently registered in the MCP manager. Pass this whenever + the caller has access to the server registry so that the check + is accurate. Returns: - True if tool name is prefixed, False otherwise + True if tool name is prefixed, False otherwise. """ - return MCP_TOOL_PREFIX_SEPARATOR in tool_name + if MCP_TOOL_PREFIX_SEPARATOR not in tool_name: + return False + + if known_server_prefixes is not None: + candidate_prefix = tool_name.split(MCP_TOOL_PREFIX_SEPARATOR, 1)[0] + return normalize_server_name(candidate_prefix) in known_server_prefixes + + # Legacy fallback – separator present somewhere in the name. + return True def validate_mcp_server_name( diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_is_tool_name_prefixed.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_is_tool_name_prefixed.py new file mode 100644 index 00000000000..d761d9c54cc --- /dev/null +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_is_tool_name_prefixed.py @@ -0,0 +1,90 @@ +""" +Tests for is_tool_name_prefixed with known_server_prefixes parameter. + +Verifies fix for https://github.com/BerriAI/litellm/issues/25081 +""" + +import pytest + +from litellm.proxy._experimental.mcp_server.utils import is_tool_name_prefixed + + +# --------------------------------------------------------------------------- +# Legacy behaviour (no known_server_prefixes passed) +# --------------------------------------------------------------------------- + + +class TestLegacyBehaviour: + """Without known_server_prefixes the function falls back to heuristic.""" + + def test_plain_name_returns_false(self): + assert is_tool_name_prefixed("get_weather") is False + + def test_hyphenated_name_returns_true_legacy(self): + """Legacy heuristic: any hyphen → True (the bug this issue reports).""" + assert is_tool_name_prefixed("text-to-speech") is True + + def test_prefixed_name_returns_true_legacy(self): + assert is_tool_name_prefixed("myserver-get_weather") is True + + +# --------------------------------------------------------------------------- +# New behaviour (known_server_prefixes supplied) +# --------------------------------------------------------------------------- + + +class TestWithKnownPrefixes: + """When known_server_prefixes is supplied, only real prefixes match.""" + + PREFIXES = {"myserver", "weather_api", "code_tools"} + + def test_known_prefix_returns_true(self): + assert ( + is_tool_name_prefixed( + "myserver-get_weather", known_server_prefixes=self.PREFIXES + ) + is True + ) + + def test_hyphenated_non_mcp_tool_returns_false(self): + """This is the core fix: 'text-to-speech' is NOT an MCP-prefixed tool.""" + assert ( + is_tool_name_prefixed( + "text-to-speech", known_server_prefixes=self.PREFIXES + ) + is False + ) + + def test_code_review_not_misclassified(self): + assert ( + is_tool_name_prefixed( + "code-review", known_server_prefixes=self.PREFIXES + ) + is False + ) + + def test_no_separator_returns_false(self): + assert ( + is_tool_name_prefixed( + "simple_tool", known_server_prefixes=self.PREFIXES + ) + is False + ) + + def test_empty_prefixes_set_rejects_all(self): + """With an empty registry, nothing can be prefixed.""" + assert ( + is_tool_name_prefixed("myserver-get_weather", known_server_prefixes=set()) + is False + ) + + def test_prefix_normalisation(self): + """Server names with spaces are normalised to underscores.""" + prefixes = {"my_server"} + # add_server_prefix_to_name normalises spaces → underscores + assert ( + is_tool_name_prefixed( + "my_server-list_files", known_server_prefixes=prefixes + ) + is True + ) From d6351a3966e5cbbddee1bacd63020bb6aa857614 Mon Sep 17 00:00:00 2001 From: Neha Prasad Date: Sun, 5 Apr 2026 07:09:37 +0530 Subject: [PATCH 004/420] fix(s3_v2): use prepared URL for SigV4-signed S3 requests (#25074) --- litellm/integrations/s3_v2.py | 13 +++--- tests/test_litellm/integrations/test_s3_v2.py | 44 +++++++++++++++++++ 2 files changed, 51 insertions(+), 6 deletions(-) diff --git a/litellm/integrations/s3_v2.py b/litellm/integrations/s3_v2.py index 405bf9698cc..f767f61be87 100644 --- a/litellm/integrations/s3_v2.py +++ b/litellm/integrations/s3_v2.py @@ -403,9 +403,8 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): # Prepare the signed headers signed_headers = dict(aws_request.headers.items()) - # Make the request response = await self.async_httpx_client.put( - url, data=json_string, headers=signed_headers + prepped.url, data=json_string, headers=signed_headers ) response.raise_for_status() except Exception as e: @@ -582,8 +581,9 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): if self.s3_verify is not None else None ) - # Make the request - response = httpx_client.put(url, data=json_string, headers=signed_headers) + response = httpx_client.put( + prepped.url, data=json_string, headers=signed_headers + ) response.raise_for_status() except Exception as e: verbose_logger.exception(f"Error uploading to s3: {str(e)}") @@ -674,8 +674,9 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): # Prepare the signed headers signed_headers = dict(aws_request.headers.items()) - # Make the request - response = await self.async_httpx_client.get(url, headers=signed_headers) + response = await self.async_httpx_client.get( + prepped.url, headers=signed_headers + ) if response.status_code != 200: verbose_logger.exception( diff --git a/tests/test_litellm/integrations/test_s3_v2.py b/tests/test_litellm/integrations/test_s3_v2.py index b53c05fa241..943fe4ec37b 100644 --- a/tests/test_litellm/integrations/test_s3_v2.py +++ b/tests/test_litellm/integrations/test_s3_v2.py @@ -292,6 +292,50 @@ class TestS3V2UnitTests: assert result == {"downloaded": "data"} + @patch("asyncio.create_task") + @patch("litellm.integrations.s3_v2.CustomBatchLogger.periodic_flush") + def test_s3_v2_put_url_encodes_spaces_in_object_key( + self, mock_periodic_flush, mock_create_task + ): + import requests + from unittest.mock import AsyncMock + + from litellm.types.integrations.s3_v2 import s3BatchLoggingElement + + mock_periodic_flush.return_value = None + mock_create_task.return_value = None + + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.raise_for_status = MagicMock() + + s3_object_key = "My Team/2025-09-14/test-key.json" + test_element = s3BatchLoggingElement( + s3_object_key=s3_object_key, + payload={"test": "data"}, + s3_object_download_filename="test-file.json", + ) + + s3_logger = S3Logger( + s3_bucket_name="test-bucket", + s3_endpoint_url="https://s3.amazonaws.com", + s3_aws_access_key_id="test-key", + s3_aws_secret_access_key="test-secret", + s3_region_name="us-east-1", + ) + s3_logger.async_httpx_client = AsyncMock() + s3_logger.async_httpx_client.put.return_value = mock_response + + asyncio.run(s3_logger.async_upload_data_to_s3(test_element)) + + call_args = s3_logger.async_httpx_client.put.call_args + assert call_args is not None + actual_url = call_args[0][0] + raw_url = f"https://s3.amazonaws.com/test-bucket/{s3_object_key}" + expected_url = requests.Request("PUT", raw_url).prepare().url + assert actual_url == expected_url + assert " " not in actual_url + @pytest.mark.asyncio async def test_async_log_event_skips_when_standard_logging_object_missing(): """ From e68cfaae0c1238d19a4944efb8af47c41dc949ce Mon Sep 17 00:00:00 2001 From: Christian Reynoso Hunter Date: Sat, 4 Apr 2026 22:40:56 -0300 Subject: [PATCH 005/420] fix(cache): Prevent "multiple values" error in get_cache_key (#20261) ## Problem When `get_cache_key(**kwargs)` is called with kwargs that already contains `preset_cache_key` (which can happen when cache key is recomputed in certain code paths), the call to `_set_preset_cache_key_in_kwargs()` fails with: ``` TypeError: _set_preset_cache_key_in_kwargs() got multiple values for keyword argument 'preset_cache_key' ``` This is because `preset_cache_key` is passed both explicitly: ```python self._set_preset_cache_key_in_kwargs( preset_cache_key=hashed_cache_key, **kwargs ) ``` And implicitly via `**kwargs` unpacking when `kwargs["preset_cache_key"]` exists. ## Solution Filter out `preset_cache_key` from kwargs before passing to `_set_preset_cache_key_in_kwargs()`: ```python kwargs_for_preset = {k: v for k, v in kwargs.items() if k != "preset_cache_key"} self._set_preset_cache_key_in_kwargs( preset_cache_key=hashed_cache_key, **kwargs_for_preset ) ``` ## Testing Added unit tests covering: - kwargs with existing preset_cache_key (the bug case) - kwargs without preset_cache_key (regression test) - Verification that preset_cache_key is correctly set in litellm_params --- litellm/caching/caching.py | 5 +- tests/local_testing/test_cache_preset_key.py | 87 ++++++++++++++++++++ 2 files changed, 91 insertions(+), 1 deletion(-) create mode 100644 tests/local_testing/test_cache_preset_key.py diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index 406a4f8c98a..6a68ba8c4d1 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -312,8 +312,11 @@ class Cache: verbose_logger.debug("\nCreated cache key: %s", cache_key) hashed_cache_key = Cache._get_hashed_cache_key(cache_key) hashed_cache_key = self._add_namespace_to_cache_key(hashed_cache_key, **kwargs) + # Remove preset_cache_key from kwargs to avoid "got multiple values" TypeError + # when kwargs already contains preset_cache_key from upstream callers + kwargs_for_preset = {k: v for k, v in kwargs.items() if k != "preset_cache_key"} self._set_preset_cache_key_in_kwargs( - preset_cache_key=hashed_cache_key, **kwargs + preset_cache_key=hashed_cache_key, **kwargs_for_preset ) return hashed_cache_key diff --git a/tests/local_testing/test_cache_preset_key.py b/tests/local_testing/test_cache_preset_key.py new file mode 100644 index 00000000000..de0ec05603c --- /dev/null +++ b/tests/local_testing/test_cache_preset_key.py @@ -0,0 +1,87 @@ +""" +Test for preset_cache_key multiple values bug fix. + +This test verifies that get_cache_key doesn't raise TypeError when kwargs +already contains preset_cache_key. + +Issue: When get_cache_key(**kwargs) is called with kwargs containing +preset_cache_key, the call to _set_preset_cache_key_in_kwargs() would fail with: + TypeError: got multiple values for keyword argument 'preset_cache_key' +""" + +import pytest +from unittest.mock import MagicMock, patch + + +class TestPresetCacheKeyFix: + """Tests for the preset_cache_key multiple values fix.""" + + def test_get_cache_key_with_preset_cache_key_in_kwargs(self): + """ + Test that get_cache_key handles kwargs that already contain preset_cache_key. + + This was causing: + TypeError: _set_preset_cache_key_in_kwargs() got multiple values + for keyword argument 'preset_cache_key' + """ + from litellm.caching.caching import Cache + + cache = Cache() + + # Simulate kwargs that already has preset_cache_key (as can happen + # when the cache key is recomputed in certain code paths) + kwargs_with_preset = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Hello"}], + "preset_cache_key": "existing_key_12345", # This caused the bug + "litellm_params": {}, + } + + # This should NOT raise TypeError + try: + result = cache.get_cache_key(**kwargs_with_preset) + assert result is not None + assert isinstance(result, str) + except TypeError as e: + if "multiple values for keyword argument" in str(e): + pytest.fail(f"Bug not fixed: {e}") + raise + + def test_get_cache_key_without_preset_cache_key(self): + """Test normal case without preset_cache_key in kwargs still works.""" + from litellm.caching.caching import Cache + + cache = Cache() + + kwargs_normal = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Hello"}], + "litellm_params": {}, + } + + result = cache.get_cache_key(**kwargs_normal) + assert result is not None + assert isinstance(result, str) + + def test_preset_cache_key_is_set_in_litellm_params(self): + """Verify that preset_cache_key is correctly set in litellm_params.""" + from litellm.caching.caching import Cache + + cache = Cache() + + litellm_params = {} + kwargs = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Hello"}], + "litellm_params": litellm_params, + } + + result = cache.get_cache_key(**kwargs) + + # The method should set preset_cache_key in litellm_params + assert "preset_cache_key" in litellm_params + assert litellm_params["preset_cache_key"] == result + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) From fc75380b88481bff17a0b25d6c7d7cf49e11c361 Mon Sep 17 00:00:00 2001 From: Dmitry Date: Sun, 5 Apr 2026 04:46:43 +0300 Subject: [PATCH 006/420] fix(presidio): use correct text positions in anonymize_text (#24998) * fix(presidio): use correct text positions in anonymize_text (#24160) The Presidio anonymizer endpoint returns items with start/end positions that reference the *anonymized output* text, not the original input. anonymize_text() was applying these positions to the original text, causing garbled output with remnants of un-masked PII data. When output_parse_pii is False, return redacted_text["text"] directly from the anonymizer response instead of manually splicing. When output_parse_pii is True, use analyze_results positions (which correctly reference the original text) to build numbered replacement tokens and the pii_tokens mapping. * address review: remove dead code, fix token numbering order - Remove unused `anon_item_by_entity` dict (Greptile P2) - Number tokens left-to-right ( first in text, not last) - Add assertion for token numbering order in test --- .../guardrails/guardrail_hooks/presidio.py | 106 ++++++----- .../guardrail_hooks/test_presidio.py | 168 ++++++++++++++++++ 2 files changed, 226 insertions(+), 48 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index 0f4ebbd4880..e048ca21cba 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -485,61 +485,71 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): redacted_text = await response.json() - new_text = text if redacted_text is not None: verbose_proxy_logger.debug("redacted_text: %s", redacted_text) - # Process items in reverse order by start position so that - # replacing later spans first does not shift earlier coordinates. - for item in sorted( - redacted_text["items"], key=lambda x: x["start"], reverse=True - ): - start = item["start"] - end = item["end"] - replacement = item["text"] # replacement token - if item["operator"] == "replace" and output_parse_pii is True: - if request_data is None: - verbose_proxy_logger.warning( - "Presidio anonymize_text called without request_data — " - "PII tokens cannot be stored per-request. " - "This may indicate a missing caller update." + + if not output_parse_pii: + # No need to build numbered tokens — just use Presidio's + # already-anonymized text directly. The old code incorrectly + # applied anonymizer item positions (which reference the + # *output* text) to the *original* text, causing offset errors. + for item in redacted_text.get("items", []): + entity_type = item.get("entity_type", None) + if entity_type is not None: + masked_entity_count[entity_type] = ( + masked_entity_count.get(entity_type, 0) + 1 ) - request_data = {} - # Store pii_tokens in metadata to avoid leaking to LLM providers. - # Providers like Anthropic reject unknown top-level fields. - if not request_data.get("metadata"): - request_data["metadata"] = {} - if "pii_tokens" not in request_data["metadata"]: - request_data["metadata"]["pii_tokens"] = {} - pii_tokens = request_data["metadata"]["pii_tokens"] + return redacted_text["text"] - # Append a sequential number to make each token unique - # per request, so unmasking maps back to the correct - # original value. Format: , - # This is LLM-friendly and degrades gracefully if the - # LLM doesn't echo the token verbatim. - seq = len(pii_tokens) + 1 - if replacement.endswith(">"): - replacement = f"{replacement[:-1]}_{seq}>" - else: - replacement = f"{replacement}_{seq}" + # output_parse_pii is True — we need sequentially numbered + # tokens and a pii_tokens mapping for later unmasking. + # Use analyze_results positions (which reference the ORIGINAL + # text) instead of anonymizer items (which reference the output). + new_text = text + if request_data is None: + verbose_proxy_logger.warning( + "Presidio anonymize_text called without request_data — " + "PII tokens cannot be stored per-request. " + "This may indicate a missing caller update." + ) + request_data = {} + if not request_data.get("metadata"): + request_data["metadata"] = {} + if "pii_tokens" not in request_data["metadata"]: + request_data["metadata"]["pii_tokens"] = {} + pii_tokens = request_data["metadata"]["pii_tokens"] - # Use ORIGINAL text (not new_text) since start/end - # reference the original text's coordinates. - pii_tokens[replacement] = text[start:end] + # Assign sequence numbers in forward (left-to-right) order so + # that is the first entity in the text, etc. + sorted_forward = sorted( + analyze_results, key=lambda x: x["start"] + ) + seq_map = {} + for idx, ar in enumerate(sorted_forward, start=1): + seq_map[(ar["start"], ar["end"])] = idx + # Apply replacements in reverse order by start position so + # that replacing later spans first does not shift earlier + # coordinates in the original text. + for ar in reversed(sorted_forward): + start = ar["start"] + end = ar["end"] + entity_type = ar["entity_type"] + replacement = f"<{entity_type}>" + + seq = seq_map[(start, end)] + if replacement.endswith(">"): + replacement = f"{replacement[:-1]}_{seq}>" + else: + replacement = f"{replacement}_{seq}" + + pii_tokens[replacement] = text[start:end] new_text = new_text[:start] + replacement + new_text[end:] - entity_type = item.get("entity_type", None) - if entity_type is not None: - masked_entity_count[entity_type] = ( - masked_entity_count.get(entity_type, 0) + 1 - ) - # When output_parse_pii is True, new_text contains sequentially - # numbered tokens (e.g. ) that match the keys - # in pii_tokens. Returning redacted_text["text"] (Presidio's - # original output) would send un-numbered tokens to the LLM, - # making unmasking impossible. - # When output_parse_pii is False, new_text == redacted_text["text"] - # because no suffix is appended. + + masked_entity_count[entity_type] = ( + masked_entity_count.get(entity_type, 0) + 1 + ) + return new_text else: raise Exception("Invalid anonymizer response: received None") diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py index 32a8c1b1070..38ea42285c1 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_presidio.py @@ -2230,3 +2230,171 @@ async def test_apply_to_output_streaming_bytes_only_logs_warning(): mock_logger.warning.assert_called_once() warning_msg = mock_logger.warning.call_args[0][0] assert "Output PII masking was skipped" in warning_msg + + +@pytest.mark.asyncio +async def test_anonymize_text_uses_correct_positions_no_parse_pii(): + """ + Regression test for anonymizer offset bug (fixes #24160). + + The Presidio anonymizer returns items with start/end positions that + reference the *anonymized output* text, not the original input text. + When output_parse_pii is False, anonymize_text must return + redacted_text["text"] directly instead of manually splicing the + original text using those positions, which produces garbled output + with remnants of original PII data. + """ + original_text = ( + "My name is John Smith, my email is john@example.com, phone 555-867-5309" + ) + # Positions as returned by the analyzer (reference original text) + analyze_results = [ + {"end": 51, "entity_type": "EMAIL_ADDRESS", "score": 1.0, "start": 35}, + {"end": 21, "entity_type": "PERSON", "score": 0.85, "start": 11}, + {"end": 71, "entity_type": "PHONE_NUMBER", "score": 0.75, "start": 59}, + ] + # Anonymizer response — positions reference the *anonymized* text + anonymizer_response = { + "text": "My name is , my email is , phone ", + "items": [ + { + "start": 56, + "end": 70, + "entity_type": "PHONE_NUMBER", + "text": "", + "operator": "replace", + }, + { + "start": 33, + "end": 48, + "entity_type": "EMAIL_ADDRESS", + "text": "", + "operator": "replace", + }, + { + "start": 11, + "end": 19, + "entity_type": "PERSON", + "text": "", + "operator": "replace", + }, + ], + } + + guardrail = _OPTIONAL_PresidioPIIMasking( + presidio_analyzer_api_base="http://test-analyzer/", + presidio_anonymizer_api_base="http://test-anonymizer/", + mock_testing=False, + ) + + mock_iterator = _make_mock_session_iterator( + json_response=anonymizer_response, + ) + + masked_entity_count = {} + with patch.object(guardrail, "_get_session_iterator", mock_iterator): + result = await guardrail.anonymize_text( + text=original_text, + analyze_results=analyze_results, + output_parse_pii=False, + masked_entity_count=masked_entity_count, + ) + + expected = "My name is , my email is , phone " + assert result == expected, ( + f"anonymize_text produced garbled output with PII remnants.\n" + f"Expected: {expected!r}\n" + f"Got: {result!r}" + ) + assert masked_entity_count == { + "PERSON": 1, + "EMAIL_ADDRESS": 1, + "PHONE_NUMBER": 1, + } + + +@pytest.mark.asyncio +async def test_anonymize_text_uses_correct_positions_with_parse_pii(): + """ + Regression test for anonymizer offset bug with output_parse_pii=True + (fixes #24160). + + When output_parse_pii is True, anonymize_text must use positions from + analyze_results (which reference the original text) to build numbered + tokens and the pii_tokens mapping, not positions from anonymizer items + (which reference the anonymized output text). + """ + original_text = ( + "My name is John Smith, my email is john@example.com, phone 555-867-5309" + ) + analyze_results = [ + {"end": 51, "entity_type": "EMAIL_ADDRESS", "score": 1.0, "start": 35}, + {"end": 21, "entity_type": "PERSON", "score": 0.85, "start": 11}, + {"end": 71, "entity_type": "PHONE_NUMBER", "score": 0.75, "start": 59}, + ] + anonymizer_response = { + "text": "My name is , my email is , phone ", + "items": [ + { + "start": 56, + "end": 70, + "entity_type": "PHONE_NUMBER", + "text": "", + "operator": "replace", + }, + { + "start": 33, + "end": 48, + "entity_type": "EMAIL_ADDRESS", + "text": "", + "operator": "replace", + }, + { + "start": 11, + "end": 19, + "entity_type": "PERSON", + "text": "", + "operator": "replace", + }, + ], + } + + guardrail = _OPTIONAL_PresidioPIIMasking( + presidio_analyzer_api_base="http://test-analyzer/", + presidio_anonymizer_api_base="http://test-anonymizer/", + mock_testing=False, + output_parse_pii=True, + ) + + mock_iterator = _make_mock_session_iterator( + json_response=anonymizer_response, + ) + + masked_entity_count = {} + request_data = {"metadata": {}} + with patch.object(guardrail, "_get_session_iterator", mock_iterator): + result = await guardrail.anonymize_text( + text=original_text, + analyze_results=analyze_results, + output_parse_pii=True, + masked_entity_count=masked_entity_count, + request_data=request_data, + ) + + # Result must not contain any remnants of original PII + assert "John" not in result + assert "john@example.com" not in result + assert "555-867-5309" not in result + + # pii_tokens must map numbered tokens back to correct original values + pii_tokens = request_data["metadata"]["pii_tokens"] + token_values = set(pii_tokens.values()) + assert "John Smith" in token_values + assert "john@example.com" in token_values + assert "555-867-5309" in token_values + + # Tokens must be numbered in left-to-right order of appearance: + # PERSON (pos 11) → _1, EMAIL_ADDRESS (pos 35) → _2, PHONE_NUMBER (pos 59) → _3 + assert pii_tokens.get("") == "John Smith" + assert pii_tokens.get("") == "john@example.com" + assert pii_tokens.get("") == "555-867-5309" From 23e702dae609aa79682b971967dcf0bc05efab59 Mon Sep 17 00:00:00 2001 From: Bohdan Kulinich Date: Sun, 5 Apr 2026 04:48:39 +0300 Subject: [PATCH 007/420] feat(prometheus): add 7m and 10m latency histogram buckets (#25071) Extend LATENCY_BUCKETS beyond 5 minutes so request/LLM latency metrics can distinguish long runs up to the typical default LLM request timeout. Made-with: Cursor --- litellm/types/integrations/prometheus.py | 2 ++ .../types/test_prometheus_latency_buckets.py | 17 +++++++++++++++++ 2 files changed, 19 insertions(+) create mode 100644 tests/test_litellm/types/test_prometheus_latency_buckets.py diff --git a/litellm/types/integrations/prometheus.py b/litellm/types/integrations/prometheus.py index 0d1501664b9..35b695cd054 100644 --- a/litellm/types/integrations/prometheus.py +++ b/litellm/types/integrations/prometheus.py @@ -156,6 +156,8 @@ LATENCY_BUCKETS = ( 180.0, 240.0, 300.0, + 420.0, # 7 minutes + 600.0, # 10 minutes (typical default LLM request timeout) float("inf"), ) diff --git a/tests/test_litellm/types/test_prometheus_latency_buckets.py b/tests/test_litellm/types/test_prometheus_latency_buckets.py new file mode 100644 index 00000000000..85670bb0b74 --- /dev/null +++ b/tests/test_litellm/types/test_prometheus_latency_buckets.py @@ -0,0 +1,17 @@ +"""LATENCY_BUCKETS covers long-running LLM calls (histograms are in seconds).""" + +import math + +from litellm.types.integrations.prometheus import LATENCY_BUCKETS + + +def test_latency_buckets_include_seven_and_ten_minutes(): + """Buckets beyond 5 min so histograms resolve requests up to default LLM timeouts.""" + assert 300.0 in LATENCY_BUCKETS + assert 420.0 in LATENCY_BUCKETS # 7 min + assert 600.0 in LATENCY_BUCKETS # 10 min + assert math.isinf(LATENCY_BUCKETS[-1]) + idx_300 = LATENCY_BUCKETS.index(300.0) + idx_420 = LATENCY_BUCKETS.index(420.0) + idx_600 = LATENCY_BUCKETS.index(600.0) + assert idx_300 < idx_420 < idx_600 From 4ca368923054ef1df73c1f3b815e39716ff2ce71 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sun, 5 Apr 2026 01:15:08 -0700 Subject: [PATCH 008/420] chore: fixes --- .../workflows/run_llm_translation_tests.py | 0 .trivyignore | 12 - ci_cd/.grype.yaml | 36 --- ci_cd/security_scans.sh | 261 ------------------ docs/my-website/.trivyignore | 7 - ui/litellm-dashboard/.trivyignore | 7 - 6 files changed, 323 deletions(-) mode change 100755 => 100644 .github/workflows/run_llm_translation_tests.py delete mode 100644 .trivyignore delete mode 100644 ci_cd/.grype.yaml delete mode 100755 ci_cd/security_scans.sh delete mode 100644 docs/my-website/.trivyignore delete mode 100644 ui/litellm-dashboard/.trivyignore diff --git a/.github/workflows/run_llm_translation_tests.py b/.github/workflows/run_llm_translation_tests.py old mode 100755 new mode 100644 diff --git a/.trivyignore b/.trivyignore deleted file mode 100644 index 0d04ecacdb5..00000000000 --- a/.trivyignore +++ /dev/null @@ -1,12 +0,0 @@ -# LiteLLM Trivy Ignore File -# CVEs listed here are temporarily allowlisted pending fixes - -# Next.js vulnerabilities in UI dashboard (next@14.2.35) -# Allowlisted: 2026-01-31, 7-day fix timeline -# Fix: Upgrade to Next.js 15.5.10+ or 16.1.5+ - -# HIGH: DoS via request deserialization -GHSA-h25m-26qc-wcjf - -# MEDIUM: Image Optimizer DoS -CVE-2025-59471 diff --git a/ci_cd/.grype.yaml b/ci_cd/.grype.yaml deleted file mode 100644 index b9bc9db58f5..00000000000 --- a/ci_cd/.grype.yaml +++ /dev/null @@ -1,36 +0,0 @@ -ignore: - - vulnerability: CVE-2026-22184 - reason: no fixed zlib package is available yet in the Wolfi repositories, so this is ignored temporarily until an upstream release exists - # Wolfi base image: Python 3.13 and Node from apk have no fixed builds in Wolfi yet / not applicable - - vulnerability: CVE-2025-55130 - reason: Node in Wolfi apk; only used for Admin UI build/prisma - - vulnerability: CVE-2025-59465 - reason: Node in Wolfi apk; only used for Admin UI build/prisma - - vulnerability: CVE-2025-55131 - reason: Node in Wolfi apk; only used for Admin UI build/prisma - - vulnerability: CVE-2025-59466 - reason: Node in Wolfi apk; only used for Admin UI build/prisma - - vulnerability: CVE-2026-21637 - reason: Node in Wolfi apk; only used for Admin UI build/prisma - - vulnerability: CVE-2025-55132 - reason: Node in Wolfi apk; only used for Admin UI build/prisma - - vulnerability: GHSA-hx9q-6w63-j58v - reason: orjson dumps recursion; allowlisted - - vulnerability: GHSA-73rr-hh4g-fpgx - reason: diff npm transitive dep; override in package.json, allowlisted - - vulnerability: CVE-2026-0865 - reason: Python 3.13 in Wolfi base; no fixed apk build yet - - vulnerability: CVE-2025-15282 - reason: Python 3.13 in Wolfi base; no fixed apk build yet - - vulnerability: CVE-2026-0672 - reason: Python 3.13 in Wolfi base; no fixed apk build yet - - vulnerability: CVE-2025-15366 - reason: Python 3.13 in Wolfi base; no fixed apk build yet - - vulnerability: CVE-2025-15367 - reason: Python 3.13 in Wolfi base; no fixed apk build yet - - vulnerability: CVE-2025-11468 - reason: Python 3.13 in Wolfi base; no fixed apk build yet - - vulnerability: CVE-2025-12781 - reason: Python 3.13 in Wolfi base; no fixed apk build yet - - vulnerability: CVE-2026-1299 - reason: Python 3.13 in Wolfi base; no fixed apk build yet diff --git a/ci_cd/security_scans.sh b/ci_cd/security_scans.sh deleted file mode 100755 index 2138fca6cd5..00000000000 --- a/ci_cd/security_scans.sh +++ /dev/null @@ -1,261 +0,0 @@ -#!/bin/bash - -# Security Scans Script for LiteLLM -# This script runs comprehensive security scans including Trivy and Grype - -set -e - -echo "Starting security scans for LiteLLM..." - -# Function to install Trivy and required tools -install_trivy() { - echo "Installing Trivy and required tools..." - TRIVY_VERSION="0.35.0" - sudo apt-get update - sudo apt-get install -y wget jq curl bsdmainutils - wget -qO trivy.deb "https://github.com/aquasecurity/trivy/releases/download/v${TRIVY_VERSION}/trivy_${TRIVY_VERSION}_Linux-64bit.deb" - sudo dpkg -i trivy.deb - rm trivy.deb - echo "Trivy ${TRIVY_VERSION} installed successfully" -} - -# Function to install Grype -install_grype() { - echo "Installing Grype..." - curl -sSfL https://raw.githubusercontent.com/anchore/grype/main/install.sh | sudo sh -s -- -b /usr/local/bin - echo "Grype installed successfully" -} - -# Function to install ggshield -install_ggshield() { - echo "Installing ggshield..." - pip3 install --upgrade pip - pip3 install ggshield - echo "ggshield installed successfully" -} - -# # Function to run secret detection scans -# run_secret_detection() { -# echo "Running secret detection scans..." - -# if ! command -v ggshield &> /dev/null; then -# install_ggshield -# fi - -# # Check if GITGUARDIAN_API_KEY is set (required for CI/CD) -# if [ -z "$GITGUARDIAN_API_KEY" ]; then -# echo "Warning: GITGUARDIAN_API_KEY environment variable is not set." -# echo "ggshield requires a GitGuardian API key to scan for secrets." -# echo "Please set GITGUARDIAN_API_KEY in your CI/CD environment variables." -# exit 1 -# fi - -# echo "Scanning codebase for secrets..." -# echo "Note: Large codebases may take several minutes due to API rate limits (50 requests/minute on free plan)" -# echo "ggshield will automatically handle rate limits and retry as needed." -# echo "Binary files, cache files, and build artifacts are excluded via .gitguardian.yaml" - -# # Use --recursive for directory scanning and auto-confirm if prompted -# # .gitguardian.yaml will automatically exclude binary files, wheel files, etc. -# # GITGUARDIAN_API_KEY environment variable will be used for authentication -# echo y | ggshield secret scan path . --recursive || { -# echo "" -# echo "==========================================" -# echo "ERROR: Secret Detection Failed" -# echo "==========================================" -# echo "ggshield has detected secrets in the codebase." -# echo "Please review discovered secrets above, revoke any actively used secrets" -# echo "from underlying systems and make changes to inject secrets dynamically at runtime." -# echo "" -# echo "For more information, see: https://docs.gitguardian.com/secrets-detection/" -# echo "==========================================" -# echo "" -# exit 1 -# } - -# echo "Secret detection scans completed successfully" -# } - -# Function to run Trivy scans -run_trivy_scans() { - echo "Running Trivy scans..." - - echo "Scanning LiteLLM Docs..." - trivy fs --ignorefile .trivyignore --scanners vuln --dependency-tree --exit-code 1 --severity HIGH,CRITICAL,MEDIUM ./docs/ - - echo "Scanning LiteLLM UI..." - trivy fs --ignorefile .trivyignore --scanners vuln --dependency-tree --exit-code 1 --severity HIGH,CRITICAL,MEDIUM ./ui/ - - echo "Trivy scans completed successfully" -} - -# Function to build and scan Docker images with Grype -run_grype_scans() { - echo "Running Grype scans..." - - # Temporarily add wheel files to .dockerignore for security scans - echo "Temporarily modifying .dockerignore to exclude problematic wheel files..." - cp .dockerignore .dockerignore.backup 2>/dev/null || touch .dockerignore.backup - echo "/*.whl" >> .dockerignore - - # Build and scan Dockerfile.database - echo "Building and scanning Dockerfile.database..." - docker build --no-cache -t litellm-database:latest -f ./docker/Dockerfile.database . - grype litellm-database:latest --config ci_cd/.grype.yaml --fail-on critical - - # Build and scan main Dockerfile - echo "Building and scanning main Dockerfile..." - docker build --no-cache -t litellm:latest . - grype litellm:latest --config ci_cd/.grype.yaml --fail-on critical - - # Restore original .dockerignore - echo "Restoring original .dockerignore..." - mv .dockerignore.backup .dockerignore - - # Scan the locally built LiteLLM image for vulnerabilities with CVSS >= 4.0 - echo "Scanning locally built LiteLLM image for high-severity vulnerabilities..." - echo "Using locally built image: litellm:latest" - - # Allowlist of CVEs to be ignored in failure threshold/reporting - # - CVE-2025-8869: Not applicable on Python >=3.13 (PEP 706 implemented); pip fallback unused; no OS-level fix - # - GHSA-4xh5-x5gv-qwph: GitHub Security Advisory alias for CVE-2025-8869 - # - GHSA-5j98-mcp5-4vw2: glob CLI command injection via -c/--cmd; glob CLI is not used in the litellm runtime image, - # and the vulnerable versions are pulled in only via OS-level/node tooling outside of our application code - ALLOWED_CVES=( - "CVE-2025-8869" - "GHSA-4xh5-x5gv-qwph" - "CVE-2025-8291" # no fix available as of Oct 11, 2025 - "GHSA-5j98-mcp5-4vw2" - "CVE-2025-13836" # Python 3.13 HTTP response reading OOM/DoS - no fix available in base image - "CVE-2025-12084" # Python 3.13 xml.dom.minidom quadratic algorithm - no fix available in base image - "CVE-2025-60876" # BusyBox wget HTTP request splitting - no fix available in Chainguard Wolfi base image - "CVE-2026-0861" # Wolfi glibc still flagged even on 2.42-r5; upstream patched build unavailable yet - "CVE-2010-4756" # glibc glob DoS - awaiting patched Wolfi glibc build - "CVE-2019-1010022" # glibc stack guard bypass - awaiting patched Wolfi glibc build - "CVE-2019-1010023" # glibc ldd remap issue - awaiting patched Wolfi glibc build - "CVE-2019-1010024" # glibc ASLR mitigation bypass - awaiting patched Wolfi glibc build - "CVE-2019-1010025" # glibc pthread heap address leak - awaiting patched Wolfi glibc build - "CVE-2026-22184" # zlib untgz buffer overflow - untgz unused + no fixed Wolfi build yet - "GHSA-58pv-8j8x-9vj2" # jaraco.context path traversal - setuptools vendored only (v5.3.0), not used in application code (using v6.1.0+) - "GHSA-34x7-hfp2-rc4v" # node-tar hardlink path traversal - not applicable, tar CLI not exposed in application code - "GHSA-r6q2-hw4h-h46w" # node-tar not used by application runtime, Linux-only container, not affect by macOS APFS-specific exploit - "GHSA-8rrh-rw8j-w5fx" # wheel is from chainguard and will be handled by then TODO: Remove this after Chainguard updates the wheel - "CVE-2025-59465" # Node only used for Admin UI build/prisma - "CVE-2025-55131" # Node only used for Admin UI build/prisma - "CVE-2025-59466" # Node only used for Admin UI build/prisma - "CVE-2025-55130" # Node only used for Admin UI build/prisma - "CVE-2025-59467" # Node only used for Admin UI build/prisma - "CVE-2026-21637" # Node only used for Admin UI build/prisma - "CVE-2025-55132" # Node only used for Admin UI build/prisma - "GHSA-hx9q-6w63-j58v" # orjson dumps recursion; allowlisted - "CVE-2025-15281" # No fix available yet - "CVE-2026-0865" # No fix available yet - "CVE-2025-15282" # No fix available yet - "CVE-2026-0672" # No fix available yet - "CVE-2025-15366" # No fix available yet - "CVE-2025-15367" # No fix available yet - "CVE-2025-12781" # No fix available yet - "CVE-2025-11468" # No fix available yet - "CVE-2026-1299" # Python 3.13 email module header injection - not applicable, LiteLLM doesn't use BytesGenerator for email serialization - "CVE-2026-0775" # npm cli incorrect permission assignment - no fix available yet, npm is only used at build/prisma-generate time - "GHSA-3ppc-4f35-3m26" # minimatch ReDoS via repeated wildcards - from nodejs_wheel bundled npm, not used in application runtime code - "GHSA-83g3-92jg-28cx" # tar arbitrary file read/write via hardlink - from nodejs_wheel bundled npm, not used in application runtime code - "CVE-2026-2297" # Python 3.13 SourcelessFileLoader audit hook bypass - no fix available in base image - "GHSA-qffp-2rhf-9h96" # tar hardlink path traversal - from nodejs_wheel bundled npm, not used in application runtime code - "CVE-2026-2673" # OpenSSL 3.6.1 TLS 1.3 key exchange group negotiation issue - no fix available yet - "CVE-2026-3644" # Python 3.13 vulnerability - no fix available in base image - "CVE-2026-4224" # Python 3.13 Expat parser stack overflow in ElementDeclHandler - no fix available in base image - ) - - # Build JSON array of allowlisted CVE IDs for jq - ALLOWED_IDS_JSON=$(printf '%s\n' "${ALLOWED_CVES[@]}" | jq -R . | jq -s .) - - echo "Checking for vulnerabilities with CVSS score >= 4.0..." - echo "Allowlisted CVEs (ignored in threshold): ${ALLOWED_CVES[*]}" - echo "" - - # Show all high-severity vulnerabilities for transparency - TOTAL_HIGH_SEVERITY=$(grype litellm:latest -o json | jq -r ' - .matches[] - | select(.vulnerability.cvss[]?.metrics.baseScore >= 4.0) - | .vulnerability.id' | wc -l) - - if [ "$TOTAL_HIGH_SEVERITY" -gt 0 ]; then - echo "Total vulnerabilities found with CVSS >= 4.0: $TOTAL_HIGH_SEVERITY" - echo "" - echo "All high-severity vulnerabilities (including allowlisted):" - grype litellm:latest -o json | jq --argjson allow "$ALLOWED_IDS_JSON" -r ' - ["Package", "Version", "Vulnerability ID", "CVSS Score", "Allowlisted"], - (.matches[] - | select(.vulnerability.cvss[]?.metrics.baseScore >= 4.0) - | [.artifact.name, .artifact.version, .vulnerability.id, .vulnerability.cvss[0].metrics.baseScore, (if (.vulnerability.id as $id | $allow | index($id)) then "YES" else "NO" end)]) - | @tsv' | column -t -s $'\t' - echo "" - fi - - HIGH_SEVERITY_COUNT=$(grype litellm:latest -o json | jq --argjson allow "$ALLOWED_IDS_JSON" -r ' - .matches[] - | select(.vulnerability.cvss[]?.metrics.baseScore >= 4.0) - | select((.vulnerability.id as $id | $allow | index($id) | not)) - | .vulnerability.id' | wc -l) - - if [ "$HIGH_SEVERITY_COUNT" -gt 0 ]; then - echo "" - echo "==========================================" - echo "ERROR: Security Scan Failed" - echo "==========================================" - echo "Found $HIGH_SEVERITY_COUNT non-allowlisted vulnerabilities with CVSS score >= 4.0 in litellm:latest" - echo "" - echo "These vulnerabilities are NOT in the allowlist and must be addressed." - echo "Current allowlisted CVEs: ${ALLOWED_CVES[*]}" - echo "" - echo "Detailed vulnerability report:" - echo "" - grype litellm:latest -o json | jq --argjson allow "$ALLOWED_IDS_JSON" -r ' - ["Package", "Version", "Vulnerability ID", "CVSS Score", "Severity", "Fix Version", "Description"], - (.matches[] - | select(.vulnerability.cvss[]?.metrics.baseScore >= 4.0) - | select((.vulnerability.id as $id | $allow | index($id) | not)) - | [.artifact.name, .artifact.version, .vulnerability.id, .vulnerability.cvss[0].metrics.baseScore, .vulnerability.severity, (.vulnerability.fix.versions[0] // "No fix available"), .vulnerability.description]) - | @tsv' | column -t -s $'\t' - echo "" - echo "==========================================" - echo "Action Required:" - echo "==========================================" - echo "1. If a fix is available, update the package to the fixed version" - echo "2. If the vulnerability is not applicable or has no fix:" - echo " - Add the CVE/GHSA ID to ALLOWED_CVES array in ci_cd/security_scans.sh" - echo " - Add a comment explaining why it's safe to ignore" - echo "" - echo "Note: Some vulnerabilities may have multiple IDs (CVE-XXXX and GHSA-XXXX)." - echo "Add all relevant IDs to the allowlist if they refer to the same issue." - echo "==========================================" - echo "" - exit 1 - else - echo "No high-severity vulnerabilities (CVSS >= 4.0) found in litellm:latest" - fi - - echo "Grype scans completed successfully" -} - -# Main execution -main() { - echo "Installing security scanning tools..." - install_trivy - install_grype - - # echo "Running secret detection scans..." - # run_secret_detection - - echo "Running filesystem vulnerability scans..." - run_trivy_scans - - echo "Running Docker image vulnerability scans..." - run_grype_scans - - echo "All security scans completed successfully!" -} - -# Execute main function -main "$@" diff --git a/docs/my-website/.trivyignore b/docs/my-website/.trivyignore deleted file mode 100644 index 977504f2670..00000000000 --- a/docs/my-website/.trivyignore +++ /dev/null @@ -1,7 +0,0 @@ -# js-yaml CVE-2025-64718 -# This vulnerability is not applicable because we've forced js-yaml to version 4.1.1 -# via npm overrides in package.json. Trivy incorrectly reports this based on -# dependency requirements in the lockfile, but the actual installed version is 4.1.1. -# Verified with: npm list js-yaml -CVE-2025-64718 - diff --git a/ui/litellm-dashboard/.trivyignore b/ui/litellm-dashboard/.trivyignore deleted file mode 100644 index 977504f2670..00000000000 --- a/ui/litellm-dashboard/.trivyignore +++ /dev/null @@ -1,7 +0,0 @@ -# js-yaml CVE-2025-64718 -# This vulnerability is not applicable because we've forced js-yaml to version 4.1.1 -# via npm overrides in package.json. Trivy incorrectly reports this based on -# dependency requirements in the lockfile, but the actual installed version is 4.1.1. -# Verified with: npm list js-yaml -CVE-2025-64718 - From f233520c44061a07e3c43ba34ba6f79c73cc0bae Mon Sep 17 00:00:00 2001 From: Hendrik Jaks Date: Mon, 6 Apr 2026 21:13:06 +0300 Subject: [PATCH 009/420] fix(ui): resolve login redirect loop when reverse proxy adds HttpOnly to cookies (#23532) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(ui): resolve login redirect loop when reverse proxy adds HttpOnly to cookies When LiteLLM is behind nginx-ingress or similar with security-hardened configs, the reverse proxy adds HttpOnly to all Set-Cookie headers. This makes the JWT token unreadable by JavaScript, causing an infinite login redirect loop. Fix by returning the JWT token in the /v2/login response body so the frontend can set a JS-accessible cookie directly. Fixes #19663 Co-Authored-By: Claude Opus 4.6 * fix: address Greptile review feedback - Add window guard to setTokenCookie for SSR consistency with clearTokenCookies - Add SSR test for window undefined case - Add code comment explaining why JWT is included in response body Co-Authored-By: Claude Opus 4.6 * fix: address second round of Greptile review feedback - Add loginCall integration tests verifying setTokenCookie is called with token and skipped when absent (backward-compatibility path) - Use encodeURIComponent/decodeURIComponent in setTokenCookie/getCookie for defense-in-depth against non-standard token formats Co-Authored-By: Claude Opus 4.6 * Update ui/litellm-dashboard/src/utils/cookieUtils.ts Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * Update ui/litellm-dashboard/src/utils/cookieUtils.ts Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * fix(ui): use sessionStorage instead of cookie for login token storage Replace setTokenCookie (which is a no-op when reverse proxy adds HttpOnly) with storeLoginToken using sessionStorage. Add sessionStorage fallback to getCookie so the token is found even when the cookie is HttpOnly. Also handle '=' in cookie values with .slice(1).join("=") and clear sessionStorage on logout. Co-Authored-By: Claude Opus 4.6 * fix(ui): use shared getCookie in page.tsx and user_dashboard.tsx Replace local getCookie functions in page.tsx and user_dashboard.tsx with the shared one from cookieUtils that has the sessionStorage fallback. Without this, the HttpOnly cookie fix was incomplete — page.tsx (the dashboard entry point) could not read the token, causing the redirect loop to persist. Also scope the sessionStorage fallback to the "token" key only, and clear sessionStorage in page.tsx deleteCookie. Co-Authored-By: Claude Opus 4.6 * fix(ui): scope deleteCookie sessionStorage cleanup to token key only Also document the sessionStorage cross-tab trade-off: per-tab scope means users behind an HttpOnly proxy must log in once per tab, but this is intentional to avoid localStorage XSS exposure. Co-Authored-By: Claude Opus 4.6 * Update ui/litellm-dashboard/src/utils/cookieUtils.ts Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * Apply suggestion from @greptile-apps[bot] Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * style: remove stray double blank line in user_dashboard.tsx Co-Authored-By: Claude Opus 4.6 * fix(ui): guard storeLoginToken against empty/whitespace-only tokens Co-Authored-By: Claude Opus 4.6 * fix(ui): preserve sessionStorage token across beforeunload clear The existing beforeunload handler calls sessionStorage.clear() to flush cached UI data on page refresh. This also wiped the token stored by storeLoginToken, re-introducing the redirect loop after any page refresh in the HttpOnly proxy scenario. Now the token is saved and restored across the clear. Co-Authored-By: Claude Opus 4.6 * fix(ui): set JS-accessible cookie at /ui path as HttpOnly workaround sessionStorage alone is unreliable. Also set the token via document.cookie at path=/ui — nginx only adds HttpOnly to server-set Set-Cookie headers, so a JS-set cookie is always readable. Co-Authored-By: Claude Opus 4.6 (1M context) * fix(ui): use dynamic cookie path based on server_root_path Hardcoded path=/ui breaks when LiteLLM is deployed with a custom server_root_path. Now derives the cookie path from serverRootPath so it works at /ui, /myapp/ui, etc. Also reuse clearTokenCookies() in deleteCookie() to avoid duplication. Co-Authored-By: Claude Opus 4.6 (1M context) * refactor(ui): remove circular dependency in cookieUtils.ts Derive the UI cookie path from window.location.pathname instead of importing serverRootPath from networking.tsx. This breaks the cookieUtils → networking → cookieUtils cycle that could cause serverRootPath to be undefined under certain bundler configurations. Co-Authored-By: Claude Opus 4.6 (1M context) * fix(ui): harden getUiCookiePath regex and add missing tests - Use regex /\/ui(?=\/|$)/ to match "/ui" only as a full path segment, preventing false matches on paths like "/my-ui-tool/login". - Add unit tests for storeLoginToken empty/whitespace guard and cookie-at-/ui-path behavior. Co-Authored-By: Claude Opus 4.6 (1M context) * style: fix Black formatting in audit_logs.py Co-Authored-By: Claude Opus 4.6 (1M context) * fix CI: formatting, test params, remove token from login JSON Co-Authored-By: Claude Opus 4.6 (1M context) * fix: reformat with Black 23.x to match CI Co-Authored-By: Claude Opus 4.6 (1M context) * fix: keep token in login JSON body for UI storeLoginToken flow Co-Authored-By: Claude Opus 4.6 (1M context) * fix: use storeLoginToken in exchangeLoginCode, add credentials include Co-Authored-By: Claude Opus 4.6 (1M context) * revert: remove unrelated changes from HttpOnly cookie fix branch Reset files not related to the login cookie fix back to main: - prometheus.py, bedrock converse, guardrail handler - auth_checks.py, reset_budget_job.py, audit_logs.py - test_user_api_key_auth.py Co-Authored-By: Claude Opus 4.6 (1M context) * Revert "revert: remove unrelated changes from HttpOnly cookie fix branch" This reverts commit 0684a1e27521ada35cf8a4afbdff7aeaa87ff41d. * Revert "fix: use storeLoginToken in exchangeLoginCode, add credentials include" This reverts commit 866405f443eb2004ee56ed2f52c9047593d8cc6e. * Revert "fix: keep token in login JSON body for UI storeLoginToken flow" This reverts commit 086c41640c5749399f11aba298321d1d16e10a92. * Revert "fix: reformat with Black 23.x to match CI" This reverts commit b2c3334c888a9dbb60fd94cce12f43e8f3aa7e82. * Revert "fix CI: formatting, test params, remove token from login JSON" This reverts commit 2905d47bd4013b81fbaa01cf2c7402d51360f1e6. --------- Co-authored-by: Claude Opus 4.6 Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- litellm/proxy/proxy_server.py | 5 +- ui/litellm-dashboard/src/app/page.tsx | 16 +--- .../src/components/networking.test.ts | 33 ++++++++ .../src/components/networking.tsx | 6 +- .../src/components/user_dashboard.tsx | 12 ++- .../src/utils/cookieUtils.test.ts | 73 ++++++++++++++++- ui/litellm-dashboard/src/utils/cookieUtils.ts | 81 ++++++++++++++++++- 7 files changed, 198 insertions(+), 28 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 9738ae4f1a2..accf669f2e3 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -11506,8 +11506,11 @@ async def login_v2(request: Request): # noqa: PLR0915 litellm_dashboard_ui += "/ui/" litellm_dashboard_ui += "?login=success" + # Token is included in the response body so the UI can set a JS-accessible + # cookie even when a reverse proxy (e.g. nginx-ingress) adds HttpOnly to the + # server-set cookie, which would otherwise cause an infinite login redirect. json_response = JSONResponse( - content={"redirect_url": litellm_dashboard_ui}, + content={"redirect_url": litellm_dashboard_ui, "token": jwt_token}, status_code=status.HTTP_200_OK, ) json_response.set_cookie(key="token", value=jwt_token) diff --git a/ui/litellm-dashboard/src/app/page.tsx b/ui/litellm-dashboard/src/app/page.tsx index 44df1b5bd41..73c27c00724 100644 --- a/ui/litellm-dashboard/src/app/page.tsx +++ b/ui/litellm-dashboard/src/app/page.tsx @@ -42,6 +42,7 @@ import ToolPoliciesView from "@/components/ToolPoliciesView"; import SpendLogsTable from "@/components/view_logs"; import ViewUserDashboard from "@/components/view_users"; import { ThemeProvider } from "@/contexts/ThemeContext"; +import { clearTokenCookies, getCookie } from "@/utils/cookieUtils"; import { isJwtExpired } from "@/utils/jwtUtils"; import { buildLoginUrlWithReturn, consumeReturnUrl, isValidReturnUrl, normalizeUrlForCompare, storeReturnUrl } from "@/utils/returnUrlUtils"; import { formatUserRole, isAdminRole } from "@/utils/roles"; @@ -51,21 +52,12 @@ import { useRouter, useSearchParams } from "next/navigation"; import { Suspense, useEffect, useMemo, useRef, useState } from "react"; import { ConfigProvider, theme } from "antd"; -function getCookie(name: string) { - // Safer cookie read + decoding; handles '=' inside values - const match = document.cookie.split("; ").find((row) => row.startsWith(name + "=")); - if (!match) return null; - const value = match.slice(name.length + 1); - try { - return decodeURIComponent(value); - } catch { - return value; - } -} - function deleteCookie(name: string, path = "/") { // Best-effort client-side clear (works for non-HttpOnly cookies without Domain) document.cookie = `${name}=; Max-Age=0; Path=${path}`; + if (name === "token") { + clearTokenCookies(); + } } interface ProxySettings { diff --git a/ui/litellm-dashboard/src/components/networking.test.ts b/ui/litellm-dashboard/src/components/networking.test.ts index c57dcb97eb9..3c107fa586f 100644 --- a/ui/litellm-dashboard/src/components/networking.test.ts +++ b/ui/litellm-dashboard/src/components/networking.test.ts @@ -5,6 +5,7 @@ import * as Networking from "./networking"; vi.mock("@/utils/cookieUtils", () => ({ clearTokenCookies: vi.fn(), getCookie: vi.fn(), + storeLoginToken: vi.fn(), })); vi.mock("./molecules/notifications_manager", () => ({ @@ -79,6 +80,38 @@ describe("networking - expired session handling", () => { }); }); +describe("loginCall - storeLoginToken integration", () => { + const originalFetch = global.fetch; + + beforeEach(() => { + vi.clearAllMocks(); + }); + + afterEach(() => { + global.fetch = originalFetch; + }); + + it("calls storeLoginToken when response includes token", async () => { + global.fetch = vi.fn().mockResolvedValue({ + ok: true, + json: async () => ({ redirect_url: "/ui/?login=success", token: "my-jwt" }), + }) as any; + const { storeLoginToken } = await import("@/utils/cookieUtils"); + await Networking.loginCall("admin", "pass"); + expect(storeLoginToken).toHaveBeenCalledWith("my-jwt"); + }); + + it("does not call storeLoginToken when response has no token", async () => { + global.fetch = vi.fn().mockResolvedValue({ + ok: true, + json: async () => ({ redirect_url: "/ui/?login=success" }), + }) as any; + const { storeLoginToken } = await import("@/utils/cookieUtils"); + await Networking.loginCall("admin", "pass"); + expect(storeLoginToken).not.toHaveBeenCalled(); + }); +}); + describe("daily activity helpers", () => { const startTime = new Date("2025-02-12T00:00:00.000Z"); const endTime = new Date("2025-02-19T00:00:00.000Z"); diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 28f8d308de7..16f35605877 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -69,7 +69,7 @@ export const getInProductNudgesCall = async (accessToken: string) => { * Helper file for calls being made to proxy */ import MessageManager from "@/components/molecules/message_manager"; -import { clearTokenCookies } from "@/utils/cookieUtils"; +import { clearTokenCookies, storeLoginToken } from "@/utils/cookieUtils"; import { TagNewRequest, TagUpdateRequest, TagListResponse, TagInfoResponse } from "./tag_management/types"; import { Team } from "./key_team_helpers/key_list"; import { UserInfo } from "./view_users/types"; @@ -9255,14 +9255,14 @@ export const loginCall = async (username: string, password: string, useV3?: bool const exchangeData: LoginResponse = await exchangeResponse.json(); if (exchangeData.token) { - document.cookie = `token=${exchangeData.token}; path=/; SameSite=Lax`; + storeLoginToken(exchangeData.token); } return exchangeData; } // Backwards compatibility: v2 or old v3 returns token directly if (data.token) { - document.cookie = `token=${data.token}; path=/; SameSite=Lax`; + storeLoginToken(data.token); } return data; diff --git a/ui/litellm-dashboard/src/components/user_dashboard.tsx b/ui/litellm-dashboard/src/components/user_dashboard.tsx index f97d8ffab04..90eac56540d 100644 --- a/ui/litellm-dashboard/src/components/user_dashboard.tsx +++ b/ui/litellm-dashboard/src/components/user_dashboard.tsx @@ -1,5 +1,5 @@ "use client"; -import { clearTokenCookies } from "@/utils/cookieUtils"; +import { clearTokenCookies, getCookie } from "@/utils/cookieUtils"; import { Col, Grid } from "@tremor/react"; import { Typography } from "antd"; import { jwtDecode } from "jwt-decode"; @@ -35,12 +35,6 @@ export type UserInfo = { spend: number; }; -function getCookie(name: string) { - console.log("COOKIES", document.cookie); - const cookieValue = document.cookie.split("; ").find((row) => row.startsWith(name + "=")); - return cookieValue ? cookieValue.split("=")[1] : null; -} - interface UserDashboardProps { userID: string | null; userRole: string | null; @@ -103,7 +97,11 @@ const UserDashboard: React.FC = ({ // They are only cleared on logout useEffect(() => { const handleBeforeUnload = () => { + const token = sessionStorage.getItem("token"); sessionStorage.clear(); + if (token) { + sessionStorage.setItem("token", token); + } }; window.addEventListener("beforeunload", handleBeforeUnload); return () => window.removeEventListener("beforeunload", handleBeforeUnload); diff --git a/ui/litellm-dashboard/src/utils/cookieUtils.test.ts b/ui/litellm-dashboard/src/utils/cookieUtils.test.ts index 8b066e6a8ea..c7bd27a6a85 100644 --- a/ui/litellm-dashboard/src/utils/cookieUtils.test.ts +++ b/ui/litellm-dashboard/src/utils/cookieUtils.test.ts @@ -1,11 +1,12 @@ import { describe, it, expect, beforeEach, vi } from "vitest"; -import { clearTokenCookies, getCookie } from "./cookieUtils"; +import { clearTokenCookies, getCookie, storeLoginToken } from "./cookieUtils"; describe("cookieUtils", () => { beforeEach(() => { document.cookie.split(";").forEach((c) => { document.cookie = c.replace(/^ +/, "").replace(/=.*/, "=;expires=" + new Date().toUTCString() + ";path=/"); }); + sessionStorage.clear(); vi.spyOn(console, "log").mockImplementation(() => {}); }); @@ -116,6 +117,55 @@ describe("cookieUtils", () => { vi.restoreAllMocks(); }); + + it("should clear sessionStorage token", () => { + sessionStorage.setItem("token", "stored-token"); + clearTokenCookies(); + expect(sessionStorage.getItem("token")).toBeNull(); + }); + }); + + describe("storeLoginToken", () => { + it("should store the token in sessionStorage", () => { + storeLoginToken("my-jwt-token"); + expect(sessionStorage.getItem("token")).toBe("my-jwt-token"); + }); + + it("should overwrite an existing token in sessionStorage", () => { + storeLoginToken("old-token"); + expect(sessionStorage.getItem("token")).toBe("old-token"); + + storeLoginToken("new-token"); + expect(sessionStorage.getItem("token")).toBe("new-token"); + }); + + it("should not throw when window is undefined (server-side rendering)", () => { + const originalWindow = global.window; + delete (global as any).window; + + expect(() => storeLoginToken("token")).not.toThrow(); + + global.window = originalWindow; + }); + + it("should not store empty string token", () => { + storeLoginToken(""); + expect(sessionStorage.getItem("token")).toBeNull(); + }); + + it("should not store whitespace-only token", () => { + storeLoginToken(" "); + expect(sessionStorage.getItem("token")).toBeNull(); + }); + + it("should set a JS-accessible cookie at /ui path", () => { + const cookieSpy = vi.spyOn(document, "cookie", "set"); + storeLoginToken("my-jwt-token"); + expect(cookieSpy).toHaveBeenCalledWith( + expect.stringContaining("path=/ui") + ); + vi.restoreAllMocks(); + }); }); describe("getCookie", () => { @@ -141,5 +191,26 @@ describe("cookieUtils", () => { expect(getCookie("token")).toBe("token-value"); expect(getCookie("other")).toBe("other-value"); }); + + it("should handle values containing '=' characters", () => { + document.cookie = "token=abc=def=ghi; path=/"; + expect(getCookie("token")).toBe("abc=def=ghi"); + }); + + it("should fall back to sessionStorage when cookie is not found", () => { + sessionStorage.setItem("token", "session-stored-jwt"); + expect(getCookie("token")).toBe("session-stored-jwt"); + }); + + it("should prefer cookie over sessionStorage", () => { + document.cookie = "token=cookie-value; path=/"; + sessionStorage.setItem("token", "session-value"); + expect(getCookie("token")).toBe("cookie-value"); + }); + + it("should not fall back to sessionStorage for non-token keys", () => { + sessionStorage.setItem("other", "other-value"); + expect(getCookie("other")).toBeNull(); + }); }); }); diff --git a/ui/litellm-dashboard/src/utils/cookieUtils.ts b/ui/litellm-dashboard/src/utils/cookieUtils.ts index 01add36542c..b4493744ad4 100644 --- a/ui/litellm-dashboard/src/utils/cookieUtils.ts +++ b/ui/litellm-dashboard/src/utils/cookieUtils.ts @@ -2,6 +2,23 @@ * Utility functions for managing cookies */ +/** + * Returns the cookie path for the UI. + * Derives the path from window.location.pathname so it works when + * LiteLLM is deployed behind a subpath (e.g. /myapp/ui instead of /ui). + * No imports from networking.tsx to avoid circular dependencies. + */ +function getUiCookiePath(): string { + if (typeof window === "undefined") return "/ui"; + // Match "/ui" only as a full path segment (followed by "/" or end of string) + // to avoid false matches like "/my-ui-tool/login" → "/my-ui". + const match = window.location.pathname.match(/\/ui(?=\/|$)/); + if (match && match.index !== undefined) { + return window.location.pathname.substring(0, match.index + 3); + } + return "/ui"; +} + /** * Clears the token cookie from both root and /ui paths */ @@ -16,7 +33,8 @@ export function clearTokenCookies() { // Clear with various combinations of path and SameSite // Include current path in case of custom server root path const currentPath = window.location.pathname; - const paths = ["/", "/ui"]; + const uiCookiePath = getUiCookiePath(); + const paths = ["/", uiCookiePath]; // Add the current path directory if it's different from root and /ui if (currentPath && currentPath !== "/" && !currentPath.startsWith("/ui")) { @@ -43,7 +61,45 @@ export function clearTokenCookies() { }); }); - console.log("After clearing cookies:", document.cookie); + try { + sessionStorage.removeItem("token"); + } catch { + // sessionStorage may be unavailable + } + +} + +/** + * Stores the login token so the UI can read it even when a reverse proxy + * (e.g. nginx-ingress) adds HttpOnly to the server-set cookie. + * + * Strategy: + * 1. Set a JS-accessible cookie at path "/ui". Because nginx only modifies + * server-set Set-Cookie headers, a cookie created via document.cookie will + * never carry HttpOnly. Using path "/ui" avoids colliding with the + * server-set HttpOnly cookie at path "/". + * 2. Also store in sessionStorage as a secondary fallback. + */ +export function storeLoginToken(token: string) { + if (typeof window === "undefined") return; + if (!token || !token.trim()) return; + + // 1. JS-accessible cookie at /ui — survives same-tab navigations and + // is readable by getCookie() via document.cookie. + try { + const secure = window.location.protocol === "https:" ? "; Secure" : ""; + const cookiePath = getUiCookiePath(); + document.cookie = `token=${encodeURIComponent(token)}; path=${cookiePath}; SameSite=Lax${secure}`; + } catch { + // cookie setting may fail in restrictive environments + } + + // 2. sessionStorage backup + try { + sessionStorage.setItem("token", token); + } catch { + // sessionStorage may be unavailable (e.g. private browsing quota exceeded) + } } /** @@ -53,6 +109,23 @@ export function clearTokenCookies() { */ export function getCookie(name: string) { if (typeof document === "undefined") return null; - const cookieValue = document.cookie.split("; ").find((row) => row.startsWith(name + "=")); - return cookieValue ? cookieValue.split("=")[1] : null; + const row = document.cookie.split("; ").find((r) => r.startsWith(name + "=")); + if (row) { + const raw = row.split("=").slice(1).join("="); + try { + return decodeURIComponent(raw); + } catch { + return raw; + } + } + // Fallback to sessionStorage — covers the case where a reverse proxy + // added HttpOnly to the server-set cookie, making it invisible to JS. + if (name === "token" && typeof window !== "undefined") { + try { + return sessionStorage.getItem(name); + } catch { + return null; + } + } + return null; } From c2f486acfe09051f69a7d56ddaabb5ce1d330829 Mon Sep 17 00:00:00 2001 From: ishaan-berri <155045088+ishaan-berri@users.noreply.github.com> Date: Mon, 6 Apr 2026 13:26:49 -0700 Subject: [PATCH 010/420] fix: batch-limit stale managed object cleanup to prevent 300K row UPDATE (#25227) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Add STALE_OBJECT_CLEANUP_BATCH_SIZE constant Configurable batch limit (default 1000) for stale managed object cleanup, preventing unbounded UPDATE queries from hitting 300K+ rows at once. * Batch-limit stale managed object cleanup with single bounded SQL query Two fixes to _cleanup_stale_managed_objects: 1. Replace unbounded update_many with a single execute_raw using a subquery LIMIT, capping each poll cycle to STALE_OBJECT_CLEANUP_BATCH_SIZE rows. Zero rows loaded into Python memory — everything stays in Postgres. Uses the same PostgreSQL raw-SQL pattern as spend_log_cleanup.py (the proxy requires PostgreSQL per schema.prisma). 2. Extract _expire_stale_rows as a separate method for testability. Keeps the file_purpose='response' filter to avoid incorrectly expiring long-running batch or fine-tune jobs that legitimately exceed the staleness cutoff. --- .../common_utils/check_responses_cost.py | 45 +++++++++++++++---- litellm/constants.py | 3 ++ 2 files changed, 40 insertions(+), 8 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py index 54fbc7abcc5..dc0168683c8 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py @@ -11,6 +11,7 @@ from litellm._logging import verbose_proxy_logger from litellm.constants import ( MANAGED_OBJECT_STALENESS_CUTOFF_DAYS, MAX_OBJECTS_PER_POLL_CYCLE, + STALE_OBJECT_CLEANUP_BATCH_SIZE, ) if TYPE_CHECKING: @@ -32,21 +33,49 @@ class CheckResponsesCost: self.prisma_client: PrismaClient = prisma_client self.llm_router: Router = llm_router + async def _expire_stale_rows( + self, cutoff: datetime, batch_size: int + ) -> int: + """Execute the bounded UPDATE that marks stale rows as 'stale_expired'. + + Isolated so it can be swapped / mocked in tests without touching the + orchestration logic in ``_cleanup_stale_managed_objects``. + + Uses PostgreSQL syntax (``$1::timestamptz``, ``LIMIT``, double-quoted + identifiers) which is the only dialect the proxy supports — every + ``schema.prisma`` in the repo sets ``provider = "postgresql"``. + Same pattern as ``spend_log_cleanup.py``. + """ + return await self.prisma_client.db.execute_raw( + """ + UPDATE "LiteLLM_ManagedObjectTable" + SET "status" = 'stale_expired' + WHERE "id" IN ( + SELECT "id" FROM "LiteLLM_ManagedObjectTable" + WHERE "file_purpose" = 'response' + AND "status" NOT IN ('completed', 'complete', 'failed', 'expired', 'cancelled', 'stale_expired') + AND "created_at" < $1::timestamptz + ORDER BY "created_at" ASC + LIMIT $2 + ) + """, + cutoff, + batch_size, + ) + async def _cleanup_stale_managed_objects(self) -> None: """ Mark managed objects older than MANAGED_OBJECT_STALENESS_CUTOFF_DAYS days in non-terminal states as 'stale_expired'. These will never complete and should not be polled. + + Runs as a single DB query with a subquery LIMIT so no rows are loaded + into Python memory. Processes at most STALE_OBJECT_CLEANUP_BATCH_SIZE + rows per invocation to avoid overwhelming the DB when there is a large + backlog. """ cutoff = datetime.now(timezone.utc) - timedelta(days=MANAGED_OBJECT_STALENESS_CUTOFF_DAYS) - result = await self.prisma_client.db.litellm_managedobjecttable.update_many( - where={ - "file_purpose": "response", - "status": {"not_in": ["completed", "complete", "failed", "expired", "cancelled", "stale_expired"]}, - "created_at": {"lt": cutoff}, - }, - data={"status": "stale_expired"}, - ) + result = await self._expire_stale_rows(cutoff, STALE_OBJECT_CLEANUP_BATCH_SIZE) if result > 0: verbose_proxy_logger.warning( f"CheckResponsesCost: marked {result} stale managed objects " diff --git a/litellm/constants.py b/litellm/constants.py index 1af53b2dae0..28c6c0cc0e3 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1367,6 +1367,9 @@ MAX_OBJECTS_PER_POLL_CYCLE = max(1, int(os.getenv("MAX_OBJECTS_PER_POLL_CYCLE", MANAGED_OBJECT_STALENESS_CUTOFF_DAYS = max( 1, int(os.getenv("MANAGED_OBJECT_STALENESS_CUTOFF_DAYS", 7)) ) +STALE_OBJECT_CLEANUP_BATCH_SIZE = max( + 1, int(os.getenv("STALE_OBJECT_CLEANUP_BATCH_SIZE", 1000)) +) # Set PROXY_BATCH_POLLING_ENABLED=false to disable the CheckBatchCost and # CheckResponsesCost background polling jobs entirely (e.g. to avoid DB load on # installations with large numbers of stale managed objects). From 8b0fadf99c95d3f3dd51a957b3716af9aff7585a Mon Sep 17 00:00:00 2001 From: Otavio Brito <69211663+otaviofbrito@users.noreply.github.com> Date: Mon, 6 Apr 2026 17:42:24 -0300 Subject: [PATCH 011/420] [bug-fix]return actual status code - /v1/messages/count_tokens endpoint (#21352) * return actual status code - /count_tokens endpoint * Apply suggestion from @greptile-apps[bot] Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * fix greptile suggestion * rollback file * add test case --------- Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Co-authored-by: ishaan-berri <155045088+ishaan-berri@users.noreply.github.com> --- litellm/proxy/proxy_server.py | 26 +++++--- .../test_proxy_token_counter.py | 64 +++++++++++++++++++ 2 files changed, 81 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 9738ae4f1a2..715fc44114b 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -9117,15 +9117,23 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False) ) is True ): - result = await provider_counter.count_tokens( - model_to_use=model_to_use or "", - messages=messages, # type: ignore - contents=contents, - deployment=deployment, - request_model=request.model, - tools=tools, - system=system, - ) + try: + result = await provider_counter.count_tokens( + model_to_use=model_to_use or "", + messages=messages, # type: ignore + contents=contents, + deployment=deployment, + request_model=request.model, + ) + except httpx.HTTPStatusError as e: + error_message = getattr(e, "message", None) or str(e) + status_code = getattr(e, "status_code", None) or e.response.status_code + raise ProxyException( + message=error_message, + type="token_counting_error", + param="model", + code=status_code + ) ######################################################### # Transfrom the Response to the well known format ######################################################### diff --git a/tests/proxy_unit_tests/test_proxy_token_counter.py b/tests/proxy_unit_tests/test_proxy_token_counter.py index cca1e843604..ea25369d8f8 100644 --- a/tests/proxy_unit_tests/test_proxy_token_counter.py +++ b/tests/proxy_unit_tests/test_proxy_token_counter.py @@ -994,6 +994,70 @@ async def test_bedrock_handler_httpx_error_status_code_propagation(): assert exc_info.value.message == "Forbidden - Invalid credentials" +@pytest.mark.asyncio +async def test_token_counter_httpx_status_error_raises_proxy_exception(): + """ + When provider_counter.count_tokens() raises httpx.HTTPStatusError, + the token_counter endpoint should catch it and raise a ProxyException + with the upstream status code and error message. + """ + + upstream_status = 429 + upstream_message = "Rate limit exceeded" + response = httpx.Response( + status_code=upstream_status, + request=httpx.Request("POST", "https://provider.example.com/count"), + ) + http_error = httpx.HTTPStatusError( + message=upstream_message, + request=response.request, + response=response, + ) + + mock_counter = MagicMock() + mock_counter.should_use_token_counting_api.return_value = True + mock_counter.count_tokens = AsyncMock(side_effect=http_error) + + # Save originals + original_get_provider_token_counter = litellm.proxy.proxy_server._get_provider_token_counter + original_router = litellm.proxy.proxy_server.llm_router + + try: + def mock_get_provider_token_counter(deployment, model_to_use): + return (mock_counter, "claude-4-6-sonnet", "vertex_ai") + + litellm.proxy.proxy_server._get_provider_token_counter = mock_get_provider_token_counter + + mock_router = MagicMock() + mock_router.async_get_available_deployment = AsyncMock( + return_value={ + "litellm_params": { + "model": "vertex_ai/claude-4-6-sonnet", + "api_key": "fake-key", + }, + "model_info": {}, + } + ) + litellm.proxy.proxy_server.llm_router = mock_router + + with pytest.raises(ProxyException) as exc_info: + await token_counter( + request=TokenCountRequest( + model="claude-4-6-sonnet", + messages=[{"role": "user", "content": "hello"}], + ), + call_endpoint=True, + ) + + assert exc_info.value.code == str(upstream_status) + assert upstream_message in exc_info.value.message + assert exc_info.value.type == "token_counting_error" + assert exc_info.value.param == "model" + finally: + litellm.proxy.proxy_server._get_provider_token_counter = original_get_provider_token_counter + litellm.proxy.proxy_server.llm_router = original_router + + @pytest.mark.asyncio async def test_proxy_token_counter_error_raises_exception_when_disabled(): """ From 414d3966bfe212d5d2376c3b6ae1c3adbb0f67ef Mon Sep 17 00:00:00 2001 From: ishaan-berri <155045088+ishaan-berri@users.noreply.github.com> Date: Mon, 6 Apr 2026 13:48:43 -0700 Subject: [PATCH 012/420] feat(teams): per-member model scope + team default_team_member_models (#24950) * fix(bedrock): strip [1m]/[200k] context window suffixes before cost lookup * test(bedrock): add test for [1m] context window suffix stripping in cost lookup * schema: add allowed_models to BudgetTable, default_team_member_models to TeamTable * migration: add allowed_models and default_team_member_models columns * types: add allowed_models to TeamMemberAddRequest, TeamMemberUpdateRequest, UpdateTeamRequest * utils: add allowed_models param to add_new_member, persist to budget table * common_utils: add allowed_models to _upsert_budget_and_membership * team endpoints: seed allowed_models on member_add, persist on member_update and team/update * auth: enforce per-member allowed_models at request time * networking: add allowed_models to Member type and teamMemberUpdateCall * TeamMemberTab: add Model Scope column showing per-member allowed_models * EditMembership: add Allowed Models multi-select field * TeamInfo: add default_team_member_models field in Settings tab * chore: sync schema.prisma copies from root * fix(team_member_update): update existing budget in-place instead of creating new one When a member already has a budget_id, patch only the fields the caller provided rather than always creating a fresh budget record. The old code ignored existing_budget_id entirely, so updating only allowed_models silently dropped the stored max_budget / tpm_limit / rpm_limit values. * fix(auth): pass llm_router to _check_team_member_model_access Without the router, _can_object_call_model cannot resolve wildcard model names (e.g. openai/*) or access-group names in allowed_models, causing legitimate requests to be denied. Thread the existing llm_router from _run_common_checks through to the new member-scope check. * feat(ui): add Team Member Settings accordion to Create Team modal Groups default_team_member_models, member budget/key duration, and tpm/rpm defaults into a single collapsible section. The model picker is filtered to only show the models selected for the team, and the copy distinguishes it from the team-level Models field. * feat(ui): consolidate Team Member Settings into accordion in edit team form Moves default_team_member_models + per-member budget/key/tpm/rpm fields into a collapsible "Team Member Settings" panel. Keeps the top-level form focused on team-wide settings (team models, team budget, tpm/rpm). * fix(ui): use tremor Accordion for Team Member Settings in edit team form * fix(ui): move Team Member Settings accordion above budget fields in Create Team * chore: fixes --------- Co-authored-by: Ishaan Jaffer Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> Co-authored-by: Yuneng Jiang --- .../migration.sql | 9 ++ .../litellm_proxy_extras/schema.prisma | 4 +- litellm/llms/bedrock/common_utils.py | 16 +- litellm/proxy/_types.py | 18 +++ litellm/proxy/auth/auth_checks.py | 74 ++++++++- .../management_endpoints/common_utils.py | 38 ++++- .../management_endpoints/team_endpoints.py | 67 ++++---- litellm/proxy/management_helpers/utils.py | 24 +-- litellm/proxy/schema.prisma | 4 +- schema.prisma | 4 +- .../llms/bedrock/test_bedrock_common_utils.py | 43 +++++- .../components/modals/CreateTeamModal.tsx | 105 +++++++++---- .../src/components/networking.tsx | 4 + .../src/components/team/EditMembership.tsx | 15 +- .../src/components/team/TeamInfo.tsx | 146 +++++++++++++----- .../src/components/team/TeamMemberTab.tsx | 39 +++++ 16 files changed, 480 insertions(+), 130 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260401000000_add_team_member_model_scope/migration.sql diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260401000000_add_team_member_model_scope/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260401000000_add_team_member_model_scope/migration.sql new file mode 100644 index 00000000000..d34482b3390 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260401000000_add_team_member_model_scope/migration.sql @@ -0,0 +1,9 @@ +-- Add per-member model scope to LiteLLM_BudgetTable +-- allowed_models: empty array = inherit team models; non-empty = enforce member-level restriction +ALTER TABLE "LiteLLM_BudgetTable" + ADD COLUMN IF NOT EXISTS "allowed_models" TEXT[] DEFAULT ARRAY[]::TEXT[]; + +-- Add default_team_member_models to LiteLLM_TeamTable +-- Seeds allowed_models for newly added team members; empty = no per-member restriction +ALTER TABLE "LiteLLM_TeamTable" + ADD COLUMN IF NOT EXISTS "default_team_member_models" TEXT[] DEFAULT ARRAY[]::TEXT[]; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index fce95465b55..03760695f30 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -17,8 +17,9 @@ model LiteLLM_BudgetTable { tpm_limit BigInt? rpm_limit BigInt? model_max_budget Json? - budget_duration String? + budget_duration String? budget_reset_at DateTime? + allowed_models String[] @default([]) // per-member model scope; empty = inherit team models created_at DateTime @default(now()) @map("created_at") created_by String updated_at DateTime @default(now()) @updatedAt @map("updated_at") @@ -140,6 +141,7 @@ model LiteLLM_TeamTable { team_member_permissions String[] @default([]) access_group_ids String[] @default([]) policies String[] @default([]) + default_team_member_models String[] @default([]) // default allowed_models for newly added team members; empty = no per-member restriction model_id Int? @unique // id for LiteLLM_ModelTable -> stores team-level model aliases allow_team_guardrail_config Boolean @default(false) // if true, team admin can configure guardrails for this team litellm_organization_table LiteLLM_OrganizationTable? @relation(fields: [organization_id], references: [organization_id]) diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 9666aa68c99..a7cb6ed9f15 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -432,12 +432,16 @@ def strip_bedrock_routing_prefix(model: str) -> str: def strip_bedrock_throughput_suffix(model: str) -> str: - """Strip throughput tier suffixes from Bedrock model names.""" + """Strip throughput tier suffixes and context window suffixes from Bedrock model names.""" import re # Pattern matches model:version:throughput where throughput is like 51k, 18k, etc. # Keep the model:version part, strip the :throughput suffix - return re.sub(r"(:\d+):\d+k$", r"\1", model) + model = re.sub(r"(:\d+):\d+k$", r"\1", model) + # Strip context window suffixes like [1m], [200k], etc. + # e.g. "us.anthropic.claude-opus-4-6-v1[1m]" -> "us.anthropic.claude-opus-4-6-v1" + model = re.sub(r"\[\w+\]$", "", model) + return model def get_bedrock_base_model(model: str) -> str: @@ -1062,9 +1066,11 @@ class CommonBatchFilesUtils: return ( dict(prepped.headers), - request_data.encode("utf-8") - if isinstance(request_data, str) - else request_data, + ( + request_data.encode("utf-8") + if isinstance(request_data, str) + else request_data + ), ) def generate_unique_job_name(self, model: str, prefix: str = "litellm") -> str: diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 441b3b836a1..452ba7f7995 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1649,6 +1649,9 @@ class TeamBase(LiteLLMPydanticObjectBase): blocked: bool = False router_settings: Optional[dict] = None access_group_ids: Optional[List[str]] = None + default_team_member_models: Optional[List[str]] = ( + None # default allowed_models seeded onto new team members + ) class NewTeamRequest(TeamBase): @@ -1742,6 +1745,9 @@ class UpdateTeamRequest(LiteLLMPydanticObjectBase): enforced_file_expires_after: Optional[dict] = None router_settings: Optional[dict] = None access_group_ids: Optional[List[str]] = None + default_team_member_models: Optional[List[str]] = ( + None # default allowed_models seeded onto new team members + ) class ResetTeamBudgetRequest(LiteLLMPydanticObjectBase): @@ -1943,6 +1949,9 @@ class LiteLLM_BudgetTable(LiteLLMPydanticObjectBase): rpm_limit: Optional[int] = None model_max_budget: Optional[dict] = None budget_duration: Optional[str] = None + allowed_models: Optional[List[str]] = ( + None # per-member model scope; empty = inherit team models + ) model_config = ConfigDict(protected_namespaces=()) @@ -3735,6 +3744,10 @@ class TeamMemberAddRequest(MemberAddRequest): default=None, description="Maximum budget allocated to this user within the team. If not set, user has unlimited budget within team limits", ) + allowed_models: Optional[List[str]] = Field( + default=None, + description="List of models this team member can access. If not set, inherits the team's default_team_member_models or all team models.", + ) class TeamMemberDeleteRequest(MemberDeleteRequest): @@ -3750,6 +3763,10 @@ class TeamMemberUpdateRequest(TeamMemberDeleteRequest): rpm_limit: Optional[int] = Field( default=None, description="Requests per minute limit for this team member" ) + allowed_models: Optional[List[str]] = Field( + default=None, + description="List of models this team member can access. Pass an empty list to remove per-member model restrictions.", + ) class TeamMemberUpdateResponse(MemberUpdateResponse): @@ -3757,6 +3774,7 @@ class TeamMemberUpdateResponse(MemberUpdateResponse): max_budget_in_team: Optional[float] = None tpm_limit: Optional[int] = None rpm_limit: Optional[int] = None + allowed_models: Optional[List[str]] = None class TeamModelAddRequest(BaseModel): diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 68bde8434a6..bf0de3d3ca8 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -8,6 +8,7 @@ Run checks for: 2. If user is in budget 3. If end_user ('user' passed to /chat/completions, /embeddings endpoint) is in budget """ + import asyncio import re import time @@ -456,9 +457,9 @@ async def common_checks( # noqa: PLR0915 model=_model, team_object=team_object, llm_router=llm_router, - team_model_aliases=valid_token.team_model_aliases - if valid_token - else None, + team_model_aliases=( + valid_token.team_model_aliases if valid_token else None + ), ): raise ProxyException( message=f"Team not allowed to access model. Team={team_object.team_id}, Model={_model}. Allowed team models = {team_object.models}", @@ -467,6 +468,21 @@ async def common_checks( # noqa: PLR0915 code=status.HTTP_401_UNAUTHORIZED, ) + # 2.2. If team member has per-member model scope, enforce it + if _model and team_object and valid_token and valid_token.user_id: + with tracer.trace( + "litellm.proxy.auth.common_checks.check_team_member_model_access" + ): + await _check_team_member_model_access( + model=_model, + team_object=team_object, + valid_token=valid_token, + llm_router=llm_router, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + # Require trace id for agent keys when agent has require_trace_id_on_calls_by_agent if valid_token is not None and valid_token.agent_id: from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry @@ -3108,6 +3124,58 @@ async def _check_team_member_budget( ) +async def _check_team_member_model_access( + model: Union[str, List[str]], + team_object: LiteLLM_TeamTable, + valid_token: UserAPIKeyAuth, + llm_router: Optional[Router], + prisma_client: Optional["PrismaClient"], + user_api_key_cache: DualCache, + proxy_logging_obj: ProxyLogging, +) -> None: + """ + Check if a team member's per-member model scope allows access to the requested model. + + Only enforced when the member's budget table has a non-empty allowed_models list. + If allowed_models is empty or absent, the team-level models list applies (no extra restriction). + """ + if valid_token.user_id is None or team_object.team_id is None: + return + + team_membership = await get_team_membership( + user_id=valid_token.user_id, + team_id=team_object.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + if ( + team_membership is None + or team_membership.litellm_budget_table is None + or not team_membership.litellm_budget_table.allowed_models + ): + return # no per-member restriction — inherit team-level check + + member_allowed_models: List[str] = ( + team_membership.litellm_budget_table.allowed_models + ) + try: + _can_object_call_model( + model=model, + llm_router=llm_router, + models=member_allowed_models, + object_type="team", + ) + except ProxyException: + raise ProxyException( + message=f"Team member not allowed to access model. User={valid_token.user_id}, Team={team_object.team_id}, Model={model}. Allowed member models = {member_allowed_models}", + type=ProxyErrorTypes.team_model_access_denied, + param="model", + code=status.HTTP_401_UNAUTHORIZED, + ) + + async def _team_max_budget_check( team_object: Optional[LiteLLM_TeamTable], valid_token: Optional[UserAPIKeyAuth], diff --git a/litellm/proxy/management_endpoints/common_utils.py b/litellm/proxy/management_endpoints/common_utils.py index efc42d3355c..07286b4fa80 100644 --- a/litellm/proxy/management_endpoints/common_utils.py +++ b/litellm/proxy/management_endpoints/common_utils.py @@ -1,4 +1,4 @@ -from typing import TYPE_CHECKING, Any, Dict, Optional, Union +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union from litellm._logging import verbose_proxy_logger from litellm.caching import DualCache @@ -354,6 +354,7 @@ async def _upsert_budget_and_membership( user_api_key_dict: UserAPIKeyAuth, tpm_limit: Optional[int] = None, rpm_limit: Optional[int] = None, + allowed_models: Optional[List[str]] = None, ): """ Helper function to Create/Update or Delete the budget within the team membership @@ -366,11 +367,17 @@ async def _upsert_budget_and_membership( user_api_key_dict: User API Key dictionary containing user information tpm_limit: Tokens per minute limit for the team member rpm_limit: Requests per minute limit for the team member + allowed_models: Per-member model scope. None = don't change. [] = remove restrictions. Non-empty list = enforce. - If max_budget, tpm_limit, and rpm_limit are all None, the user's budget is removed from the team membership. + If max_budget, tpm_limit, rpm_limit, and allowed_models are all None, the user's budget is removed from the team membership. If any of these values exist, a budget is updated or created and linked to the team membership. """ - if max_budget is None and tpm_limit is None and rpm_limit is None: + if ( + max_budget is None + and tpm_limit is None + and rpm_limit is None + and allowed_models is None + ): # disconnect the budget since all limits are None await tx.litellm_teammembership.update( where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}}, @@ -378,7 +385,27 @@ async def _upsert_budget_and_membership( ) return - # create a new budget + if existing_budget_id is not None: + # Update the existing budget in-place to preserve fields not being changed. + # Only write fields that the caller explicitly provided (non-None). + update_data: Dict[str, Any] = { + "updated_by": user_api_key_dict.user_id or "", + } + if max_budget is not None: + update_data["max_budget"] = max_budget + if tpm_limit is not None: + update_data["tpm_limit"] = tpm_limit + if rpm_limit is not None: + update_data["rpm_limit"] = rpm_limit + if allowed_models is not None: + update_data["allowed_models"] = allowed_models + await tx.litellm_budgettable.update( + where={"budget_id": existing_budget_id}, + data=update_data, + ) + return + + # No existing budget — create a new one and link it to the membership. create_data: Dict[str, Any] = { "created_by": user_api_key_dict.user_id or "", "updated_by": user_api_key_dict.user_id or "", @@ -389,12 +416,13 @@ async def _upsert_budget_and_membership( create_data["tpm_limit"] = tpm_limit if rpm_limit is not None: create_data["rpm_limit"] = rpm_limit + if allowed_models is not None: + create_data["allowed_models"] = allowed_models new_budget = await tx.litellm_budgettable.create( data=create_data, include={"team_membership": True}, ) - # upsert the team membership with the new/updated budget await tx.litellm_teammembership.upsert( where={ "user_id_team_id": { diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index e1534573789..c980fcb82f8 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -1545,12 +1545,12 @@ async def update_team( # noqa: PLR0915 updated_kv["router_settings"] = safe_dumps(updated_kv["router_settings"]) updated_kv = prisma_client.jsonify_team_object(db_data=updated_kv) - team_row: Optional[ - LiteLLM_TeamTable - ] = await prisma_client.db.litellm_teamtable.update( - where={"team_id": data.team_id}, - data=updated_kv, - include={"litellm_model_table": True}, # type: ignore + team_row: Optional[LiteLLM_TeamTable] = ( + await prisma_client.db.litellm_teamtable.update( + where={"team_id": data.team_id}, + data=updated_kv, + include={"litellm_model_table": True}, # type: ignore + ) ) if team_row is None or team_row.team_id is None: @@ -1764,6 +1764,11 @@ async def _process_team_members( else None ) + # Resolve allowed_models: explicit request value, or fall back to team's default_team_member_models + member_allowed_models = data.allowed_models + if member_allowed_models is None and complete_team_data.default_team_member_models: + member_allowed_models = complete_team_data.default_team_member_models + if isinstance(data.member, Member): try: updated_user, updated_tm = await add_new_member( @@ -1774,6 +1779,7 @@ async def _process_team_members( litellm_proxy_admin_name=litellm_proxy_admin_name, team_id=data.team_id, default_team_budget_id=default_team_budget_id, + allowed_models=member_allowed_models, ) except Exception as e: raise HTTPException( @@ -1798,6 +1804,7 @@ async def _process_team_members( litellm_proxy_admin_name=litellm_proxy_admin_name, team_id=data.team_id, default_team_budget_id=default_team_budget_id, + allowed_models=member_allowed_models, ) except Exception as e: raise HTTPException( @@ -2297,13 +2304,13 @@ async def team_member_delete( ) # Fetch keys before deletion to persist them - keys_to_delete: List[ - LiteLLM_VerificationToken - ] = await prisma_client.db.litellm_verificationtoken.find_many( - where={ - "user_id": {"in": list(user_ids_to_delete)}, - "team_id": data.team_id, - } + keys_to_delete: List[LiteLLM_VerificationToken] = ( + await prisma_client.db.litellm_verificationtoken.find_many( + where={ + "user_id": {"in": list(user_ids_to_delete)}, + "team_id": data.team_id, + } + ) ) if keys_to_delete: @@ -2435,6 +2442,7 @@ async def team_member_update( user_api_key_dict=user_api_key_dict, tpm_limit=data.tpm_limit, rpm_limit=data.rpm_limit, + allowed_models=data.allowed_models, ) ### update team member role @@ -2467,6 +2475,7 @@ async def team_member_update( max_budget_in_team=data.max_budget_in_team, tpm_limit=data.tpm_limit, rpm_limit=data.rpm_limit, + allowed_models=data.allowed_models, ) @@ -2687,10 +2696,10 @@ async def delete_team( team_rows: List[LiteLLM_TeamTable] = [] for team_id in data.team_ids: try: - team_row_base: Optional[ - BaseModel - ] = await prisma_client.db.litellm_teamtable.find_unique( - where={"team_id": team_id} + team_row_base: Optional[BaseModel] = ( + await prisma_client.db.litellm_teamtable.find_unique( + where={"team_id": team_id} + ) ) if team_row_base is None: raise Exception @@ -2749,10 +2758,10 @@ async def delete_team( _persist_deleted_verification_tokens, ) - keys_to_delete: List[ - LiteLLM_VerificationToken - ] = await prisma_client.db.litellm_verificationtoken.find_many( - where={"team_id": {"in": data.team_ids}} + keys_to_delete: List[LiteLLM_VerificationToken] = ( + await prisma_client.db.litellm_verificationtoken.find_many( + where={"team_id": {"in": data.team_ids}} + ) ) if keys_to_delete: @@ -2991,11 +3000,11 @@ async def team_info( ) try: - team_info: Optional[ - BaseModel - ] = await prisma_client.db.litellm_teamtable.find_unique( - where={"team_id": team_id}, - include={"object_permission": True}, + team_info: Optional[BaseModel] = ( + await prisma_client.db.litellm_teamtable.find_unique( + where={"team_id": team_id}, + include={"object_permission": True}, + ) ) if team_info is None: raise Exception @@ -3316,7 +3325,7 @@ async def _build_team_list_where_conditions( user_object_correct_type = await get_user_object( user_id=user_id, prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, + user_api_key_cache=user_api_key_cache, # type: ignore[arg-type] user_id_upsert=False, proxy_logging_obj=proxy_logging_obj, ) @@ -3835,9 +3844,7 @@ async def list_team( except Exception as e: team_exception = """Invalid team object for team_id: {}. team_object={}. Error: {} - """.format( - team.team_id, team.model_dump(), str(e) - ) + """.format(team.team_id, team.model_dump(), str(e)) verbose_proxy_logger.exception(team_exception) continue # Sort the responses by team_alias diff --git a/litellm/proxy/management_helpers/utils.py b/litellm/proxy/management_helpers/utils.py index 7d485fdebde..3e42d392077 100644 --- a/litellm/proxy/management_helpers/utils.py +++ b/litellm/proxy/management_helpers/utils.py @@ -2,7 +2,7 @@ ## Helper utils for the management endpoints (keys/users/teams) from datetime import datetime from functools import wraps -from typing import Optional, Tuple +from typing import List, Optional, Tuple from fastapi import HTTPException, Request @@ -148,6 +148,7 @@ async def add_new_member( user_api_key_dict: UserAPIKeyAuth, litellm_proxy_admin_name: str, default_team_budget_id: Optional[str] = None, + allowed_models: Optional[List[str]] = None, ) -> Tuple[LiteLLM_UserTable, Optional[LiteLLM_TeamMembership]]: """ Add a new member to a team @@ -206,17 +207,18 @@ async def add_new_member( }, ) - # Check if trying to set a budget for team member - - if max_budget_in_team is not None: + # Check if trying to set a budget or model scope for team member + if max_budget_in_team is not None or allowed_models is not None: # create a new budget item for this member - response = await prisma_client.db.litellm_budgettable.create( - data={ - "max_budget": max_budget_in_team, - "created_by": user_api_key_dict.user_id or litellm_proxy_admin_name, - "updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name, - } - ) + budget_data: dict = { + "created_by": user_api_key_dict.user_id or litellm_proxy_admin_name, + "updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name, + } + if max_budget_in_team is not None: + budget_data["max_budget"] = max_budget_in_team + if allowed_models is not None: + budget_data["allowed_models"] = allowed_models + response = await prisma_client.db.litellm_budgettable.create(data=budget_data) _budget_id = response.budget_id else: diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index fce95465b55..03760695f30 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -17,8 +17,9 @@ model LiteLLM_BudgetTable { tpm_limit BigInt? rpm_limit BigInt? model_max_budget Json? - budget_duration String? + budget_duration String? budget_reset_at DateTime? + allowed_models String[] @default([]) // per-member model scope; empty = inherit team models created_at DateTime @default(now()) @map("created_at") created_by String updated_at DateTime @default(now()) @updatedAt @map("updated_at") @@ -140,6 +141,7 @@ model LiteLLM_TeamTable { team_member_permissions String[] @default([]) access_group_ids String[] @default([]) policies String[] @default([]) + default_team_member_models String[] @default([]) // default allowed_models for newly added team members; empty = no per-member restriction model_id Int? @unique // id for LiteLLM_ModelTable -> stores team-level model aliases allow_team_guardrail_config Boolean @default(false) // if true, team admin can configure guardrails for this team litellm_organization_table LiteLLM_OrganizationTable? @relation(fields: [organization_id], references: [organization_id]) diff --git a/schema.prisma b/schema.prisma index fce95465b55..03760695f30 100644 --- a/schema.prisma +++ b/schema.prisma @@ -17,8 +17,9 @@ model LiteLLM_BudgetTable { tpm_limit BigInt? rpm_limit BigInt? model_max_budget Json? - budget_duration String? + budget_duration String? budget_reset_at DateTime? + allowed_models String[] @default([]) // per-member model scope; empty = inherit team models created_at DateTime @default(now()) @map("created_at") created_by String updated_at DateTime @default(now()) @updatedAt @map("updated_at") @@ -140,6 +141,7 @@ model LiteLLM_TeamTable { team_member_permissions String[] @default([]) access_group_ids String[] @default([]) policies String[] @default([]) + default_team_member_models String[] @default([]) // default allowed_models for newly added team members; empty = no per-member restriction model_id Int? @unique // id for LiteLLM_ModelTable -> stores team-level model aliases allow_team_guardrail_config Boolean @default(false) // if true, team admin can configure guardrails for this team litellm_organization_table LiteLLM_OrganizationTable? @relation(fields: [organization_id], references: [organization_id]) diff --git a/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py b/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py index 5ce291aa165..a4e621b0bbd 100644 --- a/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py +++ b/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py @@ -29,25 +29,30 @@ def test_govcloud_cross_region_inference_prefix(): Test that GovCloud models with cross-region inference prefix (us-gov.) are parsed correctly """ bedrock_model_info = BedrockModelInfo - + # Test us-gov prefix is stripped correctly for Claude models base_model = bedrock_model_info.get_base_model( model="bedrock/us-gov.anthropic.claude-haiku-4-5-20251001-v1:0" ) +<<<<<<< worktree-rustling-wishing-kite + assert base_model == "anthropic.claude-3-5-sonnet-20240620-v1:0" + +======= assert base_model == "anthropic.claude-haiku-4-5-20251001-v1:0" +>>>>>>> main # Test us-gov prefix is stripped correctly for different Claude versions base_model = bedrock_model_info.get_base_model( model="bedrock/us-gov.anthropic.claude-sonnet-4-5-20250929-v1:0" ) assert base_model == "anthropic.claude-sonnet-4-5-20250929-v1:0" - + # Test us-gov prefix is stripped correctly for Haiku models base_model = bedrock_model_info.get_base_model( model="bedrock/us-gov.anthropic.claude-3-haiku-20240307-v1:0" ) assert base_model == "anthropic.claude-3-haiku-20240307-v1:0" - + # Test us-gov prefix is stripped correctly for Meta models base_model = bedrock_model_info.get_base_model( model="bedrock/us-gov.meta.llama3-8b-instruct-v1:0" @@ -55,3 +60,35 @@ def test_govcloud_cross_region_inference_prefix(): assert base_model == "meta.llama3-8b-instruct-v1:0" +def test_context_window_suffix_stripped_for_cost_lookup(): + """ + Test that [1m], [200k] etc. context window suffixes are stripped from + Bedrock model names before cost lookup. + + Models configured like `bedrock/us.anthropic.claude-opus-4-6-v1[1m]` + should resolve to the base model name so pricing can be found. + """ + from litellm.llms.bedrock.common_utils import get_bedrock_base_model + + assert ( + get_bedrock_base_model("us.anthropic.claude-opus-4-6-v1[1m]") + == "anthropic.claude-opus-4-6-v1" + ) + assert ( + get_bedrock_base_model("us.anthropic.claude-sonnet-4-6[1m]") + == "anthropic.claude-sonnet-4-6" + ) + assert ( + get_bedrock_base_model("global.anthropic.claude-opus-4-5-20251101-v1:0[1m]") + == "anthropic.claude-opus-4-5-20251101-v1:0" + ) + # Ensure models without suffix are unaffected + assert ( + get_bedrock_base_model("us.anthropic.claude-opus-4-6-v1") + == "anthropic.claude-opus-4-6-v1" + ) + # Ensure :51k throughput suffix still works + assert ( + get_bedrock_base_model("anthropic.claude-3-5-sonnet-20241022-v2:0:51k") + == "anthropic.claude-3-5-sonnet-20241022-v2:0" + ) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/teams/components/modals/CreateTeamModal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/teams/components/modals/CreateTeamModal.tsx index 0aa42b69a04..b4f99f5ca8b 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/teams/components/modals/CreateTeamModal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/teams/components/modals/CreateTeamModal.tsx @@ -391,6 +391,80 @@ const CreateTeamModal = ({ + + + Team Member Settings + + + + Optional defaults applied when members join this team. All fields can be overridden per member. + + prev.models !== cur.models} + > + {({ getFieldValue }) => { + const teamModels: string[] = getFieldValue("models") || []; + const opts = teamModels.length > 0 ? teamModels : modelsToPick; + return ( + + Default Model Access{" "} + + + + + } + name="default_team_member_models" + > + + {opts.map((m) => ( + + {getModelDisplayName(m)} + + ))} + + + ); + }} + + (value ? Number(value) : undefined)} + tooltip="Default spend budget for each member in this team." + > + + + + + + + + + + + + + + @@ -409,7 +483,7 @@ const CreateTeamModal = ({ { if (!mcpAccessGroupsLoaded) { fetchMcpAccessGroups(); @@ -432,35 +506,6 @@ const CreateTeamModal = ({ }} /> - (value ? Number(value) : undefined)} - tooltip="This is the individual budget for a user in the team." - > - - - - - - - - - - - { @@ -3898,6 +3899,9 @@ export const teamMemberUpdateCall = async ( if (formValues.rpm_limit !== undefined && formValues.rpm_limit !== null) { requestBody.rpm_limit = formValues.rpm_limit; } + if (formValues.allowed_models !== undefined) { + requestBody.allowed_models = formValues.allowed_models; + } console.log("Final request body:", requestBody); diff --git a/ui/litellm-dashboard/src/components/team/EditMembership.tsx b/ui/litellm-dashboard/src/components/team/EditMembership.tsx index e6d9ccad60b..cc655ba52b9 100644 --- a/ui/litellm-dashboard/src/components/team/EditMembership.tsx +++ b/ui/litellm-dashboard/src/components/team/EditMembership.tsx @@ -21,7 +21,7 @@ interface ModalConfig { additionalFields?: Array<{ name: string; label: string | React.ReactNode; - type: "input" | "select" | "numerical"; + type: "input" | "select" | "numerical" | "multi-select"; options?: Array<{ label: string; value: string }>; rules?: any[]; step?: number; @@ -65,6 +65,8 @@ const MemberModal = ({ max_budget_in_team: (initialData as any).max_budget_in_team || null, tpm_limit: (initialData as any).tpm_limit || null, rpm_limit: (initialData as any).rpm_limit || null, + // Keep array values for multi-select fields + allowed_models: (initialData as any).allowed_models || [], }; console.log("Setting form values:", formValues); form.setFieldsValue(formValues); @@ -115,7 +117,7 @@ const MemberModal = ({ const renderField = (field: { name: string; label: string | React.ReactNode; - type: "input" | "select" | "numerical"; + type: "input" | "select" | "numerical" | "multi-select"; options?: Array<{ label: string; value: string }>; rules?: any[]; step?: number; @@ -144,6 +146,15 @@ const MemberModal = ({ ))} ); + case "multi-select": + return ( + - - - - - - form.setFieldValue("team_member_budget_duration", value)} - value={form.getFieldValue("team_member_budget_duration")} - /> - - - - - - - - - - - - - + + + Team Member Settings + + + + Optional defaults applied when members join this team. All fields can be overridden per member. + + + Default Model Access{" "} + + + + + } + name="default_team_member_models" + > + prev.models !== cur.models}> + {({ getFieldValue }) => { + const teamModels = getFieldValue("models") || info.models || []; + return ( + @@ -1318,6 +1362,18 @@ const TeamInfoView: React.FC = ({ ))} + {info.default_team_member_models && info.default_team_member_models.length > 0 && ( +
+ Default Member Models +
+ {info.default_team_member_models.map((model, index) => ( + + {model} + + ))} +
+
+ )}
Rate Limits
TPM: {info.tpm_limit || "Unlimited"}
@@ -1483,6 +1539,20 @@ const TeamInfoView: React.FC = ({ min: 0, placeholder: "Requests per minute limit for this member in this team", }, + { + name: "allowed_models", + label: ( + + Allowed Models{" "} + + + + + ), + type: "multi-select" as const, + options: (info.models || []).map((m: string) => ({ label: m, value: m })), + placeholder: "Leave empty to inherit all team models", + }, ], }} /> diff --git a/ui/litellm-dashboard/src/components/team/TeamMemberTab.tsx b/ui/litellm-dashboard/src/components/team/TeamMemberTab.tsx index 652d1dbcd91..e880aa49f65 100644 --- a/ui/litellm-dashboard/src/components/team/TeamMemberTab.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamMemberTab.tsx @@ -82,7 +82,45 @@ export default function TeamMemberTab({ const isUserTeamAdmin = isUserTeamAdminForSingleTeam(teamData.team_info.members_with_roles, userId || ""); const isProxyAdmin = isProxyAdminRole(userRole || ""); + const getUserAllowedModels = (userId: string | null): string[] | null => { + if (!userId) return null; + const membership = teamData.team_memberships.find((tm) => tm.user_id === userId); + const models = membership?.litellm_budget_table?.allowed_models; + return models && models.length > 0 ? models : null; + }; + const extraColumns: ColumnsType = [ + { + title: ( + + Model Scope + + + + + ), + key: "model_scope", + render: (_: unknown, record: Member) => { + const models = getUserAllowedModels(record.user_id); + if (!models) { + return (all team models); + } + const displayed = models.slice(0, 2); + const remaining = models.length - displayed.length; + return ( + + {displayed.map((m) => ( + {m} + ))} + {remaining > 0 && ( + + +{remaining} more + + )} + + ); + }, + }, { title: ( @@ -138,6 +176,7 @@ export default function TeamMemberTab({ max_budget_in_team: membership?.litellm_budget_table?.max_budget || null, tpm_limit: membership?.litellm_budget_table?.tpm_limit || null, rpm_limit: membership?.litellm_budget_table?.rpm_limit || null, + allowed_models: membership?.litellm_budget_table?.allowed_models || [], }; setSelectedEditMember(enhancedMember); setIsEditMemberModalVisible(true); From 0afffe436686da1f89a5b83b07bb14e1bf0a5bb7 Mon Sep 17 00:00:00 2001 From: ishaan-berri <155045088+ishaan-berri@users.noreply.github.com> Date: Mon, 6 Apr 2026 14:02:04 -0700 Subject: [PATCH 013/420] feat: multiple concurrent budget windows per API key and team (#24883) (#25109) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat: multiple concurrent budget windows per API key and team (#24883) * feat(proxy): add BudgetLimitEntry type and wire budget_limits into key/team models * feat(schema): add budget_limits Json column to VerificationToken and TeamTable * feat(migrations): add migration for budget_limits column on keys and teams * feat(keys): initialize budget_limits windows with reset_at on key create/update * feat(teams): initialize budget_limits windows with reset_at on team create/update * feat(auth): add _virtual_key_multi_budget_check and _team_multi_budget_check * feat(auth): call multi-budget checks from common_checks for keys and teams * feat(proxy): increment per-window Redis spend counters after each request * feat(budget): reset individual budget windows on schedule via reset_budget_job * feat(ui): add hourly option to BudgetDurationDropdown * feat(ui): add budget_limits field to KeyResponse type * feat(ui): add Budget Windows editor to key edit view * feat(ui): add Budget Windows editor to create key form * fix(proxy): strip budget_limits=None before Prisma upsert to fix login 500 Prisma rejects nullable JSON fields (Json? without @default) when passed as Python None — it needs the field omitted entirely so the DB stores NULL via the column's nullable constraint. This was breaking /v2/login because the UI session key creation path hit the upsert with budget_limits=None. * ui(key-edit): use antd InputNumber+Button for budget windows, add reset hints * ui(create-key): use antd InputNumber+Button for budget windows, add reset hints * docs(users): add multiple budget windows section with API + dashboard walkthrough * fix: BudgetExceededError returns HTTP 429 instead of 400 - Add status_code=429 to BudgetExceededError class - auth_exception_handler hardcoded code=400 → code=429 * fix: no-op else branch in multi-budget auth checks causes KeyError - BudgetLimitEntry objects must be coerced via model_dump() not left as-is - Move _virtual_key_multi_budget_check into common_checks (was asymmetric with _team_multi_budget_check which already lived there) * fix: len() on JSON string returns char count not window count Guard with isinstance check + json.loads() before iterating per-window Redis counters in increment_spend_counters * fix: silent except:pass hides Redis reset failures in reset_budget_windows Log Redis counter reset failures as warnings so they are observable * test: add unit tests for multi-budget window enforcement 5 tests covering: no budget_limits passes, under budget passes, over hourly window raises 429, over monthly window raises 429, BudgetLimitEntry objects coerced without KeyError * fix: key per-window counters stable across reorders (duration key, not index) * fix: team+key per-window spend increments use duration key, not index * fix: budget window reset uses duration key; log failures instead of swallowing * refactor: extract BudgetWindowsEditor to shared component * refactor: key_edit_view imports BudgetWindowsEditor from shared component * refactor: create_key_button imports BudgetWindowsEditor from shared component --------- Co-authored-by: Ishaan Jaffer * fix(reset_budget_job): extract _reset_expired_window helper to fix PLR0915 too many statements * feat(skills): Skills Registry & Hub — register skills, browse in AI Hub, public skill hub (#25118) * feat(skills): add domain and namespace fields to plugin types * feat(skills): store and return domain/namespace inside manifest_json * feat(skills): add /public/skill_hub endpoint for unauthenticated access * feat(skills): whitelist /public/skill_hub from auth requirements * feat(skills): add domain, namespace to Plugin and RegisterPluginRequest types * feat(skills): smart URL parser — paste github URL, auto-detect source type and name * feat(skills): replace enable toggle with Public badge, make rows clickable * feat(skills): add skill detail view with Overview and How to Use tabs * feat(skills): add MakeSkillPublicForm modal for publishing skills to the hub * feat(skills): rename panel to Skills, wire in skill detail view on row click * feat(skills): add skill hub table columns — name, description, domain, source, status * feat(skills): add SkillHubDashboard with stats row, domain dropdown filter, and table * feat(skills): add Skill Hub tab to AI Hub with Select Skills to Make Public button * feat(skills): move Skills to top-level nav item directly under MCP Servers * feat(skills): add skillHubPublicCall and NEXT_PUBLIC_BASE_URL support * feat(skills): add Skill Hub tab to public AI Hub page * feat(skills): add skills page routing in main app router * feat(skills): add /skills page route * chore: update package-lock after npm install * docs(skills): add Skills Gateway doc page with mermaid architecture diagram * docs(skills): add Skills Gateway to sidebar under Agent & MCP Gateway * docs(skills): add loom walkthrough video to Skills Gateway doc * chore: fixes --------- Co-authored-by: Ishaan Jaffer Co-authored-by: Yuneng Jiang --- docs/my-website/docs/proxy/users.md | 61 +++ docs/my-website/docs/skills_gateway.md | 111 ++++++ docs/my-website/sidebars.js | 7 + .../migration.sql | 5 + .../litellm_proxy_extras/schema.prisma | 2 + litellm/exceptions.py | 3 +- litellm/proxy/_types.py | 20 + .../claude_code_marketplace.py | 6 + litellm/proxy/auth/auth_checks.py | 81 ++++ litellm/proxy/auth/auth_exception_handler.py | 2 +- litellm/proxy/auth/user_api_key_auth.py | 12 +- .../proxy/common_utils/reset_budget_job.py | 109 +++++- .../key_management_endpoints.py | 65 +++- .../management_endpoints/team_endpoints.py | 25 ++ litellm/proxy/proxy_server.py | 40 ++ .../public_endpoints/public_endpoints.py | 43 ++ litellm/proxy/schema.prisma | 2 + litellm/proxy/utils.py | 22 +- litellm/types/proxy/claude_code_endpoints.py | 4 + .../proxy/auth/test_multi_budget_windows.py | 137 +++++++ ui/litellm-dashboard/package-lock.json | 82 +--- .../src/app/(dashboard)/skills/page.tsx | 17 + ui/litellm-dashboard/src/app/page.tsx | 2 +- .../src/components/AIHub/ModelHubTable.tsx | 63 ++- .../components/AIHub/SkillHubDashboard.tsx | 139 +++++++ .../src/components/claude_code_plugins.tsx | 115 +++--- .../MakeSkillPublicForm.tsx | 249 ++++++++++++ .../claude_code_plugins/add_plugin_form.tsx | 331 ++++++++-------- .../claude_code_plugins/plugin_table.tsx | 71 +--- .../claude_code_plugins/skill_detail.tsx | 366 ++++++++++++++++++ .../components/claude_code_plugins/types.ts | 14 +- .../budget_duration_dropdown.tsx | 2 + .../key_team_helpers/BudgetWindowsEditor.tsx | 84 ++++ .../components/key_team_helpers/key_list.tsx | 1 + .../src/components/leftnav.tsx | 15 +- .../src/components/networking.tsx | 26 +- .../organisms/create_key_button.tsx | 30 +- .../src/components/public_model_hub.tsx | 27 ++ .../components/skill_hub_table_columns.tsx | 114 ++++++ .../components/templates/key_edit_view.tsx | 27 +- 40 files changed, 2103 insertions(+), 429 deletions(-) create mode 100644 docs/my-website/docs/skills_gateway.md create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260401000000_add_budget_limits/migration.sql create mode 100644 tests/test_litellm/proxy/auth/test_multi_budget_windows.py create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/skills/page.tsx create mode 100644 ui/litellm-dashboard/src/components/AIHub/SkillHubDashboard.tsx create mode 100644 ui/litellm-dashboard/src/components/claude_code_plugins/MakeSkillPublicForm.tsx create mode 100644 ui/litellm-dashboard/src/components/claude_code_plugins/skill_detail.tsx create mode 100644 ui/litellm-dashboard/src/components/key_team_helpers/BudgetWindowsEditor.tsx create mode 100644 ui/litellm-dashboard/src/components/skill_hub_table_columns.tsx diff --git a/docs/my-website/docs/proxy/users.md b/docs/my-website/docs/proxy/users.md index 88a7a0f1e07..0e36e84c208 100644 --- a/docs/my-website/docs/proxy/users.md +++ b/docs/my-website/docs/proxy/users.md @@ -333,6 +333,67 @@ curl 'http://0.0.0.0:4000/key/generate' \ }' ``` +#### **Set multiple budget windows on a key** + +Apply multiple concurrent budget limits at different time scales on the same key — for example, cap a key at **$10/day** AND **$100/month**. + +**When is this useful?** + +A single `budget_duration` window can't prevent a bad day from burning your entire month. Multiple budget windows let you: + +- Block a runaway usage spike within the day while still allowing normal monthly spend. +- Give Claude Code rollouts a daily guardrail (`24h`) and a monthly ceiling (`30d`) so a single heavy session doesn't exhaust the whole month. +- Layer fine-grained hourly limits for bursty workloads on top of a weekly cap. + +:::info + +See [User Budget docs](https://docs.litellm.ai/docs/proxy/users) for more on how budgets work across keys, teams, and users. + +::: + +**Via API** + +Pass `budget_limits` as a list of `{budget_duration, max_budget}` objects: + +```bash +curl 'http://0.0.0.0:4000/key/generate' \ +--header 'Authorization: Bearer ' \ +--header 'Content-Type: application/json' \ +--data-raw '{ + "budget_limits": [ + {"budget_duration": "24h", "max_budget": 10}, + {"budget_duration": "30d", "max_budget": 100} + ] +}' +``` + +Each window is tracked independently and resets on its own schedule: + +| `budget_duration` | Resets | +|---|---| +| `1h` | Every hour | +| `24h` | Daily at midnight UTC | +| `7d` | Every Sunday at midnight UTC | +| `30d` | 1st of every month at midnight UTC | + +**Via Dashboard** + +Open **Virtual Keys → Create Key → Optional Settings → Budget Windows**. + +![Step 1 - open key settings](https://colony-recorder.s3.amazonaws.com/files/2026-04-01/18930ba5-67c0-4031-afc0-57f37b4e59e4/ascreenshot_ef79d8a000bb41cdacf1bd9827732ee8_text_export.jpeg) + +Click **+ Add Budget Window** to add a row, choose the period from the dropdown, and enter the spend cap. + +![Step 2 - add a window](https://colony-recorder.s3.amazonaws.com/files/2026-04-01/5ae8c0b3-2d03-41ad-a63c-47b20c350dfe/ascreenshot_1a7dc6c7d65544f38fd8a65604674f22_text_export.jpeg) + +Add a second row for a different time period (e.g. monthly $100 on top of a daily $10). + +![Step 3 - add second window](https://colony-recorder.s3.amazonaws.com/files/2026-04-01/cbded3a7-1086-4e20-8f0f-de154b76146c/ascreenshot_c51c18752c3b4f8b976d28799b2638b6_text_export.jpeg) + +Each window shows the reset schedule below the input so it's always clear when spend resets. + +![Step 4 - reset hints](https://colony-recorder.s3.amazonaws.com/files/2026-04-01/8754f121-1640-4892-9dd0-fd4a870418bf/ascreenshot_8079eb0df2194e8f99e5258ba4b3c082_text_export.jpeg) + ### ✨ Virtual Key (Model Specific) diff --git a/docs/my-website/docs/skills_gateway.md b/docs/my-website/docs/skills_gateway.md new file mode 100644 index 00000000000..d0eb8107579 --- /dev/null +++ b/docs/my-website/docs/skills_gateway.md @@ -0,0 +1,111 @@ +# Skills Gateway + + + +LiteLLM acts as a **Skills Registry** — a central place to register, manage, and discover Claude Code skills across your organization. Teams can publish skills once and have agents and developers find them through a single hub. + +## How it works + +```mermaid +graph TD + Dev["👨‍💻 Developer
registers a skill
(GitHub URL or subdir)"] -->|POST /claude-code/plugins| Proxy["LiteLLM Proxy
(Skills Registry)"] + + Admin["🔑 Admin
publishes skill
(marks as public)"] -->|enable via UI or API| Proxy + + Proxy -->|GET /public/skill_hub| SkillHub["🗂️ Skill Hub
(AI Hub → Skill Hub tab)"] + Proxy -->|GET /claude-code/marketplace.json| Marketplace["📦 Claude Code
Marketplace endpoint"] + + SkillHub --> Human["🧑 Human
browses & discovers skills
in AI Hub UI"] + Marketplace --> Agent["🤖 Agent / Claude Code
installs skill with
/plugin marketplace add <name>"] + + style Proxy fill:#1a73e8,color:#fff + style SkillHub fill:#e8f0fe,color:#1a73e8 + style Marketplace fill:#e8f0fe,color:#1a73e8 +``` + +## Quick start + +### 1. Register a skill + +Paste any GitHub URL into the Skills UI — LiteLLM auto-detects the source type and skill name. + +```bash +curl -X POST https://your-proxy/claude-code/plugins \ + -H "Authorization: Bearer $LITELLM_KEY" \ + -H "Content-Type: application/json" \ + -d '{ + "name": "grill-me", + "source": { + "source": "git-subdir", + "url": "https://github.com/mattpocock/skills", + "path": "grill-me" + }, + "description": "Interview skill for relentless questioning", + "domain": "Productivity", + "namespace": "interviews" + }' +``` + +Skills nested in subdirectories (e.g. `github.com/org/repo/tree/main/skill-name`) are supported — LiteLLM parses the URL automatically in the UI. + +### 2. Publish to hub + +In the Admin UI: **AI Hub → Skill Hub → Select Skills to Make Public**. + +Or via API: + +```bash +curl -X POST https://your-proxy/claude-code/plugins/grill-me/enable \ + -H "Authorization: Bearer $LITELLM_KEY" +``` + +### 3. Browse the hub + +Public skills appear at: +- **Admin UI**: AI Hub → Skill Hub tab +- **Public page**: `/ui/model_hub` → Skill Hub tab (no login required) +- **API**: `GET /public/skill_hub` + +### 4. Install in Claude Code + +Point Claude Code at your proxy marketplace once: + +```json title="~/.claude/settings.json" +{ + "extraKnownMarketplaces": { + "my-org": { + "source": "url", + "url": "https://your-proxy/claude-code/marketplace.json" + } + } +} +``` + +Then install any skill: + +``` +/plugin marketplace add grill-me +``` + +## Skill fields + +| Field | Description | +|-------|-------------| +| `name` | Unique skill identifier (used in `/plugin marketplace add`) | +| `source` | Git source — `github`, `url`, or `git-subdir` | +| `description` | Short description shown in the hub | +| `domain` | Category for grouping (e.g. `Engineering`, `Productivity`) | +| `namespace` | Subcategory within a domain (e.g. `quality`, `meetings`) | +| `keywords` | Tags for search and filtering | +| `version` | Semver string | + +## API reference + +| Endpoint | Auth | Description | +|----------|------|-------------| +| `POST /claude-code/plugins` | Required | Register a skill | +| `GET /claude-code/plugins` | Required | List all skills (admin) | +| `POST /claude-code/plugins/{name}/enable` | Required | Publish a skill | +| `POST /claude-code/plugins/{name}/disable` | Required | Unpublish a skill | +| `GET /public/skill_hub` | None | List public skills | +| `GET /claude-code/marketplace.json` | None | Claude Code marketplace manifest | diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index ab8f257c7d5..bb863a0e771 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -333,6 +333,13 @@ const sidebars = { }, ], }, + { + type: "category", + label: "Skills Gateway", + items: [ + "skills_gateway", + ], + }, ], }, { diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260401000000_add_budget_limits/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260401000000_add_budget_limits/migration.sql new file mode 100644 index 00000000000..fdd65543a3a --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260401000000_add_budget_limits/migration.sql @@ -0,0 +1,5 @@ +-- AlterTable: add budget_limits column to LiteLLM_VerificationToken +ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN IF NOT EXISTS "budget_limits" JSONB; + +-- AlterTable: add budget_limits column to LiteLLM_TeamTable +ALTER TABLE "LiteLLM_TeamTable" ADD COLUMN IF NOT EXISTS "budget_limits" JSONB; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 03760695f30..d570b1dc351 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -141,6 +141,7 @@ model LiteLLM_TeamTable { team_member_permissions String[] @default([]) access_group_ids String[] @default([]) policies String[] @default([]) + budget_limits Json? // multiple concurrent budget windows [{budget_duration, max_budget, reset_at}] default_team_member_models String[] @default([]) // default allowed_models for newly added team members; empty = no per-member restriction model_id Int? @unique // id for LiteLLM_ModelTable -> stores team-level model aliases allow_team_guardrail_config Boolean @default(false) // if true, team admin can configure guardrails for this team @@ -402,6 +403,7 @@ model LiteLLM_VerificationToken { rotation_interval String? // How often to rotate (e.g., "30d", "90d") last_rotation_at DateTime? // When this key was last rotated key_rotation_at DateTime? // When this key should next be rotated + budget_limits Json? // multiple concurrent budget windows [{budget_duration, max_budget, reset_at}] litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id]) litellm_organization_table LiteLLM_OrganizationTable? @relation(fields: [organization_id], references: [organization_id]) litellm_project_table LiteLLM_ProjectTable? @relation(fields: [project_id], references: [project_id]) diff --git a/litellm/exceptions.py b/litellm/exceptions.py index abdba09dd8d..51810c5643f 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -281,7 +281,7 @@ class Timeout(openai.APITimeoutError): # type: ignore return _message -class PermissionDeniedError(openai.PermissionDeniedError): # type:ignore +class PermissionDeniedError(openai.PermissionDeniedError): # type: ignore def __init__( self, message, @@ -847,6 +847,7 @@ class BudgetExceededError(Exception): ): self.current_cost = current_cost self.max_budget = max_budget + self.status_code = 429 message = ( message or f"Budget has been exceeded! Current cost: {current_cost}, Max budget: {max_budget}" diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 452ba7f7995..12b0c542589 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -596,6 +596,7 @@ class LiteLLMRoutes(enum.Enum): "/public/model_hub", "/public/agent_hub", "/public/mcp_hub", + "/public/skill_hub", "/public/litellm_model_cost_map", ] ) @@ -868,6 +869,14 @@ class LiteLLM_ObjectPermissionBase(LiteLLMPydanticObjectBase): models: Optional[List[str]] = None +class BudgetLimitEntry(LiteLLMPydanticObjectBase): + """A single budget window with its own limit and independent reset schedule.""" + + budget_duration: str # e.g. "24h", "7d", "30d" + max_budget: float # max spend in USD for this window + reset_at: Optional[datetime] = None # populated at creation/reset time + + class GenerateRequestBase(LiteLLMPydanticObjectBase): """ Overlapping schema between key and user generate/update requests @@ -887,6 +896,9 @@ class GenerateRequestBase(LiteLLMPydanticObjectBase): rpm_limit: Optional[int] = None budget_duration: Optional[str] = None + budget_limits: Optional[List[BudgetLimitEntry]] = ( + None # multiple concurrent budget windows + ) allowed_cache_controls: Optional[list] = [] config: Optional[dict] = {} permissions: Optional[dict] = {} @@ -991,6 +1003,7 @@ class GenerateKeyResponse(KeyRequestBase): "permissions", "model_max_budget", "router_settings", + "budget_limits", ] for field in dict_fields: value = values.get(field) @@ -1644,6 +1657,9 @@ class TeamBase(LiteLLMPydanticObjectBase): max_budget: Optional[float] = None soft_budget: Optional[float] = None budget_duration: Optional[str] = None + budget_limits: Optional[List[BudgetLimitEntry]] = ( + None # multiple concurrent budget windows + ) models: list = [] blocked: bool = False @@ -1745,6 +1761,8 @@ class UpdateTeamRequest(LiteLLMPydanticObjectBase): enforced_file_expires_after: Optional[dict] = None router_settings: Optional[dict] = None access_group_ids: Optional[List[str]] = None + budget_limits: Optional[List[BudgetLimitEntry]] = ( + None # multiple concurrent budget windows default_team_member_models: Optional[List[str]] = ( None # default allowed_models seeded onto new team members ) @@ -1893,6 +1911,7 @@ class LiteLLM_TeamTable(TeamBase): "model_max_budget", "model_aliases", "router_settings", + "budget_limits", ] if isinstance(values, BaseModel): @@ -2378,6 +2397,7 @@ class LiteLLM_VerificationToken(LiteLLMPydanticObjectBase): last_rotation_at: Optional[datetime] = None # When this key was last rotated key_rotation_at: Optional[datetime] = None # When this key should next be rotated router_settings: Optional[dict] = None + budget_limits: Optional[List[dict]] = None # multiple concurrent budget windows model_config = ConfigDict(protected_namespaces=()) diff --git a/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py b/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py index abd3ce5661d..78cce865a21 100644 --- a/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py +++ b/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py @@ -260,6 +260,10 @@ async def register_plugin( manifest["keywords"] = request.keywords if request.category: manifest["category"] = request.category + if request.domain: + manifest["domain"] = request.domain + if request.namespace: + manifest["namespace"] = request.namespace # Check if plugin exists existing = await prisma_client.db.litellm_claudecodeplugintable.find_unique( @@ -362,6 +366,8 @@ async def list_plugins( homepage=manifest.get("homepage"), keywords=manifest.get("keywords"), category=manifest.get("category"), + domain=manifest.get("domain"), + namespace=manifest.get("namespace"), enabled=p.enabled, created_at=p.created_at.isoformat() if p.created_at else None, updated_at=p.updated_at.isoformat() if p.updated_at else None, diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index bf0de3d3ca8..8e0cf62e56f 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -534,6 +534,17 @@ async def common_checks( # noqa: PLR0915 valid_token=valid_token, ) + # 3.1. Multi-window budget check for team + with tracer.trace("litellm.proxy.auth.common_checks.team_multi_budget_check"): + await _team_multi_budget_check(team_object=team_object) + + # 3.2. Multi-window budget check for key + with tracer.trace( + "litellm.proxy.auth.common_checks.virtual_key_multi_budget_check" + ): + if valid_token is not None: + await _virtual_key_multi_budget_check(valid_token=valid_token) + # 3.0.5. If team is over soft budget (alert only, doesn't block) with tracer.trace("litellm.proxy.auth.common_checks.team_soft_budget_check"): await _team_soft_budget_check( @@ -2982,6 +2993,43 @@ async def _virtual_key_max_budget_check( ) +async def _virtual_key_multi_budget_check( + valid_token: UserAPIKeyAuth, +): + """ + Raises BudgetExceededError if any budget window in valid_token.budget_limits is exceeded. + + Each window has its own Redis counter keyed by spend:key:{token}:window:{budget_duration}. + Using budget_duration (not list index) keeps counters stable when windows are reordered + or removed during a key update. + + Note: counters are not seeded from DB on Redis cold-start. After a Redis flush, + per-window spend resets to zero within the current window period. This is an acceptable + trade-off: the DB stores reset_at timestamps but not per-window accumulated spend. + """ + if not valid_token.budget_limits: + return + + from litellm.proxy.proxy_server import get_current_spend + + for window in valid_token.budget_limits: + w: dict = window if isinstance(window, dict) else window.model_dump() + counter_key = f"spend:key:{valid_token.token}:window:{w['budget_duration']}" + window_spend = await get_current_spend( + counter_key=counter_key, + fallback_spend=0.0, + ) + if window_spend >= w["max_budget"]: + raise litellm.BudgetExceededError( + current_cost=window_spend, + max_budget=w["max_budget"], + message=( + f"ExceededBudget: Key over {w['budget_duration']} budget. " + f"Spend=${window_spend:.4f}, Limit=${w['max_budget']:.2f}" + ), + ) + + async def _virtual_key_soft_budget_check( valid_token: UserAPIKeyAuth, proxy_logging_obj: ProxyLogging, @@ -3223,6 +3271,39 @@ async def _team_max_budget_check( ) +async def _team_multi_budget_check( + team_object: Optional[LiteLLM_TeamTable], +): + """ + Raises BudgetExceededError if any budget window in team_object.budget_limits is exceeded. + + Each window has its own Redis counter keyed by spend:team:{team_id}:window:{budget_duration}. + Using budget_duration (not list index) keeps counters stable when windows are reordered + or removed during a team update. + """ + if team_object is None or not team_object.budget_limits: + return + + from litellm.proxy.proxy_server import get_current_spend + + for window in team_object.budget_limits: + w: dict = window if isinstance(window, dict) else window.model_dump() + counter_key = f"spend:team:{team_object.team_id}:window:{w['budget_duration']}" + window_spend = await get_current_spend( + counter_key=counter_key, + fallback_spend=0.0, + ) + if window_spend >= w["max_budget"]: + raise litellm.BudgetExceededError( + current_cost=window_spend, + max_budget=w["max_budget"], + message=( + f"ExceededBudget: Team={team_object.team_id} over {w['budget_duration']} budget. " + f"Spend=${window_spend:.4f}, Limit=${w['max_budget']:.2f}" + ), + ) + + async def _team_soft_budget_check( team_object: Optional[LiteLLM_TeamTable], valid_token: Optional[UserAPIKeyAuth], diff --git a/litellm/proxy/auth/auth_exception_handler.py b/litellm/proxy/auth/auth_exception_handler.py index 9c306acd2c6..b9b614994c2 100644 --- a/litellm/proxy/auth/auth_exception_handler.py +++ b/litellm/proxy/auth/auth_exception_handler.py @@ -106,7 +106,7 @@ class UserAPIKeyAuthExceptionHandler: message=e.message, type=ProxyErrorTypes.budget_exceeded, param=None, - code=400, + code=429, ) if isinstance(e, HTTPException): raise ProxyException( diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 046c39a9101..1d5b913f9b8 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -984,9 +984,9 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 route=route, ) if _end_user_object is not None: - end_user_params[ - "allowed_model_region" - ] = _end_user_object.allowed_model_region + end_user_params["allowed_model_region"] = ( + _end_user_object.allowed_model_region + ) if _end_user_object.litellm_budget_table is not None: _apply_budget_limits_to_end_user_params( end_user_params=end_user_params, @@ -1540,9 +1540,9 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 if _end_user_object is not None: valid_token_dict.update(end_user_params) - valid_token_dict[ - "end_user_object_permission" - ] = _end_user_object.object_permission + valid_token_dict["end_user_object_permission"] = ( + _end_user_object.object_permission + ) # check if token is from litellm-ui, litellm ui makes keys to allow users to login with sso. These keys can only be used for LiteLLM UI functions # sso/login, ui/login, /key functions and /user functions diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index bcfaed24398..695facd704f 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -2,7 +2,7 @@ import asyncio import json import time from datetime import datetime, timedelta, timezone -from typing import List, Literal, Optional, Union +from typing import Any, List, Literal, Optional, Union from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ( @@ -48,6 +48,9 @@ class ResetBudgetJob: ### RESET ENDUSER (Customer) BUDGET and corresponding Budget duration ### await self.reset_budget_for_litellm_budget_table() + ### RESET MULTI-WINDOW BUDGETS ### + await self.reset_budget_windows() + async def reset_budget_for_litellm_team_members( self, budgets_to_reset: List[LiteLLM_BudgetTableFull] ): @@ -549,6 +552,102 @@ class ResetBudgetJob: ) verbose_proxy_logger.exception("Failed to reset budget for teams: %s", e) + @staticmethod + async def _reset_expired_window( + window: dict, + counter_key: str, + spend_counter_cache: Any, + now: datetime, + ) -> bool: + """Reset a single budget window if expired. Returns True if the window was reset.""" + from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time + + reset_at_str = window.get("reset_at") + if not reset_at_str: + return False + reset_at = datetime.fromisoformat( + reset_at_str.replace("Z", "+00:00") + ).replace(tzinfo=None) + if reset_at > now: + return False + spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=0.0) + if spend_counter_cache.redis_cache is not None: + try: + await spend_counter_cache.redis_cache.async_set_cache( + key=counter_key, value=0.0 + ) + except Exception as redis_err: + verbose_proxy_logger.warning( + "Failed to reset Redis counter %s: %s", counter_key, redis_err + ) + window["reset_at"] = get_budget_reset_time( + budget_duration=window["budget_duration"] + ).isoformat() + return True + + async def reset_budget_windows(self) -> None: + """ + For keys and teams with budget_limits, reset any individual windows where + reset_at <= now. Only the expired windows are reset; other windows are untouched. + """ + from litellm.proxy.proxy_server import spend_counter_cache + + now = datetime.utcnow() + + # --- Keys --- + try: + all_keys = await self.prisma_client.db.litellm_verificationtoken.find_many( + where={"budget_limits": {"not": None}} # type: ignore[arg-type] + ) + for key in all_keys: + raw = key.budget_limits # type: ignore[attr-defined] + if not raw: + continue + windows: list = raw if isinstance(raw, list) else json.loads(raw) + changed = False + for window in windows: + counter_key = f"spend:key:{key.token}:window:{window['budget_duration']}" + if await ResetBudgetJob._reset_expired_window( + window, counter_key, spend_counter_cache, now + ): + changed = True + if changed: + await self.prisma_client.db.litellm_verificationtoken.update( + where={"token": key.token}, + data={"budget_limits": json.dumps(windows)}, # type: ignore[arg-type] + ) + except Exception as e: + verbose_proxy_logger.exception( + "Failed to reset budget windows for keys: %s", e + ) + + # --- Teams --- + try: + all_teams = await self.prisma_client.db.litellm_teamtable.find_many( + where={"budget_limits": {"not": None}} # type: ignore[arg-type] + ) + for team in all_teams: + raw = team.budget_limits # type: ignore[attr-defined] + if not raw: + continue + windows = raw if isinstance(raw, list) else json.loads(raw) + changed = False + for window in windows: + counter_key = f"spend:team:{team.team_id}:window:{window['budget_duration']}" + if await ResetBudgetJob._reset_expired_window( + window, counter_key, spend_counter_cache, now + ): + changed = True + if changed: + await self.prisma_client.db.litellm_teamtable.update( + where={"team_id": team.team_id}, + data={"budget_limits": json.dumps(windows)}, # type: ignore[arg-type] + ) + except Exception as e: + verbose_proxy_logger.exception( + "Failed to reset budget windows for teams: %s", e + ) + @staticmethod async def _reset_budget_common( item: Union[LiteLLM_TeamTable, LiteLLM_UserTable, LiteLLM_VerificationToken], @@ -570,14 +669,14 @@ class ResetBudgetJob: from litellm.proxy.proxy_server import spend_counter_cache counter_key = None - if item_type == "key" and hasattr(item, "token") and item.token is not None: - counter_key = f"spend:key:{item.token}" + if item_type == "key" and hasattr(item, "token") and item.token is not None: # type: ignore[union-attr] + counter_key = f"spend:key:{item.token}" # type: ignore[union-attr] elif ( item_type == "team" and hasattr(item, "team_id") - and item.team_id is not None + and item.team_id is not None # type: ignore[union-attr] ): - counter_key = f"spend:team:{item.team_id}" + counter_key = f"spend:team:{item.team_id}" # type: ignore[union-attr] if counter_key is not None: # Always reset in-memory (local fallback) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 00d8ce182ec..89c72131760 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -744,9 +744,9 @@ async def _common_key_generation_helper( # noqa: PLR0915 request_type="key", **data_json, table_name="key" ) - response[ - "soft_budget" - ] = data.soft_budget # include the user-input soft budget in the response + response["soft_budget"] = ( + data.soft_budget + ) # include the user-input soft budget in the response response = GenerateKeyResponse(**response) @@ -1571,6 +1571,19 @@ async def prepare_key_update_data( non_default_values["budget_reset_at"] = key_reset_at non_default_values["budget_duration"] = budget_duration + if "budget_limits" in non_default_values and non_default_values["budget_limits"]: + from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time + + raw_windows = non_default_values["budget_limits"] + initialized_windows = [] + for window in raw_windows: + w = window if isinstance(window, dict) else window.model_dump() + w["reset_at"] = get_budget_reset_time( + budget_duration=w["budget_duration"] + ).isoformat() + initialized_windows.append(w) + non_default_values["budget_limits"] = json.dumps(initialized_windows) + if "object_permission" in non_default_values: non_default_values = await _handle_update_object_permission( data_json=non_default_values, @@ -2832,6 +2845,7 @@ async def generate_key_helper_fn( # noqa: PLR0915 rotation_interval: Optional[str] = None, router_settings: Optional[dict] = None, access_group_ids: Optional[list] = None, + budget_limits: Optional[list] = None, # multiple concurrent budget windows ): from litellm.proxy.proxy_server import premium_user, prisma_client @@ -2863,6 +2877,18 @@ async def generate_key_helper_fn( # noqa: PLR0915 else: reset_at = get_budget_reset_time(budget_duration=budget_duration) + # Initialize reset_at for each budget window + budget_limits_json: Optional[str] = None + if budget_limits: + initialized_windows = [] + for window in budget_limits: + w = dict(window) if not isinstance(window, dict) else {**window} + w["reset_at"] = get_budget_reset_time( + budget_duration=w["budget_duration"] + ).isoformat() + initialized_windows.append(w) + budget_limits_json = json.dumps(initialized_windows) + aliases_json = json.dumps(aliases) config_json = json.dumps(config) permissions_json = json.dumps(permissions) @@ -2945,6 +2971,7 @@ async def generate_key_helper_fn( # noqa: PLR0915 "organization_id": organization_id, "budget_id": budget_id, "blocked": blocked, + "budget_limits": budget_limits_json, "created_by": created_by, "updated_by": updated_by, "allowed_routes": allowed_routes or [], @@ -3222,10 +3249,10 @@ async def delete_verification_tokens( try: if prisma_client: tokens = [_hash_token_if_needed(token=key) for key in tokens] - _keys_being_deleted: List[ - LiteLLM_VerificationToken - ] = await prisma_client.db.litellm_verificationtoken.find_many( - where={"token": {"in": tokens}} + _keys_being_deleted: List[LiteLLM_VerificationToken] = ( + await prisma_client.db.litellm_verificationtoken.find_many( + where={"token": {"in": tokens}} + ) ) if len(_keys_being_deleted) == 0: @@ -3425,9 +3452,9 @@ async def _rotate_master_key( # noqa: PLR0915 from litellm.proxy.proxy_server import proxy_config try: - models: Optional[ - List - ] = await prisma_client.db.litellm_proxymodeltable.find_many() + models: Optional[List] = ( + await prisma_client.db.litellm_proxymodeltable.find_many() + ) except Exception: models = None # 2. process model table @@ -4067,11 +4094,11 @@ async def validate_key_list_check( param="user_id", code=status.HTTP_403_FORBIDDEN, ) - complete_user_info_db_obj: Optional[ - BaseModel - ] = await prisma_client.db.litellm_usertable.find_unique( - where={"user_id": user_api_key_dict.user_id}, - include={"organization_memberships": True}, + complete_user_info_db_obj: Optional[BaseModel] = ( + await prisma_client.db.litellm_usertable.find_unique( + where={"user_id": user_api_key_dict.user_id}, + include={"organization_memberships": True}, + ) ) if complete_user_info_db_obj is None: @@ -4154,10 +4181,10 @@ async def _fetch_user_team_objects( if complete_user_info is None or not complete_user_info.teams: return [] - teams: Optional[ - List[BaseModel] - ] = await prisma_client.db.litellm_teamtable.find_many( - where={"team_id": {"in": complete_user_info.teams}} + teams: Optional[List[BaseModel]] = ( + await prisma_client.db.litellm_teamtable.find_many( + where={"team_id": {"in": complete_user_info.teams}} + ) ) if teams is None: return [] diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index c980fcb82f8..15b7f9c6461 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -978,6 +978,19 @@ async def new_team( # noqa: PLR0915 budget_duration=complete_team_data.budget_duration, ) + # If budget_limits is set, initialize reset_at for each window + if complete_team_data.budget_limits: + from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time + + initialized_windows = [] + for window in complete_team_data.budget_limits: + w = window if isinstance(window, dict) else window.model_dump() + w["reset_at"] = get_budget_reset_time( + budget_duration=w["budget_duration"] + ).isoformat() + initialized_windows.append(w) + complete_team_data.budget_limits = initialized_windows + ## Add Team Member Budget Table members_with_roles: List[Member] = [] if complete_team_data.members_with_roles is not None: @@ -1593,6 +1606,18 @@ def _set_budget_reset_at(data: UpdateTeamRequest, updated_kv: dict) -> None: reset_at = get_budget_reset_time(budget_duration=data.budget_duration) updated_kv["budget_reset_at"] = reset_at + if data.budget_limits is not None and len(data.budget_limits) > 0: + from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time + + initialized_windows = [] + for window in data.budget_limits: + w = window if isinstance(window, dict) else window.model_dump() + w["reset_at"] = get_budget_reset_time( + budget_duration=w["budget_duration"] + ).isoformat() + initialized_windows.append(w) + updated_kv["budget_limits"] = json.dumps(initialized_windows) + async def handle_update_object_permission( data_json: dict, existing_team_row: LiteLLM_TeamTable diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 715fc44114b..ef4807275e4 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1781,6 +1781,26 @@ async def increment_spend_counters( increment=response_cost, ) + # Increment per-window budget counters for multi-budget keys + key_obj = await user_api_key_cache.async_get_cache(key=hashed_token) + if key_obj is not None: + key_budget_limits = getattr(key_obj, "budget_limits", None) or ( + key_obj.get("budget_limits") if isinstance(key_obj, dict) else None + ) + if isinstance(key_budget_limits, str): + key_budget_limits = json.loads(key_budget_limits) + if isinstance(key_budget_limits, list): + for window in key_budget_limits: + duration = ( + window["budget_duration"] + if isinstance(window, dict) + else window.budget_duration + ) + await spend_counter_cache.async_increment_cache( + key=f"spend:key:{hashed_token}:window:{duration}", + value=response_cost, + ) + if team_id is not None: await _init_and_increment_spend_counter( counter_key=f"spend:team:{team_id}", @@ -1788,6 +1808,26 @@ async def increment_spend_counters( increment=response_cost, ) + # Increment per-window budget counters for multi-budget teams + team_obj = await user_api_key_cache.async_get_cache(key=f"team_id:{team_id}") + if team_obj is not None: + team_budget_limits = getattr(team_obj, "budget_limits", None) or ( + team_obj.get("budget_limits") if isinstance(team_obj, dict) else None + ) + if isinstance(team_budget_limits, str): + team_budget_limits = json.loads(team_budget_limits) + if isinstance(team_budget_limits, list): + for window in team_budget_limits: + duration = ( + window["budget_duration"] + if isinstance(window, dict) + else window.budget_duration + ) + await spend_counter_cache.async_increment_cache( + key=f"spend:team:{team_id}:window:{duration}", + value=response_cost, + ) + if user_id is not None and team_id is not None: await _init_and_increment_spend_counter( counter_key=f"spend:team_member:{user_id}:{team_id}", diff --git a/litellm/proxy/public_endpoints/public_endpoints.py b/litellm/proxy/public_endpoints/public_endpoints.py index eb9ed59055c..10b0a6644e7 100644 --- a/litellm/proxy/public_endpoints/public_endpoints.py +++ b/litellm/proxy/public_endpoints/public_endpoints.py @@ -247,6 +247,49 @@ async def get_mcp_servers(): ] +@router.get( + "/public/skill_hub", + tags=["public", "Claude Code Marketplace"], +) +async def public_skill_hub(): + """Return enabled (public) Claude Code skills — no auth required.""" + from litellm.proxy.anthropic_endpoints.claude_code_endpoints.claude_code_marketplace import ( + _get_prisma_client, + ) + from litellm.types.proxy.claude_code_endpoints import ListPluginsResponse, PluginListItem + + try: + prisma_client = await _get_prisma_client() + plugins = await prisma_client.db.litellm_claudecodeplugintable.find_many( + where={"enabled": True} + ) + items = [] + for plugin in plugins: + raw = plugin.manifest_json or {} + manifest = json.loads(raw) if isinstance(raw, str) else raw + items.append( + PluginListItem( + id=plugin.id, + name=plugin.name, + enabled=plugin.enabled, + created_at=str(plugin.created_at) if plugin.created_at else None, + updated_at=str(plugin.updated_at) if plugin.updated_at else None, + source=manifest.get("source", {}), + description=manifest.get("description"), + version=manifest.get("version"), + category=manifest.get("category"), + keywords=manifest.get("keywords"), + author=manifest.get("author"), + homepage=manifest.get("homepage"), + domain=manifest.get("domain"), + namespace=manifest.get("namespace"), + ) + ) + return ListPluginsResponse(plugins=items, count=len(items)) + except Exception as e: + raise HTTPException(status_code=500, detail=str(e)) + + @router.get( "/public/model_hub/info", tags=["public", "model management"], diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 03760695f30..d570b1dc351 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -141,6 +141,7 @@ model LiteLLM_TeamTable { team_member_permissions String[] @default([]) access_group_ids String[] @default([]) policies String[] @default([]) + budget_limits Json? // multiple concurrent budget windows [{budget_duration, max_budget, reset_at}] default_team_member_models String[] @default([]) // default allowed_models for newly added team members; empty = no per-member restriction model_id Int? @unique // id for LiteLLM_ModelTable -> stores team-level model aliases allow_team_guardrail_config Boolean @default(false) // if true, team admin can configure guardrails for this team @@ -402,6 +403,7 @@ model LiteLLM_VerificationToken { rotation_interval String? // How often to rotate (e.g., "30d", "90d") last_rotation_at DateTime? // When this key was last rotated key_rotation_at DateTime? // When this key should next be rotated + budget_limits Json? // multiple concurrent budget windows [{budget_duration, max_budget, reset_at}] litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id]) litellm_organization_table LiteLLM_OrganizationTable? @relation(fields: [organization_id], references: [organization_id]) litellm_project_table LiteLLM_ProjectTable? @relation(fields: [project_id], references: [project_id]) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index ec98cfd4d1e..8d1ac631910 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1898,9 +1898,9 @@ class ProxyLogging: normalized_call_type = CallTypes.aembedding.value if normalized_call_type is not None: litellm_logging_obj.call_type = normalized_call_type - litellm_logging_obj.model_call_details[ - "call_type" - ] = normalized_call_type + litellm_logging_obj.model_call_details["call_type"] = ( + normalized_call_type + ) # Pass-through endpoints are logged via the callback loop's # async_post_call_failure_hook — skip pre_call and failure handlers. if litellm_logging_obj.call_type == CallTypes.pass_through.value: @@ -2524,8 +2524,7 @@ class PrismaClient: required_view = "LiteLLM_VerificationTokenView" expected_views_str = ", ".join(f"'{view}'" for view in expected_views) pg_schema = os.getenv("DATABASE_SCHEMA", "public") - ret = await self.db.query_raw( - f""" + ret = await self.db.query_raw(f""" WITH existing_views AS ( SELECT viewname FROM pg_views @@ -2537,8 +2536,7 @@ class PrismaClient: (SELECT COUNT(*) FROM existing_views) AS view_count, ARRAY_AGG(viewname) AS view_names FROM existing_views - """ - ) + """) expected_total_views = len(expected_views) if ret[0]["view_count"] == expected_total_views: verbose_proxy_logger.info("All necessary views exist!") @@ -2547,8 +2545,7 @@ class PrismaClient: ## check if required view exists ## if ret[0]["view_names"] and required_view not in ret[0]["view_names"]: await self.health_check() # make sure we can connect to db - await self.db.execute_raw( - """ + await self.db.execute_raw(""" CREATE VIEW "LiteLLM_VerificationTokenView" AS SELECT v.*, @@ -2558,8 +2555,7 @@ class PrismaClient: t.rpm_limit AS team_rpm_limit FROM "LiteLLM_VerificationToken" v LEFT JOIN "LiteLLM_TeamTable" t ON v.team_id = t.team_id; - """ - ) + """) verbose_proxy_logger.info( "LiteLLM_VerificationTokenView Created in DB!" @@ -3142,6 +3138,10 @@ class PrismaClient: hashed_token = self.hash_token(token=token) db_data = self.jsonify_object(data=data) db_data["token"] = hashed_token + # Prisma rejects nullable JSON fields set to None (no default). + # Strip them so the DB stores NULL via the column's nullable constraint. + if db_data.get("budget_limits") is None: + db_data.pop("budget_limits", None) print_verbose( "PrismaClient: Before upsert into litellm_verificationtoken" ) diff --git a/litellm/types/proxy/claude_code_endpoints.py b/litellm/types/proxy/claude_code_endpoints.py index bdf4f122e0e..e5d82a4503b 100644 --- a/litellm/types/proxy/claude_code_endpoints.py +++ b/litellm/types/proxy/claude_code_endpoints.py @@ -49,6 +49,8 @@ class RegisterPluginRequest(BaseModel): homepage: Optional[str] = Field(None, description="Plugin homepage URL") keywords: Optional[List[str]] = Field(None, description="Search keywords") category: Optional[str] = Field(None, description="Plugin category") + domain: Optional[str] = Field(None, description="Skill domain (e.g., 'Productivity')") + namespace: Optional[str] = Field(None, description="Skill namespace within domain (e.g., 'workflows')") class PluginResponse(BaseModel): @@ -82,6 +84,8 @@ class PluginListItem(BaseModel): homepage: Optional[str] = None keywords: Optional[List[str]] = None category: Optional[str] = None + domain: Optional[str] = None + namespace: Optional[str] = None enabled: bool created_at: Optional[str] updated_at: Optional[str] diff --git a/tests/test_litellm/proxy/auth/test_multi_budget_windows.py b/tests/test_litellm/proxy/auth/test_multi_budget_windows.py new file mode 100644 index 00000000000..ed94fca837b --- /dev/null +++ b/tests/test_litellm/proxy/auth/test_multi_budget_windows.py @@ -0,0 +1,137 @@ +""" +Unit tests for multi-budget-window enforcement on API keys. +""" + +from unittest.mock import AsyncMock, patch + +import pytest + +import litellm +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.auth_checks import _virtual_key_multi_budget_check + + +def _make_valid_token(**kwargs) -> UserAPIKeyAuth: + defaults = dict( + token="sk-test-token", + key_name="test", + spend=0.0, + max_budget=None, + budget_limits=[], + ) + defaults.update(kwargs) + return UserAPIKeyAuth(**defaults) + + +@pytest.mark.asyncio +async def test_no_budget_limits_passes(): + """Keys with empty budget_limits should pass without raising.""" + token = _make_valid_token(budget_limits=[]) + # Should not raise + await _virtual_key_multi_budget_check(valid_token=token) + + +@pytest.mark.asyncio +async def test_under_budget_passes(): + """Key with spend under all windows should pass.""" + token = _make_valid_token( + budget_limits=[ + {"budget_duration": "24h", "max_budget": 10.0, "reset_at": None}, + {"budget_duration": "30d", "max_budget": 100.0, "reset_at": None}, + ] + ) + with patch( + "litellm.proxy.proxy_server.get_current_spend", + new_callable=AsyncMock, + return_value=1.0, # well under both windows + ): + await _virtual_key_multi_budget_check(valid_token=token) + + +@pytest.mark.asyncio +async def test_over_first_window_raises(): + """Key exceeding the first (daily) window should raise BudgetExceededError.""" + token = _make_valid_token( + budget_limits=[ + {"budget_duration": "24h", "max_budget": 5.0, "reset_at": None}, + {"budget_duration": "30d", "max_budget": 100.0, "reset_at": None}, + ] + ) + + spend_by_window = [6.0, 6.0] # over daily, under monthly + + call_count = 0 + + async def fake_get_spend(counter_key, fallback_spend): + nonlocal call_count + val = spend_by_window[call_count] + call_count += 1 + return val + + with patch( + "litellm.proxy.proxy_server.get_current_spend", side_effect=fake_get_spend + ): + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _virtual_key_multi_budget_check(valid_token=token) + + err = exc_info.value + assert err.status_code == 429 + assert "24h" in str(err) + assert "Key over" in str(err) + + +@pytest.mark.asyncio +async def test_over_second_window_raises(): + """Key exceeding only the monthly window should raise BudgetExceededError referencing 30d.""" + token = _make_valid_token( + budget_limits=[ + {"budget_duration": "24h", "max_budget": 50.0, "reset_at": None}, + {"budget_duration": "30d", "max_budget": 5.0, "reset_at": None}, + ] + ) + + spend_by_window = [1.0, 10.0] # under daily, over monthly + + call_count = 0 + + async def fake_get_spend(counter_key, fallback_spend): + nonlocal call_count + val = spend_by_window[call_count] + call_count += 1 + return val + + with patch( + "litellm.proxy.proxy_server.get_current_spend", side_effect=fake_get_spend + ): + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _virtual_key_multi_budget_check(valid_token=token) + + err = exc_info.value + assert err.status_code == 429 + assert "30d" in str(err) + + +@pytest.mark.asyncio +async def test_budget_limit_entry_objects_coerced(): + """BudgetLimitEntry Pydantic objects (not dicts) must be handled without KeyError. + + While budget_limits is normally serialized as List[dict], the auth check must + tolerate BudgetLimitEntry objects in case they arrive without prior serialization. + """ + from litellm.proxy._types import BudgetLimitEntry + + token = _make_valid_token(budget_limits=[]) + # Bypass Pydantic validation to simulate BudgetLimitEntry objects reaching the check + object.__setattr__( + token, + "budget_limits", + [BudgetLimitEntry(budget_duration="24h", max_budget=10.0)], + ) + + with patch( + "litellm.proxy.proxy_server.get_current_spend", + new_callable=AsyncMock, + return_value=1.0, + ): + # Should not raise TypeError / KeyError — model_dump() coerces the object + await _virtual_key_multi_budget_check(valid_token=token) diff --git a/ui/litellm-dashboard/package-lock.json b/ui/litellm-dashboard/package-lock.json index 58b0a6b6117..af2edafd831 100644 --- a/ui/litellm-dashboard/package-lock.json +++ b/ui/litellm-dashboard/package-lock.json @@ -91,7 +91,6 @@ "version": "5.2.0", "resolved": "https://registry.npmjs.org/@alloc/quick-lru/-/quick-lru-5.2.0.tgz", "integrity": "sha512-UrcABB+4bUrFABwbluTIBErXwvbsU/V7TZWfmbgJfbkwiBuziS9gxdODUyuiecfdGQ85jglMW6juS3+z5TsKLw==", - "dev": true, "license": "MIT", "engines": { "node": ">=10" @@ -1773,7 +1772,6 @@ "version": "0.3.13", "resolved": "https://registry.npmjs.org/@jridgewell/gen-mapping/-/gen-mapping-0.3.13.tgz", "integrity": "sha512-2kkt/7niJ6MgEPxF0bYdQ6etZaA+fQvDcLKckhy1yIQOzaoKjBBjSj63/aLVjYE3qhRt5dvM+uUyfCg6UKCBbA==", - "dev": true, "license": "MIT", "dependencies": { "@jridgewell/sourcemap-codec": "^1.5.0", @@ -1784,7 +1782,6 @@ "version": "3.1.2", "resolved": "https://registry.npmjs.org/@jridgewell/resolve-uri/-/resolve-uri-3.1.2.tgz", "integrity": "sha512-bRISgCIjP20/tbWSPWMEi54QVPRZExkuD9lJL+UIxUKtwVJA8wW1Trb1jMs1RFXo1CBTNZ/5hpC9QvmKWdopKw==", - "dev": true, "license": "MIT", "engines": { "node": ">=6.0.0" @@ -1794,14 +1791,12 @@ "version": "1.5.5", "resolved": "https://registry.npmjs.org/@jridgewell/sourcemap-codec/-/sourcemap-codec-1.5.5.tgz", "integrity": "sha512-cYQ9310grqxueWbl+WuIUIaiUaDcj7WOq5fVhEljNVgRfOUhY9fy2zTvfoqWsnebh8Sl70VScFbICvJnLKB0Og==", - "dev": true, "license": "MIT" }, "node_modules/@jridgewell/trace-mapping": { "version": "0.3.31", "resolved": "https://registry.npmjs.org/@jridgewell/trace-mapping/-/trace-mapping-0.3.31.tgz", "integrity": "sha512-zzNR+SdQSDJzc8joaeP8QQoCQr8NuYx2dIIytl1QeBEZHJ9uW6hebsrYgbz8hJwUQao3TWCMtmfV8Nu1twOLAw==", - "dev": true, "license": "MIT", "dependencies": { "@jridgewell/resolve-uri": "^3.1.0", @@ -1969,7 +1964,6 @@ "version": "2.1.5", "resolved": "https://registry.npmjs.org/@nodelib/fs.scandir/-/fs.scandir-2.1.5.tgz", "integrity": "sha512-vq24Bq3ym5HEQm2NKCr3yXDwjc7vTsEThRDnkp2DK9p1uqLR+DHurm/NOTo0KG7HYHU7eppKZj3MyqYuMBf62g==", - "dev": true, "license": "MIT", "dependencies": { "@nodelib/fs.stat": "2.0.5", @@ -1983,7 +1977,6 @@ "version": "2.0.5", "resolved": "https://registry.npmjs.org/@nodelib/fs.stat/-/fs.stat-2.0.5.tgz", "integrity": "sha512-RkhPPp2zrqDAQA/2jNhnztcPAlv64XdhIp7a7454A5ovI7Bukxgt7MX7udwAu3zg1DcpPU0rz3VV1SeaqvY4+A==", - "dev": true, "license": "MIT", "engines": { "node": ">= 8" @@ -1993,7 +1986,6 @@ "version": "1.2.8", "resolved": "https://registry.npmjs.org/@nodelib/fs.walk/-/fs.walk-1.2.8.tgz", "integrity": "sha512-oGB+UxlgWcgQkgwo8GcEGwemoTFt3FIO9ababBmaGwXIoBKZ+GTy0pP185beGg7Llih/NSHSV2XAs1lnznocSg==", - "dev": true, "license": "MIT", "dependencies": { "@nodelib/fs.scandir": "2.1.5", @@ -2317,7 +2309,7 @@ "version": "1.58.1", "resolved": "https://registry.npmjs.org/@playwright/test/-/test-1.58.1.tgz", "integrity": "sha512-6LdVIUERWxQMmUSSQi0I53GgCBYgM2RpGngCPY7hSeju+VrKjq3lvs7HpJoPbDiY5QM5EYRtRX5fvrinnMAz3w==", - "dev": true, + "devOptional": true, "license": "Apache-2.0", "dependencies": { "playwright": "1.58.1" @@ -3422,14 +3414,12 @@ "version": "15.7.15", "resolved": "https://registry.npmjs.org/@types/prop-types/-/prop-types-15.7.15.tgz", "integrity": "sha512-F6bEyamV9jKGAFBEmlQnesRPGOQqS2+Uwi0Em15xenOxHaf2hv6L8YCVn3rPdPJOiJfPiCnLIRyvwVaqMY3MIw==", - "dev": true, "license": "MIT" }, "node_modules/@types/react": { "version": "18.2.48", "resolved": "https://registry.npmjs.org/@types/react/-/react-18.2.48.tgz", "integrity": "sha512-qboRCl6Ie70DQQG9hhNREz81jqC1cs9EVNcjQ1AU+jH6NFfSAhVVbrrY/+nSF+Bsk4AOwm9Qa61InvMCyV+H3w==", - "dev": true, "license": "MIT", "dependencies": { "@types/prop-types": "*", @@ -3471,7 +3461,6 @@ "version": "0.26.0", "resolved": "https://registry.npmjs.org/@types/scheduler/-/scheduler-0.26.0.tgz", "integrity": "sha512-WFHp9YUJQ6CKshqoC37iOlHnQSmxNc795UhB26CyBBttrN9svdIrUjl/NjnNmfcwtncN0h/0PPAFWv9ovP8mLA==", - "dev": true, "license": "MIT" }, "node_modules/@types/unist": { @@ -4332,14 +4321,12 @@ "version": "1.3.0", "resolved": "https://registry.npmjs.org/any-promise/-/any-promise-1.3.0.tgz", "integrity": "sha512-7UvmKalWRt1wgjL1RrGxoSJW/0QZFIegpeGvZG9kjp8vrRu55XTHbwnqq2GpXm9uLbcuhxm3IqX9OB4MZR1b2A==", - "dev": true, "license": "MIT" }, "node_modules/anymatch": { "version": "3.1.3", "resolved": "https://registry.npmjs.org/anymatch/-/anymatch-3.1.3.tgz", "integrity": "sha512-KMReFUr0B4t+D+OBkjR3KYqvocp2XaSzO55UcB6mgQMd3KbcE+mWTyvVV7D/zsdEbNnV6acZUutkiHQXvTr1Rw==", - "dev": true, "license": "ISC", "dependencies": { "normalize-path": "^3.0.0", @@ -4353,7 +4340,6 @@ "version": "2.3.1", "resolved": "https://registry.npmjs.org/picomatch/-/picomatch-2.3.1.tgz", "integrity": "sha512-JU3teHTNjmE2VCGFzuY8EXzCDVwEqB2a8fsIvwaStHhAWJEeVd1o1QD80CU6+ZdEXXSLbSsuLwJjkCBWqRQUVA==", - "dev": true, "license": "MIT", "engines": { "node": ">=8.6" @@ -4366,7 +4352,6 @@ "version": "5.0.2", "resolved": "https://registry.npmjs.org/arg/-/arg-5.0.2.tgz", "integrity": "sha512-PYjyFOLKQ9y57JvQ6QLo8dAgNqswh8M1RMJYdQduT6xbWSgK36P/Z/v+p888pM69jMMfS8Xd8F6I1kQ/I9HUGg==", - "dev": true, "license": "MIT" }, "node_modules/argparse": { @@ -4726,7 +4711,6 @@ "version": "2.3.0", "resolved": "https://registry.npmjs.org/binary-extensions/-/binary-extensions-2.3.0.tgz", "integrity": "sha512-Ceh+7ox5qe7LJuLHoY0feh3pHuUDHAcRUeyL2VYghZwfpkNIy/+8Ocg0a3UuSoYzavmylwuLWQOf3hl0jjMMIw==", - "dev": true, "license": "MIT", "engines": { "node": ">=8" @@ -4752,7 +4736,6 @@ "version": "3.0.3", "resolved": "https://registry.npmjs.org/braces/-/braces-3.0.3.tgz", "integrity": "sha512-yQbXgO/OSZVD2IsiLlro+7Hf6Q18EJrKSEsdoMzKePKXct3gvD8oLcOQdIzGupr5Fj+EDe8gO/lxc1BzfMpxvA==", - "dev": true, "license": "MIT", "dependencies": { "fill-range": "^7.1.1" @@ -4868,7 +4851,6 @@ "version": "2.0.1", "resolved": "https://registry.npmjs.org/camelcase-css/-/camelcase-css-2.0.1.tgz", "integrity": "sha512-QOSvevhslijgYwRx6Rv7zKdMF8lbRmx+uQGx2+vDc+KI/eBnsy9kit5aj23AgGu3pa4t9AgwbnXWqS+iOY+2aA==", - "dev": true, "license": "MIT", "engines": { "node": ">= 6" @@ -4992,7 +4974,6 @@ "version": "3.6.0", "resolved": "https://registry.npmjs.org/chokidar/-/chokidar-3.6.0.tgz", "integrity": "sha512-7VT13fmjotKpGipCW9JEQAusEPE+Ei8nl6/g4FBAmIm0GOOLMua9NDDo/DWp0ZAxCr3cPq5ZpBqmPAQgDda2Pw==", - "dev": true, "license": "MIT", "dependencies": { "anymatch": "~3.1.2", @@ -5017,7 +4998,6 @@ "version": "5.1.2", "resolved": "https://registry.npmjs.org/glob-parent/-/glob-parent-5.1.2.tgz", "integrity": "sha512-AOIgSQCepiJYwP3ARnGx+5VnTu2HBYdzbGP45eLw1vr3zB3vZLeyed1sC9hnbcOc9/SrMyM5RPQrkGz4aS9Zow==", - "dev": true, "license": "ISC", "dependencies": { "is-glob": "^4.0.1" @@ -5093,7 +5073,6 @@ "version": "4.1.1", "resolved": "https://registry.npmjs.org/commander/-/commander-4.1.1.tgz", "integrity": "sha512-NOKm8xhkzAjzFx8B2v5OAHT+u5pRQc2UCa2Vq9jYL/31o2wi9mxBA7LIFs3sV5VSC49z6pEhfbMULvShKj26WA==", - "dev": true, "license": "MIT", "engines": { "node": ">= 6" @@ -5154,7 +5133,6 @@ "version": "3.0.0", "resolved": "https://registry.npmjs.org/cssesc/-/cssesc-3.0.0.tgz", "integrity": "sha512-/Tb/JcjK111nNScGob5MNtsntNM1aCNUDipB/TkwZFhyDrrE47SOx/18wF2bbjgc3ZzCSKW1T5nt5EbFoAz/Vg==", - "dev": true, "license": "MIT", "bin": { "cssesc": "bin/cssesc" @@ -5568,14 +5546,12 @@ "version": "1.2.2", "resolved": "https://registry.npmjs.org/didyoumean/-/didyoumean-1.2.2.tgz", "integrity": "sha512-gxtyfqMg7GKyhQmb056K7M3xszy/myH8w+B4RT+QXBQsvAOdc3XymqDDPHx1BgPgsdAA5SIifona89YtRATDzw==", - "dev": true, "license": "Apache-2.0" }, "node_modules/dlv": { "version": "1.1.3", "resolved": "https://registry.npmjs.org/dlv/-/dlv-1.1.3.tgz", "integrity": "sha512-+HlytyjlPKnIG8XuRG8WvmBP8xs8P71y+SKKS6ZXWoEgLuePxtDoUEiH7WkdePWrQ5JBpE6aoVqfZfJUQkjXwA==", - "dev": true, "license": "MIT" }, "node_modules/doctrine": { @@ -6489,7 +6465,6 @@ "version": "1.20.1", "resolved": "https://registry.npmjs.org/fastq/-/fastq-1.20.1.tgz", "integrity": "sha512-GGToxJ/w1x32s/D2EKND7kTil4n8OVk/9mycTc4VDza13lOvpUZTGX3mFSCtV9ksdGBVzvsyAVLM6mHFThxXxw==", - "dev": true, "license": "ISC", "dependencies": { "reusify": "^1.0.4" @@ -6522,7 +6497,6 @@ "version": "6.5.0", "resolved": "https://registry.npmjs.org/fdir/-/fdir-6.5.0.tgz", "integrity": "sha512-tIbYtZbucOs0BRGqPJkshJUYdL+SDH7dVM8gjy+ERp3WAUjLEFJE+02kanyHtwjWOnwrKYBiwAmM0p4kLJAnXg==", - "dev": true, "license": "MIT", "engines": { "node": ">=12.0.0" @@ -6560,7 +6534,6 @@ "version": "7.1.1", "resolved": "https://registry.npmjs.org/fill-range/-/fill-range-7.1.1.tgz", "integrity": "sha512-YsGpe3WHLK8ZYi4tWDg2Jy3ebRz2rXowDxnld4bkQB00cc/1Zw9AWnC0i9ztDJitivtQvaI9KaLyKrc+hBW0yg==", - "dev": true, "license": "MIT", "dependencies": { "to-regex-range": "^5.0.1" @@ -6700,7 +6673,6 @@ "version": "2.3.2", "resolved": "https://registry.npmjs.org/fsevents/-/fsevents-2.3.2.tgz", "integrity": "sha512-xiqMQR4xAeHTuB9uWm+fFRcIOgKBMiOBP+eXiyT7jsgVCq1bkVygt00oASowB7EdtpOHaaPgKt812P9ab+DDKA==", - "dev": true, "hasInstallScript": true, "license": "MIT", "optional": true, @@ -6851,7 +6823,6 @@ "version": "6.0.2", "resolved": "https://registry.npmjs.org/glob-parent/-/glob-parent-6.0.2.tgz", "integrity": "sha512-XxwI8EOhVQgWp6iDL+3b0r86f4d6AX6zSU55HfB4ydCEuXLXc5FcYeOu+nnGftS4TEju/11rt4KJPTMgbfmv4A==", - "dev": true, "license": "ISC", "dependencies": { "is-glob": "^4.0.3" @@ -7349,7 +7320,6 @@ "version": "2.1.0", "resolved": "https://registry.npmjs.org/is-binary-path/-/is-binary-path-2.1.0.tgz", "integrity": "sha512-ZMERYes6pDydyuGidse7OsHxtbI7WVeUEozgR/g7rd0xUimYNlvZRE/K2MgZTjWy725IfelLeVcEM97mmtRGXw==", - "dev": true, "license": "MIT", "dependencies": { "binary-extensions": "^2.0.0" @@ -7402,7 +7372,6 @@ "version": "2.16.1", "resolved": "https://registry.npmjs.org/is-core-module/-/is-core-module-2.16.1.tgz", "integrity": "sha512-UfoeMA6fIJ8wTYFEUjelnaGI67v6+N7qXJEvQuIGa99l4xsCruSYOVSQ0uPANn4dAzm8lkYPaKLrrijLq7x23w==", - "dev": true, "license": "MIT", "dependencies": { "hasown": "^2.0.2" @@ -7463,7 +7432,6 @@ "version": "2.1.1", "resolved": "https://registry.npmjs.org/is-extglob/-/is-extglob-2.1.1.tgz", "integrity": "sha512-SbKbANkN603Vi4jEZv49LeVJMn4yGwsbzZworEoyEiutsN3nJYdbO36zfhGJ6QEDpOZIFkDtnq5JRxmvl3jsoQ==", - "dev": true, "license": "MIT", "engines": { "node": ">=0.10.0" @@ -7509,7 +7477,6 @@ "version": "4.0.3", "resolved": "https://registry.npmjs.org/is-glob/-/is-glob-4.0.3.tgz", "integrity": "sha512-xelSayHH36ZgE7ZWhli7pW34hNbNl8Ojv5KVmkJD4hBdD3th8Tfk9vYasLM+mXWOZhFkgZfxhLSnrwRr4elSSg==", - "dev": true, "license": "MIT", "dependencies": { "is-extglob": "^2.1.1" @@ -7558,7 +7525,6 @@ "version": "7.0.0", "resolved": "https://registry.npmjs.org/is-number/-/is-number-7.0.0.tgz", "integrity": "sha512-41Cifkg6e8TylSpdtTpeLVMqvSBEVzTttHvERD741+pnZ8ANv0004MRL43QKPDlK9cGvNp6NZWZUBlbGXYxxng==", - "dev": true, "license": "MIT", "engines": { "node": ">=0.12.0" @@ -7835,7 +7801,6 @@ "version": "1.21.7", "resolved": "https://registry.npmjs.org/jiti/-/jiti-1.21.7.tgz", "integrity": "sha512-/imKNG4EbWNrVjoNC/1H5/9GFy+tqjGBHCaSsN+P2RnPqjsLmv6UD3Ej+Kj8nBWaRAwyk7kK5ZUc+OEatnTR3A==", - "dev": true, "license": "MIT", "bin": { "jiti": "bin/jiti.js" @@ -8121,7 +8086,6 @@ "version": "3.1.3", "resolved": "https://registry.npmjs.org/lilconfig/-/lilconfig-3.1.3.tgz", "integrity": "sha512-/vlFKAoH5Cgt3Ie+JLhRbwOsCQePABiU3tJ1egGvyQ+33R/vcwM2Zl2QR/LzjsBeItPt3oSVXapn+m4nQDvpzw==", - "dev": true, "license": "MIT", "engines": { "node": ">=14" @@ -8134,7 +8098,6 @@ "version": "1.2.4", "resolved": "https://registry.npmjs.org/lines-and-columns/-/lines-and-columns-1.2.4.tgz", "integrity": "sha512-7ylylesZQ/PV29jhEDl3Ufjo6ZX7gCqJr5F7PKrqc93v7fzSymt1BpwEU8nAUXs8qzzvqhbjhK5QZg6Mt/HkBg==", - "dev": true, "license": "MIT" }, "node_modules/locate-path": { @@ -8588,7 +8551,6 @@ "version": "1.4.1", "resolved": "https://registry.npmjs.org/merge2/-/merge2-1.4.1.tgz", "integrity": "sha512-8q7VEgMJW4J8tcfVPy8g09NcQwZdbwFEqhe/WZkoIzjn/3TGDwtOCYtXGxA3O8tPzpczCCDgv+P2P5y00ZJOOg==", - "dev": true, "license": "MIT", "engines": { "node": ">= 8" @@ -9161,7 +9123,6 @@ "version": "4.0.8", "resolved": "https://registry.npmjs.org/micromatch/-/micromatch-4.0.8.tgz", "integrity": "sha512-PXwfBhYu0hBCPw8Dn0E+WDYb7af3dSLVWKi3HGv84IdF4TyFoC0ysxFd0Goxw7nSv4T/PzEJQxsYsEiFCKo2BA==", - "dev": true, "license": "MIT", "dependencies": { "braces": "^3.0.3", @@ -9175,7 +9136,6 @@ "version": "2.3.1", "resolved": "https://registry.npmjs.org/picomatch/-/picomatch-2.3.1.tgz", "integrity": "sha512-JU3teHTNjmE2VCGFzuY8EXzCDVwEqB2a8fsIvwaStHhAWJEeVd1o1QD80CU6+ZdEXXSLbSsuLwJjkCBWqRQUVA==", - "dev": true, "license": "MIT", "engines": { "node": ">=8.6" @@ -9290,7 +9250,6 @@ "version": "2.7.0", "resolved": "https://registry.npmjs.org/mz/-/mz-2.7.0.tgz", "integrity": "sha512-z81GNO7nnYMEhrGh9LeymoE4+Yr0Wn5McHIZMK5cfQCl+NDX08sCZgUc9/6MHni9IWuFLm1Z3HTCXu2z9fN62Q==", - "dev": true, "license": "MIT", "dependencies": { "any-promise": "^1.0.0", @@ -9502,7 +9461,6 @@ "version": "3.0.0", "resolved": "https://registry.npmjs.org/normalize-path/-/normalize-path-3.0.0.tgz", "integrity": "sha512-6eZs5Ls3WtCisHWp9S2GUy8dqkpGi4BVSz3GaqiE6ezub0512ESztXUwUB6C6IKbQkY2Pnb/mD4WYojCRwcwLA==", - "dev": true, "license": "MIT", "engines": { "node": ">=0.10.0" @@ -9521,7 +9479,6 @@ "version": "3.0.0", "resolved": "https://registry.npmjs.org/object-hash/-/object-hash-3.0.0.tgz", "integrity": "sha512-RSn9F68PjH9HqtltsSnqYC1XXoWe9Bju5+213R98cNGttag9q9yAOTzdbsqvIa7aNm5WffBZFpWYr2aWrklWAw==", - "dev": true, "license": "MIT", "engines": { "node": ">= 6" @@ -9866,7 +9823,6 @@ "version": "1.0.7", "resolved": "https://registry.npmjs.org/path-parse/-/path-parse-1.0.7.tgz", "integrity": "sha512-LDJzPVEEEPR+y48z93A0Ed0yXb8pAByGWo/k5YYdYgpY2/2EsOsksJrq7lOHxryrVOn1ejG6oAp8ahvOIQD8sw==", - "dev": true, "license": "MIT" }, "node_modules/path-scurry": { @@ -9913,7 +9869,6 @@ "version": "4.0.3", "resolved": "https://registry.npmjs.org/picomatch/-/picomatch-4.0.3.tgz", "integrity": "sha512-5gTmgEY/sqK6gFXLIsQNH19lWb4ebPDLA4SdLP7dsWkIXHWlG66oPuVvXSGFPppYZz8ZDZq0dYYrbHfBCVUb1Q==", - "dev": true, "license": "MIT", "engines": { "node": ">=12" @@ -9926,7 +9881,6 @@ "version": "2.3.0", "resolved": "https://registry.npmjs.org/pify/-/pify-2.3.0.tgz", "integrity": "sha512-udgsAY+fTnvv7kI7aaxbqwWNb0AHiB0qBO89PZKPkoTmGOgdbrHDKD+0B2X4uTfJ/FT1R09r9gTsjUjNJotuog==", - "dev": true, "license": "MIT", "engines": { "node": ">=0.10.0" @@ -9936,7 +9890,6 @@ "version": "4.0.7", "resolved": "https://registry.npmjs.org/pirates/-/pirates-4.0.7.tgz", "integrity": "sha512-TfySrs/5nm8fQJDcBDuUng3VOUKsd7S+zqvbOTiGXHfxX4wK31ard+hoNuvkicM/2YFzlpDgABOevKSsB4G/FA==", - "dev": true, "license": "MIT", "engines": { "node": ">= 6" @@ -9946,7 +9899,7 @@ "version": "1.58.1", "resolved": "https://registry.npmjs.org/playwright/-/playwright-1.58.1.tgz", "integrity": "sha512-+2uTZHxSCcxjvGc5C891LrS1/NlxglGxzrC4seZiVjcYVQfUa87wBL6rTDqzGjuoWNjnBzRqKmF6zRYGMvQUaQ==", - "dev": true, + "devOptional": true, "license": "Apache-2.0", "dependencies": { "playwright-core": "1.58.1" @@ -9965,7 +9918,7 @@ "version": "1.58.1", "resolved": "https://registry.npmjs.org/playwright-core/-/playwright-core-1.58.1.tgz", "integrity": "sha512-bcWzOaTxcW+VOOGBCQgnaKToLJ65d6AqfLVKEWvexyS3AS6rbXl+xdpYRMGSRBClPvyj44njOWoxjNdL/H9UNg==", - "dev": true, + "devOptional": true, "license": "Apache-2.0", "bin": { "playwright-core": "cli.js" @@ -9988,7 +9941,6 @@ "version": "8.5.6", "resolved": "https://registry.npmjs.org/postcss/-/postcss-8.5.6.tgz", "integrity": "sha512-3Ybi1tAuwAP9s0r1UQ2J4n5Y0G05bJkpUIO0/bI9MhwmD70S5aTWbXGBwxHrelT+XM1k6dM0pk+SwNkpTRN7Pg==", - "dev": true, "funding": [ { "type": "opencollective", @@ -10017,7 +9969,6 @@ "version": "15.1.0", "resolved": "https://registry.npmjs.org/postcss-import/-/postcss-import-15.1.0.tgz", "integrity": "sha512-hpr+J05B2FVYUAXHeK1YyI267J/dDDhMU6B6civm8hSY1jYJnBXxzKDKDswzJmtLHryrjhnDjqqp/49t8FALew==", - "dev": true, "license": "MIT", "dependencies": { "postcss-value-parser": "^4.0.0", @@ -10035,7 +9986,6 @@ "version": "4.1.0", "resolved": "https://registry.npmjs.org/postcss-js/-/postcss-js-4.1.0.tgz", "integrity": "sha512-oIAOTqgIo7q2EOwbhb8UalYePMvYoIeRY2YKntdpFQXNosSu3vLrniGgmH9OKs/qAkfoj5oB3le/7mINW1LCfw==", - "dev": true, "funding": [ { "type": "opencollective", @@ -10061,7 +10011,6 @@ "version": "6.0.1", "resolved": "https://registry.npmjs.org/postcss-load-config/-/postcss-load-config-6.0.1.tgz", "integrity": "sha512-oPtTM4oerL+UXmx+93ytZVN82RrlY/wPUV8IeDxFrzIjXOLF1pN+EmKPLbubvKHT2HC20xXsCAH2Z+CKV6Oz/g==", - "dev": true, "funding": [ { "type": "opencollective", @@ -10104,7 +10053,6 @@ "version": "6.2.0", "resolved": "https://registry.npmjs.org/postcss-nested/-/postcss-nested-6.2.0.tgz", "integrity": "sha512-HQbt28KulC5AJzG+cZtj9kvKB93CFCdLvog1WFLf1D+xmMvPGlBstkpTEZfK5+AN9hfJocyBFCNiqyS48bpgzQ==", - "dev": true, "funding": [ { "type": "opencollective", @@ -10130,7 +10078,6 @@ "version": "6.1.2", "resolved": "https://registry.npmjs.org/postcss-selector-parser/-/postcss-selector-parser-6.1.2.tgz", "integrity": "sha512-Q8qQfPiZ+THO/3ZrOrO0cJJKfpYCagtMUkXbnEfmgUjwXg6z/WBeOyS9APBBPCTSiDV+s4SwQGu8yFsiMRIudg==", - "dev": true, "license": "MIT", "dependencies": { "cssesc": "^3.0.0", @@ -10144,7 +10091,6 @@ "version": "4.2.0", "resolved": "https://registry.npmjs.org/postcss-value-parser/-/postcss-value-parser-4.2.0.tgz", "integrity": "sha512-1NNCs6uurfkVbeXG4S8JFT9t19m45ICnif8zWLd5oPSZ50QnwMfK+H3jv408d4jw/7Bttv5axS5IiHoLaVNHeQ==", - "dev": true, "license": "MIT" }, "node_modules/prelude-ls": { @@ -10251,7 +10197,6 @@ "version": "1.2.3", "resolved": "https://registry.npmjs.org/queue-microtask/-/queue-microtask-1.2.3.tgz", "integrity": "sha512-NuaNSa6flKT5JaSYQzJok04JzTL1CA6aGhv5rfLW3PgqA+M2ChpZQnAC8h8i4ZFkBS8X5RqkDBHA7r4hej3K9A==", - "dev": true, "funding": [ { "type": "github", @@ -11040,7 +10985,6 @@ "version": "1.0.0", "resolved": "https://registry.npmjs.org/read-cache/-/read-cache-1.0.0.tgz", "integrity": "sha512-Owdv/Ft7IjOgm/i0xvNDZ1LrRANRfew4b2prF3OWMQLxLfu3bS8FVhCsrSCMK4lR56Y9ya+AThoTpDCTxCmpRA==", - "dev": true, "license": "MIT", "dependencies": { "pify": "^2.3.0" @@ -11050,7 +10994,6 @@ "version": "3.6.0", "resolved": "https://registry.npmjs.org/readdirp/-/readdirp-3.6.0.tgz", "integrity": "sha512-hOS089on8RduqdbhvQ5Z37A0ESjsqz6qnRcffsMU3495FuTdqSm+7bhJ29JvIOsBDEEnan5DPu9t3To9VRlMzA==", - "dev": true, "license": "MIT", "dependencies": { "picomatch": "^2.2.1" @@ -11063,7 +11006,6 @@ "version": "2.3.1", "resolved": "https://registry.npmjs.org/picomatch/-/picomatch-2.3.1.tgz", "integrity": "sha512-JU3teHTNjmE2VCGFzuY8EXzCDVwEqB2a8fsIvwaStHhAWJEeVd1o1QD80CU6+ZdEXXSLbSsuLwJjkCBWqRQUVA==", - "dev": true, "license": "MIT", "engines": { "node": ">=8.6" @@ -11361,7 +11303,6 @@ "version": "1.22.11", "resolved": "https://registry.npmjs.org/resolve/-/resolve-1.22.11.tgz", "integrity": "sha512-RfqAvLnMl313r7c9oclB1HhUEAezcpLjz95wFH4LVuhk9JF/r22qmVP9AMmOU4vMX7Q8pN8jwNg/CSpdFnMjTQ==", - "dev": true, "license": "MIT", "dependencies": { "is-core-module": "^2.16.1", @@ -11402,7 +11343,6 @@ "version": "1.1.0", "resolved": "https://registry.npmjs.org/reusify/-/reusify-1.1.0.tgz", "integrity": "sha512-g6QUff04oZpHs0eG5p83rFLhHeV00ug/Yf9nZM6fLeUrPguBTkTQOdpAWWspMh55TZfVQDPaN3NQJfbVRAxdIw==", - "dev": true, "license": "MIT", "engines": { "iojs": ">=1.0.0", @@ -11458,7 +11398,6 @@ "version": "1.2.0", "resolved": "https://registry.npmjs.org/run-parallel/-/run-parallel-1.2.0.tgz", "integrity": "sha512-5l4VyZR86LZ/lDxZTR6jqL8AFE2S0IFLMP26AbjsLVADxHdhB/c0GUsH+y39UfCi3dzz8OlQuPmnaJOMoDHQBA==", - "dev": true, "funding": [ { "type": "github", @@ -12099,7 +12038,6 @@ "version": "3.35.1", "resolved": "https://registry.npmjs.org/sucrase/-/sucrase-3.35.1.tgz", "integrity": "sha512-DhuTmvZWux4H1UOnWMB3sk0sbaCVOoQZjv8u1rDoTV0HTdGem9hkAZtl4JZy8P2z4Bg0nT+YMeOFyVr4zcG5Tw==", - "dev": true, "license": "MIT", "dependencies": { "@jridgewell/gen-mapping": "^0.3.2", @@ -12135,7 +12073,6 @@ "version": "1.0.0", "resolved": "https://registry.npmjs.org/supports-preserve-symlinks-flag/-/supports-preserve-symlinks-flag-1.0.0.tgz", "integrity": "sha512-ot0WnXS9fgdkgIcePe6RHNk1WA8+muPa6cSjeR3V8K27q9BB1rTE3R1p7Hv0z1ZyAc8s6Vvv8DIyWf681MAt0w==", - "dev": true, "license": "MIT", "engines": { "node": ">= 0.4" @@ -12171,7 +12108,6 @@ "version": "3.4.19", "resolved": "https://registry.npmjs.org/tailwindcss/-/tailwindcss-3.4.19.tgz", "integrity": "sha512-3ofp+LL8E+pK/JuPLPggVAIaEuhvIz4qNcf3nA1Xn2o/7fb7s/TYpHhwGDv1ZU3PkBluUVaF8PyCHcm48cKLWQ==", - "dev": true, "license": "MIT", "dependencies": { "@alloc/quick-lru": "^5.2.0", @@ -12209,7 +12145,6 @@ "version": "3.3.3", "resolved": "https://registry.npmjs.org/fast-glob/-/fast-glob-3.3.3.tgz", "integrity": "sha512-7MptL8U0cqcFdzIzwOTHoilX9x5BrNqye7Z/LuC7kCMRio1EMSyqRK3BEAUD7sXRq4iT4AzTVuZdhgQ2TCvYLg==", - "dev": true, "license": "MIT", "dependencies": { "@nodelib/fs.stat": "^2.0.2", @@ -12226,7 +12161,6 @@ "version": "5.1.2", "resolved": "https://registry.npmjs.org/glob-parent/-/glob-parent-5.1.2.tgz", "integrity": "sha512-AOIgSQCepiJYwP3ARnGx+5VnTu2HBYdzbGP45eLw1vr3zB3vZLeyed1sC9hnbcOc9/SrMyM5RPQrkGz4aS9Zow==", - "dev": true, "license": "ISC", "dependencies": { "is-glob": "^4.0.1" @@ -12254,7 +12188,6 @@ "version": "3.3.1", "resolved": "https://registry.npmjs.org/thenify/-/thenify-3.3.1.tgz", "integrity": "sha512-RVZSIV5IG10Hk3enotrhvz0T9em6cyHBLkH/YAZuKqd8hRkKhSfCGIcP2KUY0EPxndzANBmNllzWPwak+bheSw==", - "dev": true, "license": "MIT", "dependencies": { "any-promise": "^1.0.0" @@ -12264,7 +12197,6 @@ "version": "1.6.0", "resolved": "https://registry.npmjs.org/thenify-all/-/thenify-all-1.6.0.tgz", "integrity": "sha512-RNxQH/qI8/t3thXJDwcstUO4zeqo64+Uy/+sNVRBx4Xn2OX+OZ9oP+iJnNFqplFra2ZUVeKCSa2oVWi3T4uVmA==", - "dev": true, "license": "MIT", "dependencies": { "thenify": ">= 3.1.0 < 4" @@ -12306,7 +12238,6 @@ "version": "0.2.15", "resolved": "https://registry.npmjs.org/tinyglobby/-/tinyglobby-0.2.15.tgz", "integrity": "sha512-j2Zq4NyQYG5XMST4cbs02Ak8iJUdxRM0XI5QyxXuZOzKOINmWurp3smXu3y5wDcJrptwpSjgXHzIQxR0omXljQ==", - "dev": true, "license": "MIT", "dependencies": { "fdir": "^6.5.0", @@ -12373,7 +12304,6 @@ "version": "5.0.1", "resolved": "https://registry.npmjs.org/to-regex-range/-/to-regex-range-5.0.1.tgz", "integrity": "sha512-65P7iz6X5yEr1cwcgvQxbbIw7Uk3gOy5dIdtZ4rDveLqhrdJP+Li/Hx6tyK0NEb+2GCyneCMJiGqrADCSNk8sQ==", - "dev": true, "license": "MIT", "dependencies": { "is-number": "^7.0.0" @@ -12461,7 +12391,6 @@ "version": "0.1.13", "resolved": "https://registry.npmjs.org/ts-interface-checker/-/ts-interface-checker-0.1.13.tgz", "integrity": "sha512-Y/arvbn+rrz3JCKl9C4kVNfTfSm2/mEp5FSz5EsZSANGPSlQrpRI5M4PKF+mJnE52jOO90PnPSc3Ur3bTQw0gA==", - "dev": true, "license": "Apache-2.0" }, "node_modules/tsconfig-paths": { @@ -12578,7 +12507,7 @@ "version": "5.9.3", "resolved": "https://registry.npmjs.org/typescript/-/typescript-5.9.3.tgz", "integrity": "sha512-jl1vZzPDinLr9eUt3J/t7V6FgNEw9QjvBPdysz9KfQDD41fQrC2Y4vKQdiaUpFT4bXlb1RHhLpp8wtm6M5TgSw==", - "dev": true, + "devOptional": true, "license": "Apache-2.0", "bin": { "tsc": "bin/tsc", @@ -12780,7 +12709,6 @@ "version": "1.0.2", "resolved": "https://registry.npmjs.org/util-deprecate/-/util-deprecate-1.0.2.tgz", "integrity": "sha512-EPD5q1uXyFxJpCrLnCc1nHnq3gOa6DZBocAIiI2TaSCA7VCJ1UJDMagCzIkXNsUYfD1daK//LTEQ8xiIbrHtcw==", - "dev": true, "license": "MIT" }, "node_modules/uuid": { @@ -13234,7 +13162,7 @@ "version": "8.19.0", "resolved": "https://registry.npmjs.org/ws/-/ws-8.19.0.tgz", "integrity": "sha512-blAT2mjOEIi0ZzruJfIhb3nps74PRWTCz1IjglWEEpQl5XS/UNama6u2/rjFkDDouqr4L67ry+1aGIALViWjDg==", - "dev": true, + "devOptional": true, "license": "MIT", "engines": { "node": ">=10.0.0" diff --git a/ui/litellm-dashboard/src/app/(dashboard)/skills/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/skills/page.tsx new file mode 100644 index 00000000000..47d2331bed8 --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/skills/page.tsx @@ -0,0 +1,17 @@ +"use client"; + +import ClaudeCodePluginsPanel from "@/components/claude_code_plugins"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; + +const SkillsPage = () => { + const { accessToken, userRole } = useAuthorized(); + + return ( + + ); +}; + +export default SkillsPage; diff --git a/ui/litellm-dashboard/src/app/page.tsx b/ui/litellm-dashboard/src/app/page.tsx index 44df1b5bd41..746847057f1 100644 --- a/ui/litellm-dashboard/src/app/page.tsx +++ b/ui/litellm-dashboard/src/app/page.tsx @@ -624,7 +624,7 @@ function CreateKeyPageContent() { ) : page == "tag-management" ? ( - ) : page == "claude-code-plugins" ? ( + ) : page == "skills" || page == "claude-code-plugins" ? ( ) : page == "access-groups" ? ( diff --git a/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx b/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx index 537aa001b70..0be6919c1fa 100644 --- a/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx +++ b/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx @@ -5,7 +5,10 @@ import MakeModelPublicForm from "@/components/AIHub/forms/MakeModelPublicForm"; import { mcpHubColumns, MCPServerData } from "@/components/mcp_hub_table_columns"; import { modelHubColumns } from "@/components/model_hub_table_columns"; import UsefulLinksManagement from "@/components/AIHub/UsefulLinksManagement"; -import ClaudeCodeMarketplaceTab from "@/components/AIHub/ClaudeCodeMarketplaceTab"; +import { getClaudeCodePluginsList } from "@/components/networking"; +import { Plugin } from "@/components/claude_code_plugins/types"; +import SkillHubDashboard from "@/components/AIHub/SkillHubDashboard"; +import MakeSkillPublicForm from "@/components/claude_code_plugins/MakeSkillPublicForm"; import { ModelDataTable } from "@/components/model_dashboard/table"; import ModelFilters from "@/components/model_filters"; import NotificationsManager from "@/components/molecules/notifications_manager"; @@ -78,6 +81,10 @@ const ModelHubTable: React.FC = ({ accessToken, publicPage, const [selectedMcpServer, setSelectedMcpServer] = useState(null); const [isMcpModalVisible, setIsMcpModalVisible] = useState(false); const [isMakeMcpPublicModalVisible, setIsMakeMcpPublicModalVisible] = useState(false); + // Skill Hub state + const [skillHubData, setSkillHubData] = useState([]); + const [skillLoading, setSkillLoading] = useState(false); + const [isMakeSkillPublicModalVisible, setIsMakeSkillPublicModalVisible] = useState(false); const router = useRouter(); const { data: uiSettings, isLoading: isUISettingsLoading } = useUISettings(); @@ -208,6 +215,25 @@ const ModelHubTable: React.FC = ({ accessToken, publicPage, } }, [publicPage, accessToken]); + // Fetch Skill Hub data — all skills for admins, enabled-only for public page + useEffect(() => { + const fetchSkillData = async () => { + if (!accessToken) return; + try { + setSkillLoading(true); + const enabledOnly = publicPage === true; + const response = await getClaudeCodePluginsList(accessToken, enabledOnly); + setSkillHubData(response.plugins); + } catch (error) { + console.error("Error fetching skill hub data", error); + } finally { + setSkillLoading(false); + } + }; + + fetchSkillData(); + }, [accessToken, publicPage]); + const showModal = (model: ModelGroupInfo) => { setSelectedModel(model); setIsModalVisible(true); @@ -406,7 +432,7 @@ const ModelHubTable: React.FC = ({ accessToken, publicPage, Model Hub Agent Hub MCP Hub - Claude Code Plugin Marketplace + Skill Hub @@ -492,9 +518,26 @@ const ModelHubTable: React.FC = ({ accessToken, publicPage,
- {/* Plugin Marketplace Tab */} + {/* Skill Hub Tab */} - + {publicPage == false && isAdminRole(userRole || "") && ( +
+ +
+ )} + { + const response = await getClaudeCodePluginsList(accessToken || "", publicPage === true); + setSkillHubData(response.plugins); + }} + /> @@ -1055,6 +1098,18 @@ if __name__ == "__main__": mcpHubData={mcpHubData || []} onSuccess={handleMakeMcpPublicSuccess} /> + + {/* Make Skill Public Form */} + setIsMakeSkillPublicModalVisible(false)} + accessToken={accessToken || ""} + skillsList={skillHubData} + onSuccess={async () => { + const response = await getClaudeCodePluginsList(accessToken || "", publicPage === true); + setSkillHubData(response.plugins); + }} + /> ); }; diff --git a/ui/litellm-dashboard/src/components/AIHub/SkillHubDashboard.tsx b/ui/litellm-dashboard/src/components/AIHub/SkillHubDashboard.tsx new file mode 100644 index 00000000000..af6668dcbeb --- /dev/null +++ b/ui/litellm-dashboard/src/components/AIHub/SkillHubDashboard.tsx @@ -0,0 +1,139 @@ +import React, { useMemo, useState } from "react"; +import { Text } from "@tremor/react"; +import { SearchOutlined } from "@ant-design/icons"; +import { Input, Select } from "antd"; +import { Plugin } from "@/components/claude_code_plugins/types"; +import { ModelDataTable } from "@/components/model_dashboard/table"; +import { skillHubColumns } from "@/components/skill_hub_table_columns"; +import SkillDetail from "@/components/claude_code_plugins/skill_detail"; + +interface SkillHubDashboardProps { + skills: Plugin[]; + isLoading: boolean; + isAdmin?: boolean; + accessToken?: string | null; + publicPage?: boolean; + onPublishSuccess?: () => void; +} + +const SkillHubDashboard: React.FC = ({ + skills, + isLoading, + isAdmin, + accessToken, + publicPage = false, + onPublishSuccess, +}) => { + const [search, setSearch] = useState(""); + const [domainFilter, setDomainFilter] = useState(undefined); + const [selectedSkill, setSelectedSkill] = useState(null); + + const copyToClipboard = (text: string) => { + navigator.clipboard.writeText(text); + }; + + // Derived stats + const totalSkills = skills.length; + const domains = useMemo(() => [...new Set(skills.map((s) => s.domain).filter(Boolean))], [skills]); + const namespaces = useMemo(() => [...new Set(skills.map((s) => s.namespace).filter(Boolean))], [skills]); + + // Filtered table data + const filteredSkills = useMemo(() => { + let result = skills; + if (domainFilter) { + result = result.filter((s) => (s.domain || "General") === domainFilter); + } + if (search.trim()) { + const q = search.toLowerCase(); + result = result.filter( + (s) => + s.name.toLowerCase().includes(q) || + s.description?.toLowerCase().includes(q) || + s.domain?.toLowerCase().includes(q) || + s.namespace?.toLowerCase().includes(q) || + s.keywords?.some((k) => k.toLowerCase().includes(q)) + ); + } + return result; + }, [skills, search, domainFilter]); + + if (selectedSkill) { + return ( + setSelectedSkill(null)} + isAdmin={isAdmin} + accessToken={accessToken} + onPublishClick={onPublishSuccess} + /> + ); + } + + if (isLoading) { + return
Loading skills...
; + } + + return ( +
+ {/* Stats row */} +
+
+
Total Skills
+
{totalSkills}
+
+
+
Namespaces
+
{namespaces.length}
+
+
+
Domains
+
{domains.length}
+
+
+ + {/* Search + filters + table */} +
+
+

+ All {publicPage ? "Public " : ""}Skills +

+
+ } + placeholder="Search by name, namespace, or tag…" + value={search} + onChange={(e) => setSearch(e.target.value)} + style={{ width: 280 }} + allowClear + /> +
+
+ setSelectedSkill(skill), + copyToClipboard, + publicPage + )} + data={filteredSkills} + isLoading={false} + defaultSorting={[{ id: "name", desc: false }]} + /> +
+ + Showing {filteredSkills.length} of {totalSkills} skill{totalSkills !== 1 ? "s" : ""} + +
+
+
+ ); +}; + +export default SkillHubDashboard; diff --git a/ui/litellm-dashboard/src/components/claude_code_plugins.tsx b/ui/litellm-dashboard/src/components/claude_code_plugins.tsx index b2e66349a23..67d5f6e0ad5 100644 --- a/ui/litellm-dashboard/src/components/claude_code_plugins.tsx +++ b/ui/litellm-dashboard/src/components/claude_code_plugins.tsx @@ -7,8 +7,8 @@ import { } from "./networking"; import AddPluginForm from "./claude_code_plugins/add_plugin_form"; import PluginTable from "./claude_code_plugins/plugin_table"; +import SkillDetail from "./claude_code_plugins/skill_detail"; import { isAdminRole } from "@/utils/roles"; -import PluginInfoView from "./claude_code_plugins/plugin_info"; import NotificationsManager from "./molecules/notifications_manager"; import { Plugin, ListPluginsResponse } from "./claude_code_plugins/types"; @@ -29,27 +29,22 @@ const ClaudeCodePluginsPanel: React.FC = ({ name: string; displayName: string; } | null>(null); - const [selectedPluginId, setSelectedPluginId] = useState( - null - ); + const [selectedSkill, setSelectedSkill] = useState(null); const isAdmin = userRole ? isAdminRole(userRole) : false; const fetchPlugins = async () => { - if (!accessToken) { - return; - } + if (!accessToken) return; setIsLoading(true); try { const response: ListPluginsResponse = await getClaudeCodePluginsList( accessToken, - false // Get all plugins (enabled and disabled) + false ); - console.log(`Claude Code plugins: ${JSON.stringify(response)}`); setPluginsList(response.plugins); } catch (error) { - console.error("Error fetching Claude Code plugins:", error); + console.error("Error fetching skills:", error); } finally { setIsLoading(false); } @@ -59,21 +54,6 @@ const ClaudeCodePluginsPanel: React.FC = ({ fetchPlugins(); }, [accessToken]); - const handleAddPlugin = () => { - if (selectedPluginId) { - setSelectedPluginId(null); - } - setIsAddModalVisible(true); - }; - - const handleCloseModal = () => { - setIsAddModalVisible(false); - }; - - const handleSuccess = () => { - fetchPlugins(); - }; - const handleDeleteClick = (pluginName: string, displayName: string) => { setPluginToDelete({ name: pluginName, displayName }); }; @@ -84,79 +64,76 @@ const ClaudeCodePluginsPanel: React.FC = ({ setIsDeleting(true); try { await deleteClaudeCodePlugin(accessToken, pluginToDelete.name); - NotificationsManager.success( - `Plugin "${pluginToDelete.displayName}" deleted successfully` - ); + NotificationsManager.success(`Skill "${pluginToDelete.displayName}" deleted successfully`); fetchPlugins(); } catch (error) { - console.error("Error deleting plugin:", error); - NotificationsManager.error("Failed to delete plugin"); + console.error("Error deleting skill:", error); + NotificationsManager.error("Failed to delete skill"); } finally { setIsDeleting(false); setPluginToDelete(null); } }; - const handleDeleteCancel = () => { - setPluginToDelete(null); - }; - return (
-
-

Claude Code Plugins

-

- Manage Claude Code marketplace plugins. Add, enable, disable, or - delete plugins that will be available in your marketplace catalog. - Enabled plugins will appear in the public marketplace at{" "} - /claude-code/marketplace.json. -

-
- -
-
- - {selectedPluginId ? ( - setSelectedPluginId(null)} - accessToken={accessToken} + {selectedSkill ? ( + setSelectedSkill(null)} isAdmin={isAdmin} - onPluginUpdated={fetchPlugins} + accessToken={accessToken} + onPublishClick={fetchPlugins} /> ) : ( - setSelectedPluginId(id)} - /> + <> +
+

Skills

+

+ Register Claude Code skills. Published skills appear in the Skill Hub for all users and + are served via{" "} + /claude-code/marketplace.json. +

+
+ +
+
+ + { + const skill = pluginsList.find((p) => p.id === id); + if (skill) setSelectedSkill(skill); + }} + /> + )} setIsAddModalVisible(false)} accessToken={accessToken} - onSuccess={handleSuccess} + onSuccess={fetchPlugins} /> {pluginToDelete && ( setPluginToDelete(null)} confirmLoading={isDeleting} okText="Delete" okButtonProps={{ danger: true }} >

- Are you sure you want to delete plugin:{" "} + Are you sure you want to delete skill:{" "} {pluginToDelete.displayName}?

This action cannot be undone.

diff --git a/ui/litellm-dashboard/src/components/claude_code_plugins/MakeSkillPublicForm.tsx b/ui/litellm-dashboard/src/components/claude_code_plugins/MakeSkillPublicForm.tsx new file mode 100644 index 00000000000..351502d78b4 --- /dev/null +++ b/ui/litellm-dashboard/src/components/claude_code_plugins/MakeSkillPublicForm.tsx @@ -0,0 +1,249 @@ +import React, { useState, useEffect } from "react"; +import { Modal, Form, Steps, Button, Checkbox } from "antd"; +import { Text, Title, Badge } from "@tremor/react"; +import { enableClaudeCodePlugin, disableClaudeCodePlugin } from "../networking"; +import NotificationsManager from "../molecules/notifications_manager"; +import { Plugin } from "./types"; + +const { Step } = Steps; + +interface MakeSkillPublicFormProps { + visible: boolean; + onClose: () => void; + accessToken: string; + skillsList: Plugin[]; + onSuccess: () => void; +} + +const MakeSkillPublicForm: React.FC = ({ + visible, + onClose, + accessToken, + skillsList, + onSuccess, +}) => { + const [currentStep, setCurrentStep] = useState(0); + const [selectedSkills, setSelectedSkills] = useState>(new Set()); + const [loading, setLoading] = useState(false); + const [form] = Form.useForm(); + + const handleClose = () => { + setCurrentStep(0); + setSelectedSkills(new Set()); + form.resetFields(); + onClose(); + }; + + const handleNext = () => { + if (selectedSkills.size === 0) { + NotificationsManager.fromBackend("Please select at least one skill"); + return; + } + setCurrentStep(1); + }; + + const handleSkillSelection = (name: string, checked: boolean) => { + const next = new Set(selectedSkills); + if (checked) { + next.add(name); + } else { + next.delete(name); + } + setSelectedSkills(next); + }; + + const handleSelectAll = (checked: boolean) => { + if (checked) { + setSelectedSkills(new Set(skillsList.map((s) => s.name))); + } else { + setSelectedSkills(new Set()); + } + }; + + // Pre-check already-published skills when modal opens + useEffect(() => { + if (visible && skillsList.length > 0) { + setSelectedSkills(new Set(skillsList.filter((s) => s.enabled).map((s) => s.name))); + } + }, [visible, skillsList]); + + const handleSubmit = async () => { + if (selectedSkills.size === 0) { + NotificationsManager.fromBackend("Please select at least one skill"); + return; + } + + setLoading(true); + try { + const selectedSet = selectedSkills; + await Promise.all( + skillsList.map((skill) => { + const shouldBePublic = selectedSet.has(skill.name); + if (shouldBePublic && !skill.enabled) { + return enableClaudeCodePlugin(accessToken, skill.name); + } + if (!shouldBePublic && skill.enabled) { + return disableClaudeCodePlugin(accessToken, skill.name); + } + return Promise.resolve(); + }) + ); + + NotificationsManager.success(`Skill Hub updated — ${selectedSkills.size} skill(s) published`); + handleClose(); + onSuccess(); + } catch (error) { + console.error("Error publishing skills:", error); + NotificationsManager.fromBackend("Failed to update skills. Please try again."); + } finally { + setLoading(false); + } + }; + + const allSelected = + skillsList.length > 0 && skillsList.every((s) => selectedSkills.has(s.name)); + const isIndeterminate = selectedSkills.size > 0 && !allSelected; + + const renderStep1 = () => ( +
+
+ Select Skills to Publish + handleSelectAll(e.target.checked)} + disabled={skillsList.length === 0} + > + Select All ({skillsList.length}) + +
+ + + Selected skills will be visible to all users in the Skill Hub. + Deselected skills will be unpublished. + + +
+
+ {skillsList.length === 0 ? ( +
+ No skills registered yet. +
+ ) : ( + skillsList.map((skill) => ( +
+ handleSkillSelection(skill.name, e.target.checked)} + /> +
+
+ {skill.name} + {skill.enabled && ( + Public + )} +
+ {skill.description && ( + + {skill.description} + + )} +
+ {skill.domain && ( + {skill.domain} + )} +
+ )) + )} +
+
+ + {selectedSkills.size > 0 && ( +
+ + {selectedSkills.size} skill{selectedSkills.size !== 1 ? "s" : ""} will be published + +
+ )} +
+ ); + + const renderStep2 = () => ( +
+ Confirm Publish to Skill Hub + +
+ + Note: Published skills will be visible to all users in the Skill Hub tab. + Skills not in the list below will be unpublished. + +
+ +
+ Skills to be published: +
+
+ {Array.from(selectedSkills).map((name) => { + const skill = skillsList.find((s) => s.name === name); + return ( +
+ {name} + {skill?.domain && {skill.domain}} +
+ ); + })} +
+
+
+ +
+ + Total: {selectedSkills.size} skill{selectedSkills.size !== 1 ? "s" : ""} will be published + +
+
+ ); + + return ( + +
+ + + + + + {currentStep === 0 ? renderStep1() : renderStep2()} + +
+ +
+ {currentStep === 0 && ( + + )} + {currentStep === 1 && ( + + )} +
+
+
+ + ); +}; + +export default MakeSkillPublicForm; diff --git a/ui/litellm-dashboard/src/components/claude_code_plugins/add_plugin_form.tsx b/ui/litellm-dashboard/src/components/claude_code_plugins/add_plugin_form.tsx index d6a3f0f4fd9..fdbef3a766b 100644 --- a/ui/litellm-dashboard/src/components/claude_code_plugins/add_plugin_form.tsx +++ b/ui/litellm-dashboard/src/components/claude_code_plugins/add_plugin_form.tsx @@ -32,6 +32,80 @@ const PREDEFINED_CATEGORIES = [ "Documentation", ]; +interface ParsedSource { + source: "github" | "url" | "git-subdir"; + repo?: string; + url?: string; + path?: string; +} + +interface ParsePreview { + parsed: ParsedSource; + label: string; + suggestedName: string; +} + +function parseGitHubUrl(raw: string): ParsePreview | null { + // Strip protocol and trailing slashes/spaces + let s = raw.trim().replace(/^https?:\/\//, "").replace(/\/+$/, ""); + + if (!s.startsWith("github.com/")) return null; + + // Remove "github.com/" + const rest = s.slice("github.com/".length); + const parts = rest.split("/"); + + if (parts.length < 2) return null; + + const org = parts[0]; + const repo = parts[1]; + const repoBase = repo.replace(/\.git$/, ""); + + // github.com/org/repo (exactly 2 parts, or ends with .git) + if (parts.length === 2 || (parts.length === 2 && repoBase)) { + return { + parsed: { source: "github", repo: `${org}/${repoBase}` }, + label: `GitHub repo — ${org}/${repoBase}`, + suggestedName: repoBase, + }; + } + + // github.com/org/repo/tree/branch/folder or /blob/branch/folder/FILE.md + if ( + parts.length >= 5 && + (parts[2] === "tree" || parts[2] === "blob") + ) { + // parts[3] = branch, parts[4..] = path segments + const pathParts = parts.slice(4); + // If last segment looks like a file (has extension), drop it + const lastPart = pathParts[pathParts.length - 1]; + if (lastPart && lastPart.includes(".")) { + pathParts.pop(); + } + if (pathParts.length === 0) { + // Path resolved to repo root — treat as plain github source + return { + parsed: { source: "github", repo: `${org}/${repoBase}` }, + label: `GitHub repo — ${org}/${repoBase}`, + suggestedName: repoBase, + }; + } + const subPath = pathParts.join("/"); + const suggestedName = pathParts[pathParts.length - 1]; + return { + parsed: { + source: "git-subdir", + url: `https://github.com/${org}/${repoBase}`, + path: subPath, + }, + label: `GitHub subdir — ${org}/${repoBase} @ ${subPath}`, + suggestedName, + }; + } + + return null; +} + const AddPluginForm: React.FC = ({ visible, onClose, @@ -40,7 +114,20 @@ const AddPluginForm: React.FC = ({ }) => { const [form] = Form.useForm(); const [isSubmitting, setIsSubmitting] = useState(false); - const [sourceType, setSourceType] = useState<"github" | "url" | "git-subdir">("github"); + const [urlPreview, setUrlPreview] = useState(null); + + const handleUrlChange = (e: React.ChangeEvent) => { + const val = e.target.value; + const preview = parseGitHubUrl(val); + setUrlPreview(preview); + if (preview) { + // Auto-fill name only if it's currently empty + const currentName = form.getFieldValue("name"); + if (!currentName) { + form.setFieldsValue({ name: preview.suggestedName }); + } + } + }; const handleSubmit = async (values: any) => { if (!accessToken) { @@ -48,98 +135,62 @@ const AddPluginForm: React.FC = ({ return; } - // Validate plugin name + if (!urlPreview) { + MessageManager.error("Please enter a valid GitHub URL"); + return; + } + if (!validatePluginName(values.name)) { MessageManager.error( - "Plugin name must be kebab-case (lowercase letters, numbers, and hyphens only)" + "Skill name must be kebab-case (lowercase letters, numbers, and hyphens only)" ); return; } - // Validate semantic version if provided if (values.version && !isValidSemanticVersion(values.version)) { - MessageManager.error( - "Version must be in semantic versioning format (e.g., 1.0.0)" - ); + MessageManager.error("Version must be in semantic versioning format (e.g., 1.0.0)"); return; } - // Validate email if provided if (values.authorEmail && !isValidEmail(values.authorEmail)) { MessageManager.error("Invalid email format"); return; } - // Validate homepage URL if provided if (values.homepage && !isValidUrl(values.homepage)) { MessageManager.error("Invalid homepage URL format"); return; } - // Validate git URL for url/git-subdir source types - if ((sourceType === "url" || sourceType === "git-subdir") && values.url && !isValidUrl(values.url)) { - MessageManager.error("Invalid git URL format"); - return; - } - setIsSubmitting(true); try { - // Build plugin data const pluginData: any = { name: values.name.trim(), - source: - sourceType === "github" - ? { - source: "github", - repo: values.repo.trim(), - } - : sourceType === "git-subdir" - ? { - source: "git-subdir", - url: values.url.trim(), - path: values.path.trim(), - } - : { - source: "url", - url: values.url.trim(), - }, + source: urlPreview.parsed, }; - // Add optional fields - if (values.version) { - pluginData.version = values.version.trim(); - } - if (values.description) { - pluginData.description = values.description.trim(); - } + if (values.version) pluginData.version = values.version.trim(); + if (values.description) pluginData.description = values.description.trim(); if (values.authorName || values.authorEmail) { pluginData.author = {}; - if (values.authorName) { - pluginData.author.name = values.authorName.trim(); - } - if (values.authorEmail) { - pluginData.author.email = values.authorEmail.trim(); - } - } - if (values.homepage) { - pluginData.homepage = values.homepage.trim(); - } - if (values.category) { - pluginData.category = values.category; - } - if (values.keywords) { - pluginData.keywords = parseKeywords(values.keywords); + if (values.authorName) pluginData.author.name = values.authorName.trim(); + if (values.authorEmail) pluginData.author.email = values.authorEmail.trim(); } + if (values.homepage) pluginData.homepage = values.homepage.trim(); + if (values.category) pluginData.category = values.category; + if (values.keywords) pluginData.keywords = parseKeywords(values.keywords); + if (values.domain) pluginData.domain = values.domain.trim(); + if (values.namespace) pluginData.namespace = values.namespace.trim(); await registerClaudeCodePlugin(accessToken, pluginData); - MessageManager.success("Plugin registered successfully"); + MessageManager.success("Skill registered successfully"); form.resetFields(); - setSourceType("github"); + setUrlPreview(null); onSuccess(); onClose(); } catch (error) { - console.error("Error registering plugin:", error); - MessageManager.error("Failed to register plugin"); + console.error("Error registering skill:", error); + MessageManager.error("Failed to register skill"); } finally { setIsSubmitting(false); } @@ -147,19 +198,13 @@ const AddPluginForm: React.FC = ({ const handleCancel = () => { form.resetFields(); - setSourceType("github"); + setUrlPreview(null); onClose(); }; - const handleSourceTypeChange = (value: "github" | "url" | "git-subdir") => { - setSourceType(value); - // Clear repo/url/path fields when switching - form.setFieldsValue({ repo: undefined, url: undefined, path: undefined }); - }; - return ( = ({ onFinish={handleSubmit} className="mt-4" > - {/* Plugin Name */} + {/* Smart URL Input */} + + + + {/* Parsed preview */} + {urlPreview && ( +
+ Detected: {urlPreview.label} +
+ )} + + {/* Skill Name */} + - + - {/* Source Type */} - - - - - {/* GitHub Repository */} - {sourceType === "github" && ( + {/* Domain and Namespace — side by side */} +
- + - )} - - {/* Git URL */} - {(sourceType === "url" || sourceType === "git-subdir") && ( - + - )} - - {/* Git Subdir Path */} - {sourceType === "git-subdir" && ( - - - - )} - - {/* Version */} - - - +
{/* Description */}