From ba1b466480a000b559fde36c246c9d31392af1ce Mon Sep 17 00:00:00 2001 From: shivam Date: Sat, 7 Feb 2026 18:16:01 -0800 Subject: [PATCH 001/273] fix to ensure budget duration is being inherited from budget tier for keys --- .../proxy/common_utils/reset_budget_job.py | 36 +++++ .../key_management_endpoints.py | 7 + .../common_utils/test_reset_budget_job.py | 116 +++++++++++++++ .../test_key_management_endpoints.py | 140 ++++++++++++++++++ 4 files changed, 299 insertions(+) diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index fb600cee26b..6f038d127f6 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -69,6 +69,38 @@ class ResetBudgetJob: }, ) + async def reset_budget_for_keys_linked_to_budgets( + self, budgets_to_reset: List[LiteLLM_BudgetTableFull] + ): + """ + Resets the spend for keys linked to budget tiers that are being reset. + + This handles keys that have budget_id but no budget_duration set on the key + itself (e.g. keys created before the fix to inherit budget_duration from + the linked budget tier). + + Keys that have their own budget_duration are already handled by + reset_budget_for_litellm_keys() and are excluded here to avoid + double-resetting. + """ + budget_ids = [ + budget.budget_id + for budget in budgets_to_reset + if budget.budget_id is not None + ] + if not budget_ids: + return + + return await self.prisma_client.db.litellm_verificationtoken.update_many( + where={ + "budget_id": {"in": budget_ids}, + "budget_duration": None, # only keys without their own reset schedule + }, + data={ + "spend": 0, + }, + ) + async def reset_budget_for_litellm_budget_table(self): """ Resets the budget for all LiteLLM End-Users (Customers), and Team Members if their budget has expired @@ -112,6 +144,10 @@ class ResetBudgetJob: budgets_to_reset=budgets_to_reset ) + await self.reset_budget_for_keys_linked_to_budgets( + budgets_to_reset=budgets_to_reset + ) + if endusers_to_reset is not None and len(endusers_to_reset) > 0: for enduser in endusers_to_reset: try: diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 2eb6cf65281..b62ce329548 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -594,6 +594,13 @@ async def _common_key_generation_helper( # noqa: PLR0915 if "budget_duration" in data_json: data_json["key_budget_duration"] = data_json.pop("budget_duration", None) + elif _budget_id is not None and prisma_client is not None: + # Inherit budget_duration from linked budget tier if not explicitly set on the key + budget_row = await prisma_client.db.litellm_budgettable.find_unique( + where={"budget_id": _budget_id} + ) + if budget_row is not None and budget_row.budget_duration is not None: + data_json["key_budget_duration"] = budget_row.budget_duration if user_api_key_dict.user_id is not None: data_json["created_by"] = user_api_key_dict.user_id diff --git a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py index a059a3adcb1..f63c77c1fc8 100644 --- a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py +++ b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py @@ -25,9 +25,21 @@ class MockLiteLLMTeamMembership: return {"count": 1} +class MockLiteLLMVerificationToken: + def __init__(self): + self.update_many_calls: List[Dict[str, Any]] = [] + + async def update_many( + self, where: Dict[str, Any], data: Dict[str, Any] + ) -> Dict[str, Any]: + self.update_many_calls.append({"where": where, "data": data}) + return {"count": 1} + + class MockDB: def __init__(self): self.litellm_teammembership = MockLiteLLMTeamMembership() + self.litellm_verificationtoken = MockLiteLLMVerificationToken() class MockPrismaClient: @@ -320,3 +332,107 @@ def test_reset_budget_all(reset_budget_job, mock_prisma_client): assert mock_prisma_client.updated_data["user"][0].spend == 0.0 assert mock_prisma_client.updated_data["team"][0].spend == 0.0 assert mock_prisma_client.updated_data["enduser"][0].spend == 0.0 + + +def test_reset_budget_for_keys_linked_to_budgets(reset_budget_job, mock_prisma_client): + """ + Test that when a budget tier is reset, keys linked to that budget + (via budget_id) that don't have their own budget_duration also get + their spend reset. + + This covers the case where keys were created with budget_id but + budget_duration was not inherited to the key (pre-fix keys). + """ + from litellm.proxy._types import LiteLLM_BudgetTableFull + + now = datetime.now(timezone.utc) + + # Create a budget tier that is due for reset + test_budget = type( + "LiteLLM_BudgetTableFull", + (), + { + "max_budget": 10.0, + "budget_duration": "7d", + "budget_reset_at": now - timedelta(hours=1), + "budget_id": "7d-budget-tier", + "created_at": now - timedelta(days=7), + }, + ) + + budgets_to_reset = [test_budget] + + # Run the method + asyncio.run( + reset_budget_job.reset_budget_for_keys_linked_to_budgets( + budgets_to_reset=budgets_to_reset + ) + ) + + # Verify that update_many was called on litellm_verificationtoken + calls = mock_prisma_client.db.litellm_verificationtoken.update_many_calls + assert len(calls) == 1, f"Expected 1 update_many call, got {len(calls)}" + + # Verify the where clause filters by budget_id and null budget_duration + call = calls[0] + assert call["where"]["budget_id"] == {"in": ["7d-budget-tier"]} + assert call["where"]["budget_duration"] is None + + # Verify spend is reset to 0 + assert call["data"]["spend"] == 0 + + +def test_reset_budget_for_keys_linked_to_budgets_empty( + reset_budget_job, mock_prisma_client +): + """ + Test that when there are no budgets to reset, no update is performed + on the verification token table. + """ + # Run with empty list + asyncio.run( + reset_budget_job.reset_budget_for_keys_linked_to_budgets( + budgets_to_reset=[] + ) + ) + + # Verify no update_many calls were made + calls = mock_prisma_client.db.litellm_verificationtoken.update_many_calls + assert len(calls) == 0 + + +def test_budget_table_reset_also_resets_linked_keys( + reset_budget_job, mock_prisma_client +): + """ + Integration-style test: when reset_budget_for_litellm_budget_table runs, + it should also reset spend for keys linked to the expiring budget tiers + (in addition to end-users and team members). + """ + now = datetime.now(timezone.utc) + + test_budget = type( + "LiteLLM_BudgetTableFull", + (), + { + "max_budget": 10.0, + "budget_duration": "7d", + "budget_reset_at": now - timedelta(hours=1), + "budget_id": "7d-budget-tier", + "created_at": now - timedelta(days=7), + }, + ) + + mock_prisma_client.data["budget"] = [test_budget] + + # Run the full budget table reset + asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) + + # Verify that keys linked to the budget were also reset + calls = mock_prisma_client.db.litellm_verificationtoken.update_many_calls + assert len(calls) == 1, ( + "Expected reset_budget_for_litellm_budget_table to also reset keys " + f"linked to expiring budgets, but got {len(calls)} update_many calls" + ) + assert calls[0]["where"]["budget_id"] == {"in": ["7d-budget-tier"]} + assert calls[0]["data"]["spend"] == 0 diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 39f8d1cccb0..472504871ed 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -5559,3 +5559,143 @@ async def test_validate_key_list_check_key_hash_not_found(): assert exc_info.value.code == "403" or exc_info.value.code == 403 assert "Key Hash not found" in exc_info.value.message + + +@pytest.mark.asyncio +async def test_key_inherits_budget_duration_from_budget_tier(): + """ + Test that when a key is created with budget_id pointing to a budget tier + that has budget_duration, the key inherits budget_duration from the tier + even when budget_duration is not explicitly set on the key request. + + This verifies the fix for the bug where keys created with budget_id + would have null budget_duration and budget_reset_at, causing the + budget reset job to never reset their spend. + """ + from unittest.mock import AsyncMock, MagicMock, patch + + # Mock the budget tier lookup to return a budget with budget_duration="7d" + mock_budget_row = MagicMock() + mock_budget_row.budget_duration = "7d" + + mock_prisma = MagicMock() + mock_prisma.db.litellm_budgettable.find_unique = AsyncMock( + return_value=mock_budget_row + ) + + mock_generate_key = AsyncMock( + return_value={ + "key": "sk-test-key", + "expires": None, + "user_id": "test-user", + "team_id": None, + } + ) + + with patch( + "litellm.proxy.proxy_server.prisma_client", mock_prisma + ), patch( + "litellm.proxy.proxy_server.llm_router", None + ), patch( + "litellm.proxy.proxy_server.premium_user", False + ), patch( + "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" + ), patch( + "litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn", + mock_generate_key, + ): + await _common_key_generation_helper( + data=GenerateKeyRequest( + budget_id="7d-budget-tier", + max_budget=10.0, + # NOTE: budget_duration is intentionally NOT set here + ), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-1234", + user_id="admin-user", + ), + litellm_changed_by=None, + team_table=None, + ) + + # Verify generate_key_helper_fn was called + mock_generate_key.assert_awaited_once() + call_kwargs = mock_generate_key.call_args.kwargs + + # The key should have inherited key_budget_duration from the budget tier + assert call_kwargs.get("key_budget_duration") == "7d", ( + "key_budget_duration should be inherited from the linked budget tier " + f"but got: {call_kwargs.get('key_budget_duration')}" + ) + + # Verify the budget tier was looked up with the correct budget_id + mock_prisma.db.litellm_budgettable.find_unique.assert_awaited_once_with( + where={"budget_id": "7d-budget-tier"} + ) + + +@pytest.mark.asyncio +async def test_key_does_not_override_explicit_budget_duration(): + """ + Test that when a key is created with both budget_id and an explicit + budget_duration, the explicit budget_duration takes precedence over + the budget tier's budget_duration. + """ + from unittest.mock import AsyncMock, MagicMock, patch + + mock_prisma = MagicMock() + # The budget tier has budget_duration="7d" + mock_budget_row = MagicMock() + mock_budget_row.budget_duration = "7d" + mock_prisma.db.litellm_budgettable.find_unique = AsyncMock( + return_value=mock_budget_row + ) + + mock_generate_key = AsyncMock( + return_value={ + "key": "sk-test-key", + "expires": None, + "user_id": "test-user", + "team_id": None, + } + ) + + with patch( + "litellm.proxy.proxy_server.prisma_client", mock_prisma + ), patch( + "litellm.proxy.proxy_server.llm_router", None + ), patch( + "litellm.proxy.proxy_server.premium_user", False + ), patch( + "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" + ), patch( + "litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn", + mock_generate_key, + ): + await _common_key_generation_helper( + data=GenerateKeyRequest( + budget_id="7d-budget-tier", + max_budget=10.0, + budget_duration="30d", # explicit budget_duration should take precedence + ), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-1234", + user_id="admin-user", + ), + litellm_changed_by=None, + team_table=None, + ) + + mock_generate_key.assert_awaited_once() + call_kwargs = mock_generate_key.call_args.kwargs + + # The explicit budget_duration should take precedence + assert call_kwargs.get("key_budget_duration") == "30d", ( + "Explicit budget_duration should take precedence over the budget tier's value " + f"but got: {call_kwargs.get('key_budget_duration')}" + ) + + # The budget tier should NOT have been looked up since budget_duration was explicit + mock_prisma.db.litellm_budgettable.find_unique.assert_not_awaited() From cf14f0c8214951b63593e2e19e50652f253ba5c9 Mon Sep 17 00:00:00 2001 From: shivam Date: Sat, 14 Feb 2026 10:53:58 -0800 Subject: [PATCH 002/273] change logic to match Kriish's input --- .../key_management_endpoints.py | 11 ++---- .../test_key_management_endpoints.py | 38 +++++++------------ 2 files changed, 18 insertions(+), 31 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index b62ce329548..ab610a77b4d 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -592,15 +592,12 @@ async def _common_key_generation_helper( # noqa: PLR0915 if _budget_id is not None: data_json["budget_id"] = _budget_id + # Only set budget_duration on key when explicitly provided. Keys with budget_id + # but no explicit budget_duration follow their linked budget tier's schedule; + # reset_budget_for_keys_linked_to_budgets() resets them when the tier resets. + # This avoids duplicating budget_duration on keys so tier updates apply automatically. if "budget_duration" in data_json: data_json["key_budget_duration"] = data_json.pop("budget_duration", None) - elif _budget_id is not None and prisma_client is not None: - # Inherit budget_duration from linked budget tier if not explicitly set on the key - budget_row = await prisma_client.db.litellm_budgettable.find_unique( - where={"budget_id": _budget_id} - ) - if budget_row is not None and budget_row.budget_duration is not None: - data_json["key_budget_duration"] = budget_row.budget_duration if user_api_key_dict.user_id is not None: data_json["created_by"] = user_api_key_dict.user_id diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 472504871ed..21be2ad69e8 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -5562,26 +5562,19 @@ async def test_validate_key_list_check_key_hash_not_found(): @pytest.mark.asyncio -async def test_key_inherits_budget_duration_from_budget_tier(): +async def test_key_with_budget_id_does_not_store_budget_duration(): """ - Test that when a key is created with budget_id pointing to a budget tier - that has budget_duration, the key inherits budget_duration from the tier - even when budget_duration is not explicitly set on the key request. + Test that when a key is created with budget_id but without explicit + budget_duration, the key does NOT get budget_duration stored on it. - This verifies the fix for the bug where keys created with budget_id - would have null budget_duration and budget_reset_at, causing the - budget reset job to never reset their spend. + Keys with budget_id follow their linked budget tier's reset schedule; + reset_budget_for_keys_linked_to_budgets() resets them when the tier resets. + This avoids duplicating budget_duration on keys so tier updates apply + automatically to all linked keys. """ from unittest.mock import AsyncMock, MagicMock, patch - # Mock the budget tier lookup to return a budget with budget_duration="7d" - mock_budget_row = MagicMock() - mock_budget_row.budget_duration = "7d" - mock_prisma = MagicMock() - mock_prisma.db.litellm_budgettable.find_unique = AsyncMock( - return_value=mock_budget_row - ) mock_generate_key = AsyncMock( return_value={ @@ -5619,20 +5612,17 @@ async def test_key_inherits_budget_duration_from_budget_tier(): team_table=None, ) - # Verify generate_key_helper_fn was called mock_generate_key.assert_awaited_once() call_kwargs = mock_generate_key.call_args.kwargs - # The key should have inherited key_budget_duration from the budget tier - assert call_kwargs.get("key_budget_duration") == "7d", ( - "key_budget_duration should be inherited from the linked budget tier " - f"but got: {call_kwargs.get('key_budget_duration')}" + # Key should NOT have key_budget_duration - it follows the budget tier's schedule + assert call_kwargs.get("key_budget_duration") is None, ( + "key_budget_duration should be None for budget-linked keys without explicit " + f"budget_duration; got: {call_kwargs.get('key_budget_duration')}" ) - # Verify the budget tier was looked up with the correct budget_id - mock_prisma.db.litellm_budgettable.find_unique.assert_awaited_once_with( - where={"budget_id": "7d-budget-tier"} - ) + # No budget tier lookup - we don't copy budget_duration onto the key + mock_prisma.db.litellm_budgettable.find_unique.assert_not_called() @pytest.mark.asyncio @@ -5698,4 +5688,4 @@ async def test_key_does_not_override_explicit_budget_duration(): ) # The budget tier should NOT have been looked up since budget_duration was explicit - mock_prisma.db.litellm_budgettable.find_unique.assert_not_awaited() + mock_prisma.db.litellm_budgettable.find_unique.assert_not_called() From f79a8f7809ceaa1b2650c98ebdb95bc126a5beac Mon Sep 17 00:00:00 2001 From: Shivam Rawat <161387515+shivamrawat1@users.noreply.github.com> Date: Sat, 14 Feb 2026 14:42:17 -0800 Subject: [PATCH 003/273] Update litellm/proxy/common_utils/reset_budget_job.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- litellm/proxy/common_utils/reset_budget_job.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 6f038d127f6..c63309da1dc 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -95,11 +95,13 @@ class ResetBudgetJob: where={ "budget_id": {"in": budget_ids}, "budget_duration": None, # only keys without their own reset schedule + "spend": {"gt": 0}, # only reset keys that have accumulated spend }, data={ "spend": 0, }, ) + ) async def reset_budget_for_litellm_budget_table(self): """ From ea87dd216281155bc60da8a2751ae69122a97cce Mon Sep 17 00:00:00 2001 From: Shivam Rawat <161387515+shivamrawat1@users.noreply.github.com> Date: Sat, 14 Feb 2026 19:01:35 -0800 Subject: [PATCH 004/273] Update litellm/proxy/common_utils/reset_budget_job.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- litellm/proxy/common_utils/reset_budget_job.py | 1 - 1 file changed, 1 deletion(-) diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index c63309da1dc..b926fe28bed 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -101,7 +101,6 @@ class ResetBudgetJob: "spend": 0, }, ) - ) async def reset_budget_for_litellm_budget_table(self): """ From 2d9508ec96253557c97e6d6cc229d0a1d9e702d8 Mon Sep 17 00:00:00 2001 From: Shivam Rawat <161387515+shivamrawat1@users.noreply.github.com> Date: Sat, 14 Feb 2026 19:09:26 -0800 Subject: [PATCH 005/273] Update litellm/proxy/common_utils/reset_budget_job.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- litellm/proxy/common_utils/reset_budget_job.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index b926fe28bed..4933e679d71 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -76,13 +76,14 @@ class ResetBudgetJob: Resets the spend for keys linked to budget tiers that are being reset. This handles keys that have budget_id but no budget_duration set on the key - itself (e.g. keys created before the fix to inherit budget_duration from - the linked budget tier). + itself. Keys with budget_id rely on their linked budget tier's reset schedule + rather than having their own budget_duration. Keys that have their own budget_duration are already handled by reset_budget_for_litellm_keys() and are excluded here to avoid double-resetting. """ + """ budget_ids = [ budget.budget_id for budget in budgets_to_reset From 0117b35a6bc1c2877f074794664e9ece1294114d Mon Sep 17 00:00:00 2001 From: shivam Date: Mon, 16 Feb 2026 09:59:19 -0800 Subject: [PATCH 006/273] added more tests, fixed tests --- .../proxy/common_utils/reset_budget_job.py | 3 +- .../test_proxy_budget_reset.py | 16 +++++++ .../common_utils/test_reset_budget_job.py | 43 +++++++++++++++++++ 3 files changed, 60 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 4933e679d71..8ce73d29c84 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -83,7 +83,6 @@ class ResetBudgetJob: reset_budget_for_litellm_keys() and are excluded here to avoid double-resetting. """ - """ budget_ids = [ budget.budget_id for budget in budgets_to_reset @@ -617,4 +616,4 @@ class ResetBudgetJob: await ResetBudgetJob._reset_budget_common( item=key, current_time=current_time, item_type="key" ) - return key + return key \ No newline at end of file diff --git a/tests/litellm_utils_tests/test_proxy_budget_reset.py b/tests/litellm_utils_tests/test_proxy_budget_reset.py index 7cddde30421..34423a88da4 100644 --- a/tests/litellm_utils_tests/test_proxy_budget_reset.py +++ b/tests/litellm_utils_tests/test_proxy_budget_reset.py @@ -229,6 +229,10 @@ async def test_reset_budget_endusers_partial_failure(): prisma_client.get_data.side_effect = get_data_mock prisma_client.update_data = AsyncMock() + # Mock db.litellm_verificationtoken.update_many (used by reset_budget_for_keys_linked_to_budgets) + prisma_client.db.litellm_verificationtoken.update_many = AsyncMock( + return_value={"count": 0} + ) proxy_logging_obj = MagicMock() proxy_logging_obj.service_logging_obj = MagicMock() @@ -389,6 +393,10 @@ async def test_reset_budget_continues_other_categories_on_failure(): prisma_client.get_data = AsyncMock(side_effect=fake_get_data) prisma_client.update_data = AsyncMock() + # Mock db.litellm_verificationtoken.update_many (used by reset_budget_for_keys_linked_to_budgets) + prisma_client.db.litellm_verificationtoken.update_many = AsyncMock( + return_value={"count": 0} + ) proxy_logging_obj = MagicMock() proxy_logging_obj.service_logging_obj = MagicMock() @@ -863,6 +871,10 @@ async def test_service_logger_endusers_success(): prisma_client = MagicMock() prisma_client.get_data = AsyncMock(side_effect=fake_get_data) prisma_client.update_data = AsyncMock() + # Mock db.litellm_verificationtoken.update_many (used by reset_budget_for_keys_linked_to_budgets) + prisma_client.db.litellm_verificationtoken.update_many = AsyncMock( + return_value={"count": 0} + ) proxy_logging_obj = MagicMock() proxy_logging_obj.service_logging_obj = MagicMock() @@ -938,6 +950,10 @@ async def test_service_logger_endusers_failure(): prisma_client = MagicMock() prisma_client.get_data = AsyncMock(side_effect=fake_get_data) prisma_client.update_data = AsyncMock() + # Mock db.litellm_verificationtoken.update_many (used by reset_budget_for_keys_linked_to_budgets) + prisma_client.db.litellm_verificationtoken.update_many = AsyncMock( + return_value={"count": 0} + ) proxy_logging_obj = MagicMock() proxy_logging_obj.service_logging_obj = MagicMock() diff --git a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py index f63c77c1fc8..f975460836a 100644 --- a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py +++ b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py @@ -382,6 +382,49 @@ def test_reset_budget_for_keys_linked_to_budgets(reset_budget_job, mock_prisma_c assert call["data"]["spend"] == 0 +def test_reset_budget_for_keys_linked_to_budgets_excludes_keys_with_own_budget_duration( + reset_budget_job, mock_prisma_client +): + """ + Test that keys with BOTH budget_id AND budget_duration are excluded from + reset_budget_for_keys_linked_to_budgets. Such keys have their own reset + schedule and are handled only by reset_budget_for_litellm_keys(). The + budget_duration=None filter ensures they are NOT double-reset when the + linked budget tier expires. + """ + now = datetime.now(timezone.utc) + + test_budget = type( + "LiteLLM_BudgetTableFull", + (), + { + "max_budget": 10.0, + "budget_duration": "7d", + "budget_reset_at": now - timedelta(hours=1), + "budget_id": "7d-budget-tier", + "created_at": now - timedelta(days=7), + }, + ) + + budgets_to_reset = [test_budget] + + asyncio.run( + reset_budget_job.reset_budget_for_keys_linked_to_budgets( + budgets_to_reset=budgets_to_reset + ) + ) + + calls = mock_prisma_client.db.litellm_verificationtoken.update_many_calls + assert len(calls) == 1 + call = calls[0] + + # Critical: budget_duration must be None so keys with their own budget_duration + # (e.g. key has budget_id="X" AND budget_duration=60) are excluded. + # Those keys are reset only by reset_budget_for_litellm_keys() - no double-reset. + assert call["where"]["budget_duration"] is None + assert call["where"]["budget_id"] == {"in": ["7d-budget-tier"]} + + def test_reset_budget_for_keys_linked_to_budgets_empty( reset_budget_job, mock_prisma_client ): From 9f7f19067079977ae6d8fd3ecbd82a6f39362f02 Mon Sep 17 00:00:00 2001 From: shivam Date: Mon, 16 Feb 2026 10:24:06 -0800 Subject: [PATCH 007/273] resolved greptile comment --- tests/litellm_utils_tests/test_proxy_budget_reset.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/litellm_utils_tests/test_proxy_budget_reset.py b/tests/litellm_utils_tests/test_proxy_budget_reset.py index 34423a88da4..34b2043261c 100644 --- a/tests/litellm_utils_tests/test_proxy_budget_reset.py +++ b/tests/litellm_utils_tests/test_proxy_budget_reset.py @@ -1042,6 +1042,9 @@ async def test_reset_budget_for_litellm_team_members_called(): prisma_client.db.litellm_teammembership.update_many = AsyncMock( return_value={"count": 2} ) + prisma_client.db.litellm_verificationtoken.update_many = AsyncMock( + return_value={"count": 0} + ) proxy_logging_obj = MagicMock() proxy_logging_obj.service_logging_obj = MagicMock() From 8d5db4f712cf94eeacee130eb3557b910155096d Mon Sep 17 00:00:00 2001 From: jtsaw <166962251+jtsaw@users.noreply.github.com> Date: Tue, 17 Feb 2026 21:10:50 -0800 Subject: [PATCH 008/273] fix handling of ResponseApplyPatchToolCall in completion bridge (#20913) * fix handling of ResponseApplyPatchToolCall in completion bridge * refactor * style: fix black formatting * fix: clean up lint errors in test file (unused imports, print statements, formatting) * refactor: extract _map_optional_params_to_responses_api to fix PLR0915 * what * this linter cannot be me * revert cause idk what's going on * weird * idk why this got removed * revert more stuff * revert pt 3 --- .../transformation.py | 19 +- .../transformation.py | 77 +++++--- ...responses_transformation_transformation.py | 174 ++++++++++++++++-- 3 files changed, 225 insertions(+), 45 deletions(-) diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index e546a0dbb02..5de9a489854 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -401,6 +401,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): ResponseOutputMessage, ResponseReasoningItem, ) + from openai.types.responses.response_output_item import ResponseApplyPatchToolCall from litellm.types.utils import Choices, Message @@ -457,6 +458,18 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): accumulated_tool_calls.append(tool_call_dict) tool_call_index += 1 + elif isinstance(item, ResponseApplyPatchToolCall): + from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, + ) + + tool_call_dict = LiteLLMCompletionResponsesConfig.convert_apply_patch_tool_call_to_chat_completion_tool_call( + tool_call_item=item, + index=tool_call_index, + ) + accumulated_tool_calls.append(tool_call_dict) + tool_call_index += 1 + elif isinstance(item, dict) and handle_raw_dict_callback is not None: # Handle raw dict responses (e.g., from GPT-5 Codex) choice, index = handle_raw_dict_callback(item=item, index=index) @@ -533,7 +546,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): raw_response.usage ), ) - + # Preserve hidden params from the ResponsesAPIResponse, especially the headers # which contain important provider information like x-request-id raw_response_hidden_params = getattr(raw_response, "_hidden_params", {}) @@ -550,7 +563,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): model_response._hidden_params[key] = merged_headers else: model_response._hidden_params[key] = value - + return model_response def get_model_response_iterator( @@ -855,7 +868,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): return {"format": {"type": "text"}} return None - + @staticmethod def _convert_annotations_to_chat_format( annotations: Optional[List[Any]], diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index b8379b28c30..8daa8e49d1e 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -291,14 +291,14 @@ class LiteLLMCompletionResponsesConfig: ) _messages = litellm_completion_request.get("messages") or [] session_messages = chat_completion_session.get("messages") or [] - + # If session messages are empty (e.g., no database in test environment), # we still need to process the new input messages # Store original _messages before combining for safety check original_new_messages = _messages.copy() if _messages else [] - + combined_messages = session_messages + _messages - + # Fix: Ensure tool_results have corresponding tool_calls in previous assistant message # Pass tools parameter to help reconstruct tool_calls if not in cache tools = litellm_completion_request.get("tools") or [] @@ -306,7 +306,7 @@ class LiteLLMCompletionResponsesConfig: messages=combined_messages, tools=tools ) - + # Safety check: Ensure we don't end up with empty messages # This can happen when using previous_response_id without a database (e.g., in tests) # and session messages are empty but new input messages exist @@ -337,7 +337,7 @@ class LiteLLMCompletionResponsesConfig: model=litellm_completion_request.get("model", ""), llm_provider=litellm_completion_request.get("custom_llm_provider", ""), ) - + litellm_completion_request["messages"] = combined_messages litellm_completion_request["litellm_trace_id"] = chat_completion_session.get( "litellm_session_id" @@ -385,8 +385,8 @@ class LiteLLMCompletionResponsesConfig: ######################################################### # If Input Item is a Tool Call Output, add it to the tool_call_output_messages list - # preserving the ordering of tool call outputs. Some models require the tool - # result to immediately follow the assistant tool call. + # preserving the ordering of tool call outputs. Some models require the tool + # result to immediately follow the assistant tool call. ######################################################### if LiteLLMCompletionResponsesConfig._is_input_item_tool_call_output( input_item=_input @@ -743,47 +743,47 @@ class LiteLLMCompletionResponsesConfig: ) -> List[Union[AllMessageValues, GenericChatCompletionMessage, ChatCompletionResponseMessage]]: """ Ensure that tool_result messages have corresponding tool_calls in the previous assistant message. - + This is critical for Anthropic API which requires that each tool_result block has a corresponding tool_use block in the previous assistant message. - + Args: messages: List of messages that may include tool_result messages tools: Optional list of tools that can be used to reconstruct tool_calls if not in cache - + Returns: List of messages with tool_calls added to assistant messages when needed """ if not messages: return messages - + # Create a deep copy to avoid modifying the original import copy fixed_messages = copy.deepcopy(messages) messages_to_remove = [] - + # Count non-tool messages to avoid removing all messages # This prevents empty messages list when using previous_response_id without a database non_tool_messages_count = sum( 1 for msg in fixed_messages if msg.get("role") != "tool" ) - + for i, message in enumerate(fixed_messages): # Only process tool messages - check role first to narrow the type if message.get("role") != "tool": continue - + # At this point, we know it's a tool message, so it should have tool_call_id # Use get() with default to safely access tool_call_id tool_call_id_raw = message.get("tool_call_id") if isinstance(message, dict) else getattr(message, "tool_call_id", None) tool_call_id: str = ( str(tool_call_id_raw) if tool_call_id_raw is not None else "" ) - + prev_assistant_idx = LiteLLMCompletionResponsesConfig._find_previous_assistant_idx( fixed_messages, i ) - + # Try to recover empty tool_call_id from previous assistant message if not tool_call_id and prev_assistant_idx is not None: prev_assistant = fixed_messages[prev_assistant_idx] @@ -798,7 +798,7 @@ class LiteLLMCompletionResponsesConfig: message_dict["tool_call_id"] = tool_call_id elif hasattr(message, "tool_call_id"): setattr(message, "tool_call_id", tool_call_id) - + # Only remove messages with empty tool_call_id if we have other non-tool messages # This prevents ending up with an empty messages list when using previous_response_id # without a database (e.g., in tests where session messages are empty) @@ -810,7 +810,7 @@ class LiteLLMCompletionResponsesConfig: # If no non-tool messages, keep the tool message even with empty call_id # The API will return a proper error message about the missing tool_use block continue - + # Check if the previous assistant message has the corresponding tool_call # This needs to run for ALL tool messages with a valid tool_call_id, # not just those that had an empty tool_call_id initially @@ -819,12 +819,12 @@ class LiteLLMCompletionResponsesConfig: tool_calls = LiteLLMCompletionResponsesConfig._get_tool_calls_list( prev_assistant ) - + if not LiteLLMCompletionResponsesConfig._check_tool_call_exists( tool_calls, tool_call_id ): _tool_use_definition = TOOL_CALLS_CACHE.get_cache(key=tool_call_id) - + if not _tool_use_definition and tools: _tool_use_definition = ( LiteLLMCompletionResponsesConfig._reconstruct_tool_call_from_tools( @@ -849,11 +849,11 @@ class LiteLLMCompletionResponsesConfig: LiteLLMCompletionResponsesConfig._add_tool_call_to_assistant( prev_assistant, tool_call_chunk ) - + # Remove messages with empty tool_call_id that couldn't be fixed for idx in reversed(messages_to_remove): fixed_messages.pop(idx) - + return fixed_messages @staticmethod @@ -1454,6 +1454,39 @@ class LiteLLMCompletionResponsesConfig: return tool_call_dict + @staticmethod + def convert_apply_patch_tool_call_to_chat_completion_tool_call( + tool_call_item: Any, + index: int = 0, + ) -> Dict[str, Any]: + """ + Convert ResponseApplyPatchToolCall to ChatCompletionToolCallChunk format. + + The operation (create_file / update_file / delete_file) is serialised + as JSON so it appears in function.arguments, just like any other + tool call. + + Args: + tool_call_item: ResponseApplyPatchToolCall object with call_id and operation + index: The index of this tool call + + Returns: + Dictionary in ChatCompletionToolCallChunk format + """ + import json + + operation_dict = tool_call_item.operation.model_dump() + tool_call_dict: Dict[str, Any] = { + "id": tool_call_item.call_id, + "function": { + "name": "apply_patch", + "arguments": json.dumps(operation_dict), + }, + "type": "function", + "index": index, + } + return tool_call_dict + @staticmethod def transform_chat_completion_response_to_responses_api_response( request_input: Union[str, ResponseInputParam], diff --git a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py index f8a082ee30c..25e8a1f3304 100644 --- a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py +++ b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py @@ -1012,11 +1012,11 @@ def test_multiple_tool_calls_in_single_choice(): def test_map_reasoning_effort_adds_summary_detailed(): """ Test that _map_reasoning_effort behavior with reasoning_auto_summary flag. - + By default (flag=False), summary should NOT be added to avoid: 1. Breaking for users without verified OpenAI orgs (400 errors) 2. Making requests more expensive by including summary reasoning tokens - + When flag is enabled (flag=True or env var), summary="detailed" is added. """ import os @@ -1030,64 +1030,64 @@ def test_map_reasoning_effort_adds_summary_detailed(): # Test all string effort levels - DEFAULT BEHAVIOR (no summary) effort_levels = ["none", "low", "medium", "high", "xhigh", "minimal"] - + # Save original flag value original_flag = litellm.reasoning_auto_summary original_env = os.environ.get("LITELLM_REASONING_AUTO_SUMMARY") - + try: # Test 1: Default behavior (flag=False, no env var) - NO summary litellm.reasoning_auto_summary = False if "LITELLM_REASONING_AUTO_SUMMARY" in os.environ: del os.environ["LITELLM_REASONING_AUTO_SUMMARY"] - + for effort in effort_levels: result = handler._map_reasoning_effort(effort) - + assert result is not None, f"Result should not be None for effort={effort}" assert result["effort"] == effort, f"Effort should be {effort}" assert "summary" not in result, f"Summary should NOT be present by default for effort={effort}" - + print(f"✓ reasoning_effort='{effort}' correctly maps to effort='{effort}' (no summary by default)") - + # Test 2: With flag enabled - summary IS added litellm.reasoning_auto_summary = True - + for effort in effort_levels: result = handler._map_reasoning_effort(effort) - + assert result is not None, f"Result should not be None for effort={effort}" assert result["effort"] == effort, f"Effort should be {effort}" assert result["summary"] == "detailed", f"Summary should be 'detailed' when flag is enabled for effort={effort}" - + print(f"✓ reasoning_effort='{effort}' correctly maps to effort='{effort}', summary='detailed' (flag enabled)") - + # Test 3: With env var enabled (flag disabled) - summary IS added litellm.reasoning_auto_summary = False os.environ["LITELLM_REASONING_AUTO_SUMMARY"] = "true" - + result = handler._map_reasoning_effort("high") assert result["summary"] == "detailed", "Summary should be 'detailed' when env var is enabled" print("✓ LITELLM_REASONING_AUTO_SUMMARY env var works correctly") - + # Test 4: Dict input is passed through as-is (no modification) litellm.reasoning_auto_summary = False if "LITELLM_REASONING_AUTO_SUMMARY" in os.environ: del os.environ["LITELLM_REASONING_AUTO_SUMMARY"] - + dict_input = {"effort": "high", "summary": "custom_summary"} result_dict = handler._map_reasoning_effort(dict_input) assert result_dict["effort"] == "high" assert result_dict["summary"] == "custom_summary" print("✓ Dict input is passed through without modification") - + # Test 5: None/unknown values return None result_unknown = handler._map_reasoning_effort("unknown_value") assert result_unknown is None print("✓ Unknown reasoning_effort values return None") - + print("✓ All reasoning_effort behaviors work correctly with flag/env var control") - + finally: # Restore original values litellm.reasoning_auto_summary = original_flag @@ -1100,10 +1100,10 @@ def test_map_reasoning_effort_adds_summary_detailed(): def test_transform_response_preserves_annotations(): """ Test that annotations from Responses API are preserved when transforming to Chat Completions format. - + This is a regression test for the bug where annotations (like url_citation) were being dropped during the transformation from ResponsesAPIResponse to ModelResponse. - + The fix ensures annotations are extracted from ResponseOutputText content items and passed through to the Message object in the Chat Completions response. """ @@ -1278,3 +1278,137 @@ def test_transform_response_preserves_annotations(): assert result.usage.total_tokens == 30 print("✓ Annotations from Responses API are correctly preserved in Chat Completions format") + + +def test_apply_patch_tool_call_converted_to_chat_completion_tool_call(): + """ + Test that ResponseApplyPatchToolCall items from the Responses API are + correctly converted to ChatCompletions-style tool calls by the bridge. + + This is a regression test for a bug where litellm.completion() with a + responses/ model prefix crashed when the model returned an + apply_patch_call, because _convert_response_output_to_choices did not + handle ResponseApplyPatchToolCall items. The model DID use the tool, + but the bridge silently dropped it (or raised an error), while the + native litellm.responses() path worked correctly. + """ + import json + from unittest.mock import Mock + + from openai.types.responses.response_apply_patch_tool_call import ( + OperationCreateFile, + ) + from openai.types.responses.response_output_item import ( + ResponseApplyPatchToolCall, + ) + + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + LiteLLMResponsesTransformationHandler, + ) + from litellm.types.llms.openai import ( + InputTokensDetails, + OutputTokensDetails, + ResponseAPIUsage, + ResponsesAPIResponse, + ) + from litellm.types.utils import ModelResponse, Usage + + handler = LiteLLMResponsesTransformationHandler() + + # Build an apply_patch_call item like the model would return + operation = OperationCreateFile( + diff="--- /dev/null\n+++ b/hello.py\n@@ -0,0 +1 @@\n+print('hello world')\n", + path="hello.py", + type="create_file", + ) + apply_patch_item = ResponseApplyPatchToolCall( + id="apc_001", + call_id="call_patch_hello", + operation=operation, + status="completed", + type="apply_patch_call", + ) + + # Minimal usage + usage = ResponseAPIUsage( + input_tokens=30, + input_tokens_details=InputTokensDetails(cached_tokens=0), + output_tokens=40, + output_tokens_details=OutputTokensDetails(reasoning_tokens=0), + total_tokens=70, + ) + + raw_response = ResponsesAPIResponse( + id="resp_apply_patch_test", + created_at=1234567890, + error=None, + incomplete_details=None, + instructions=None, + metadata={}, + model="gpt-5.2-codex", + object="response", + output=[apply_patch_item], + parallel_tool_calls=True, + temperature=1.0, + tool_choice="auto", + tools=[], + top_p=1.0, + max_output_tokens=None, + previous_response_id=None, + reasoning=None, + status="completed", + text=None, + truncation="disabled", + usage=usage, + user=None, + store=True, + background=False, + ) + + model_response = ModelResponse( + id="chatcmpl-apply-patch", + created=1234567890, + model=None, + object="chat.completion", + choices=[], + usage=Usage(completion_tokens=0, prompt_tokens=0, total_tokens=0), + ) + + logging_obj = Mock() + + result = handler.transform_response( + model="gpt-5.2-codex", + raw_response=raw_response, + model_response=model_response, + logging_obj=logging_obj, + request_data={"model": "gpt-5.2-codex"}, + messages=[ + {"role": "system", "content": "You are a coding assistant."}, + {"role": "user", "content": "Create hello.py"}, + ], + optional_params={}, + litellm_params={}, + encoding=Mock(), + ) + + # Should have exactly one choice with finish_reason="tool_calls" + assert len(result.choices) == 1, f"Expected 1 choice, got {len(result.choices)}" + + choice = result.choices[0] + assert choice.finish_reason == "tool_calls" + + # The choice should contain one tool call for apply_patch + tool_calls = choice.message.tool_calls + assert tool_calls is not None, "tool_calls should not be None" + assert len(tool_calls) == 1, f"Expected 1 tool_call, got {len(tool_calls)}" + + tc = tool_calls[0] + assert tc["id"] == "call_patch_hello" + assert tc["type"] == "function" + assert tc["function"]["name"] == "apply_patch" + + # The operation should be serialised as JSON in arguments + args = json.loads(tc["function"]["arguments"]) + assert args["type"] == "create_file" + assert args["path"] == "hello.py" + assert "print('hello world')" in args["diff"] From ae613b2d36f92a700077884234b0076af67cfb85 Mon Sep 17 00:00:00 2001 From: Atharva Jaiswal <92455570+AtharvaJaiswal005@users.noreply.github.com> Date: Wed, 18 Feb 2026 12:28:08 +0530 Subject: [PATCH 009/273] fix(router): break retry loop on non-retryable errors (#21370) The retry loop in async_function_with_retries catches all exceptions blindly and continues retrying even for non-retryable errors like 400 ContextWindowExceeded or 404 NotFoundError. This causes the original retryable error to be raised instead of the actual non-retryable one. Changes: - Update original_exception to latest error on each retry attempt - Add should_retry_this_error() check inside the retry loop to break out immediately on non-retryable errors - Respect _retry_policy_applies precedence Fixes #21343 --- litellm/router.py | 22 ++ .../test_router_retry_non_retryable_errors.py | 251 ++++++++++++++++++ 2 files changed, 273 insertions(+) create mode 100644 tests/test_litellm/test_router_retry_non_retryable_errors.py diff --git a/litellm/router.py b/litellm/router.py index 888c97ca0b1..3fac761ce60 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -5149,6 +5149,10 @@ class Router: return response except Exception as e: + # Always track the latest error so we raise the most + # recent exception instead of the first one. + original_exception = e + ## LOGGING kwargs = self.log_retry(kwargs=kwargs, e=e) remaining_retries = num_retries - current_attempt - 1 @@ -5163,6 +5167,24 @@ class Router: ) else: _healthy_deployments = [] + + # Check if this error is non-retryable (e.g., 400 context + # window exceeded). If so, raise immediately instead of + # continuing the retry loop. Respect retry policy + # precedence - only check when no retry policy applies. + if not _retry_policy_applies: + try: + self.should_retry_this_error( + error=e, + healthy_deployments=_healthy_deployments, + all_deployments=_all_deployments, + context_window_fallbacks=context_window_fallbacks, + regular_fallbacks=fallbacks, + content_policy_fallbacks=content_policy_fallbacks, + ) + except Exception: + raise e + _timeout = self._time_to_sleep_before_retry( e=e, remaining_retries=remaining_retries, diff --git a/tests/test_litellm/test_router_retry_non_retryable_errors.py b/tests/test_litellm/test_router_retry_non_retryable_errors.py new file mode 100644 index 00000000000..20a1c979a04 --- /dev/null +++ b/tests/test_litellm/test_router_retry_non_retryable_errors.py @@ -0,0 +1,251 @@ +""" +Test that the Router retry loop correctly handles non-retryable errors. + +Verifies that: +1. Non-retryable errors (e.g., 400 ContextWindowExceeded) inside the retry loop + break out immediately instead of being swallowed. +2. original_exception is updated to the latest error, not stuck on the first. +3. Retryable errors (e.g., 429 RateLimitError) still retry normally. + +Regression tests for https://github.com/BerriAI/litellm/issues/21343 +""" + +from unittest.mock import AsyncMock, patch + +import pytest + +import litellm +from litellm import Router + + +def _make_rate_limit_error(message="Rate limited"): + """Create a RateLimitError for testing.""" + return litellm.RateLimitError( + message=message, + llm_provider="bedrock", + model="anthropic.claude-v2", + ) + + +def _make_context_window_error(message="prompt is too long: 1205821 tokens > 200000"): + """Create a ContextWindowExceededError for testing.""" + return litellm.ContextWindowExceededError( + message=message, + llm_provider="vertex_ai", + model="claude-3-opus", + ) + + +def _make_bad_request_error(message="Invalid request"): + """Create a BadRequestError for testing.""" + return litellm.BadRequestError( + message=message, + llm_provider="openai", + model="gpt-4", + ) + + +def _make_not_found_error(message="Model not found"): + """Create a NotFoundError for testing.""" + return litellm.NotFoundError( + message=message, + llm_provider="openai", + model="gpt-99", + ) + + +def _create_router(num_retries=2): + """Create a Router with two deployments for testing.""" + return Router( + model_list=[ + { + "model_name": "test-model", + "litellm_params": { + "model": "openai/gpt-4", + "api_key": "fake-key-1", + }, + }, + { + "model_name": "test-model", + "litellm_params": { + "model": "openai/gpt-4", + "api_key": "fake-key-2", + }, + }, + ], + num_retries=num_retries, + ) + + +def _base_kwargs(): + """Return kwargs required by async_function_with_retries.""" + return { + "model": "test-model", + "messages": [{"role": "user", "content": "test"}], + "original_function": AsyncMock(), + "metadata": {}, + } + + +@pytest.mark.asyncio +async def test_non_retryable_error_in_retry_loop_raises_immediately(): + """ + When a non-retryable error (400 ContextWindowExceeded) occurs inside the + retry loop, the router should raise it immediately instead of swallowing it + and raising the original error. + + Scenario: First call -> 429, Retry -> 400 (non-retryable) + Expected: ContextWindowExceededError is raised, NOT RateLimitError + """ + router = _create_router(num_retries=2) + + rate_limit_error = _make_rate_limit_error() + context_window_error = _make_context_window_error() + + call_count = 0 + + async def mock_make_call(*args, **kwargs): + nonlocal call_count + call_count += 1 + if call_count == 1: + raise rate_limit_error + else: + raise context_window_error + + with patch.object(router, "make_call", side_effect=mock_make_call), \ + patch.object(router, "_async_get_healthy_deployments", + return_value=(["d1", "d2"], ["d1", "d2"])), \ + patch.object(router, "_time_to_sleep_before_retry", return_value=0), \ + patch.object(router, "log_retry", side_effect=lambda kwargs, e: kwargs): + with pytest.raises(litellm.ContextWindowExceededError): + await router.async_function_with_retries( + num_retries=2, + **_base_kwargs(), + ) + + +@pytest.mark.asyncio +async def test_bad_request_error_in_retry_loop_raises_immediately(): + """ + A generic 400 BadRequestError inside the retry loop should also break out + immediately since 400 is not retryable. + """ + router = _create_router(num_retries=2) + + rate_limit_error = _make_rate_limit_error() + bad_request_error = _make_bad_request_error() + + call_count = 0 + + async def mock_make_call(*args, **kwargs): + nonlocal call_count + call_count += 1 + if call_count == 1: + raise rate_limit_error + else: + raise bad_request_error + + with patch.object(router, "make_call", side_effect=mock_make_call), \ + patch.object(router, "_async_get_healthy_deployments", + return_value=(["d1", "d2"], ["d1", "d2"])), \ + patch.object(router, "_time_to_sleep_before_retry", return_value=0), \ + patch.object(router, "log_retry", side_effect=lambda kwargs, e: kwargs): + with pytest.raises(litellm.BadRequestError): + await router.async_function_with_retries( + num_retries=2, + **_base_kwargs(), + ) + + +@pytest.mark.asyncio +async def test_original_exception_updated_to_latest_error(): + """ + When all retries are exhausted with retryable errors, the LAST error + should be raised, not the first one. + """ + router = _create_router(num_retries=2) + + call_count = 0 + + async def mock_make_call(*args, **kwargs): + nonlocal call_count + call_count += 1 + raise _make_rate_limit_error(f"Rate limit attempt {call_count}") + + with patch.object(router, "make_call", side_effect=mock_make_call), \ + patch.object(router, "_async_get_healthy_deployments", + return_value=(["d1", "d2"], ["d1", "d2"])), \ + patch.object(router, "_time_to_sleep_before_retry", return_value=0), \ + patch.object(router, "log_retry", side_effect=lambda kwargs, e: kwargs): + with pytest.raises(litellm.RateLimitError) as exc_info: + await router.async_function_with_retries( + num_retries=2, + **_base_kwargs(), + ) + # Should be the LAST error, not the first + assert "Rate limit attempt 3" in str(exc_info.value) + + +@pytest.mark.asyncio +async def test_retryable_errors_still_retry_normally(): + """ + Retryable errors (429 RateLimitError) should still be retried the + configured number of times before raising. + """ + router = _create_router(num_retries=3) + + call_count = 0 + + async def mock_make_call(*args, **kwargs): + nonlocal call_count + call_count += 1 + raise _make_rate_limit_error(f"Rate limit attempt {call_count}") + + with patch.object(router, "make_call", side_effect=mock_make_call), \ + patch.object(router, "_async_get_healthy_deployments", + return_value=(["d1", "d2"], ["d1", "d2"])), \ + patch.object(router, "_time_to_sleep_before_retry", return_value=0), \ + patch.object(router, "log_retry", side_effect=lambda kwargs, e: kwargs): + with pytest.raises(litellm.RateLimitError): + await router.async_function_with_retries( + num_retries=3, + **_base_kwargs(), + ) + + # Initial call + 3 retries = 4 total calls + assert call_count == 4 + + +@pytest.mark.asyncio +async def test_not_found_error_in_retry_loop_raises_immediately(): + """ + A 404 NotFoundError inside the retry loop should break out immediately. + """ + router = _create_router(num_retries=2) + + rate_limit_error = _make_rate_limit_error() + not_found_error = _make_not_found_error() + + call_count = 0 + + async def mock_make_call(*args, **kwargs): + nonlocal call_count + call_count += 1 + if call_count == 1: + raise rate_limit_error + else: + raise not_found_error + + with patch.object(router, "make_call", side_effect=mock_make_call), \ + patch.object(router, "_async_get_healthy_deployments", + return_value=(["d1", "d2"], ["d1", "d2"])), \ + patch.object(router, "_time_to_sleep_before_retry", return_value=0), \ + patch.object(router, "log_retry", side_effect=lambda kwargs, e: kwargs): + with pytest.raises(litellm.NotFoundError): + await router.async_function_with_retries( + num_retries=2, + **_base_kwargs(), + ) + + # Only 2 calls: initial + first retry that hits non-retryable + assert call_count == 2 From 42afba9cdd3ec78270c84da0e6e915d158105434 Mon Sep 17 00:00:00 2001 From: Atharva Jaiswal <92455570+AtharvaJaiswal005@users.noreply.github.com> Date: Wed, 18 Feb 2026 12:29:01 +0530 Subject: [PATCH 010/273] Fix invalid OpenAPI schema for /spend/calculate and /credentials endpoints (#21369) - /spend/calculate: wrap response in proper OpenAPI 3.x content structure - /credentials: split stacked route decorators into separate handlers to eliminate path parameter conflict between by_name and by_model routes --- .../proxy/credential_endpoints/endpoints.py | 97 ++++++------ .../spend_management_endpoints.py | 20 ++- .../proxy/test_openapi_schema_validation.py | 142 ++++++++++++++++++ 3 files changed, 209 insertions(+), 50 deletions(-) create mode 100644 tests/test_litellm/proxy/test_openapi_schema_validation.py diff --git a/litellm/proxy/credential_endpoints/endpoints.py b/litellm/proxy/credential_endpoints/endpoints.py index 9f228bb1184..5fa9546e006 100644 --- a/litellm/proxy/credential_endpoints/endpoints.py +++ b/litellm/proxy/credential_endpoints/endpoints.py @@ -142,17 +142,47 @@ async def get_credentials( tags=["credential management"], response_model=CredentialItem, ) +async def get_credential_by_name( + request: Request, + fastapi_response: Response, + credential_name: str = Path(..., description="The credential name, percent-decoded; may contain slashes"), + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + [BETA] endpoint. This might change unexpectedly. + """ + try: + for credential in litellm.credential_list: + if credential.credential_name == credential_name: + masked_credential = CredentialItem( + credential_name=credential.credential_name, + credential_values=_get_masked_values( + credential.credential_values, + unmasked_length=4, + number_of_asterisks=4, + ), + credential_info=credential.credential_info, + ) + return masked_credential + raise HTTPException( + status_code=404, + detail="Credential not found. Got credential name: " + credential_name, + ) + except Exception as e: + verbose_proxy_logger.exception(e) + raise handle_exception_on_proxy(e) + + @router.get( "/credentials/by_model/{model_id}", dependencies=[Depends(user_api_key_auth)], tags=["credential management"], response_model=CredentialItem, ) -async def get_credential( +async def get_credential_by_model( request: Request, fastapi_response: Response, - credential_name: str = Path(..., description="The credential name, percent-decoded; may contain slashes"), - model_id: Optional[str] = None, + model_id: str = Path(..., description="The model ID to look up credentials for"), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ @@ -161,48 +191,25 @@ async def get_credential( from litellm.proxy.proxy_server import llm_router try: - if model_id: - if llm_router is None: - raise HTTPException(status_code=500, detail="LLM router not found") - model = llm_router.get_deployment(model_id) - if model is None: - raise HTTPException(status_code=404, detail="Model not found") - credential_values = llm_router.get_deployment_credentials(model_id) - if credential_values is None: - raise HTTPException(status_code=404, detail="Model not found") - masked_credential_values = _get_masked_values( - credential_values, - unmasked_length=4, - number_of_asterisks=4, - ) - credential = CredentialItem( - credential_name="{}-credential-{}".format(model.model_name, model_id), - credential_values=masked_credential_values, - credential_info={}, - ) - # return credential object - return credential - elif credential_name: - for credential in litellm.credential_list: - if credential.credential_name == credential_name: - masked_credential = CredentialItem( - credential_name=credential.credential_name, - credential_values=_get_masked_values( - credential.credential_values, - unmasked_length=4, - number_of_asterisks=4, - ), - credential_info=credential.credential_info, - ) - return masked_credential - raise HTTPException( - status_code=404, - detail="Credential not found. Got credential name: " + credential_name, - ) - else: - raise HTTPException( - status_code=404, detail="Credential name or model ID required" - ) + if llm_router is None: + raise HTTPException(status_code=500, detail="LLM router not found") + model = llm_router.get_deployment(model_id) + if model is None: + raise HTTPException(status_code=404, detail="Model not found") + credential_values = llm_router.get_deployment_credentials(model_id) + if credential_values is None: + raise HTTPException(status_code=404, detail="Model not found") + masked_credential_values = _get_masked_values( + credential_values, + unmasked_length=4, + number_of_asterisks=4, + ) + credential = CredentialItem( + credential_name="{}-credential-{}".format(model.model_name, model_id), + credential_values=masked_credential_values, + credential_info={}, + ) + return credential except Exception as e: verbose_proxy_logger.exception(e) raise handle_exception_on_proxy(e) diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 08aaa851691..92770a5c803 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -1461,11 +1461,21 @@ async def _get_spend_report_for_time_range( dependencies=[Depends(user_api_key_auth)], responses={ 200: { - "cost": { - "description": "The calculated cost", - "example": 0.0, - "type": "float", - } + "description": "The calculated cost", + "content": { + "application/json": { + "schema": { + "type": "object", + "properties": { + "cost": { + "type": "number", + "description": "The calculated cost", + "example": 0.0, + } + }, + } + } + }, } }, ) diff --git a/tests/test_litellm/proxy/test_openapi_schema_validation.py b/tests/test_litellm/proxy/test_openapi_schema_validation.py new file mode 100644 index 00000000000..aafe08f3033 --- /dev/null +++ b/tests/test_litellm/proxy/test_openapi_schema_validation.py @@ -0,0 +1,142 @@ +""" +Test that the OpenAPI schema generated by FastAPI is valid for specific endpoints. + +Validates fixes for: +- /spend/calculate response schema (must use proper OpenAPI 3.x content wrapper) +- /credentials/by_model/{model_id} path parameter (must not leak credential_name) + +Related issue: https://github.com/BerriAI/litellm/issues/21305 +""" + +import pytest + + +class TestSpendCalculateOpenAPISchema: + """Test /spend/calculate response schema is valid OpenAPI 3.x.""" + + def test_response_schema_has_description(self): + """The 200 response must have a 'description' field per OpenAPI 3.x spec.""" + from litellm.proxy.spend_tracking.spend_management_endpoints import router + + for route in router.routes: + if hasattr(route, "path") and route.path == "/spend/calculate": + responses = route.responses or {} + response_200 = responses.get(200, {}) + assert "description" in response_200, ( + "/spend/calculate 200 response must have a 'description' field" + ) + break + else: + pytest.fail("/spend/calculate route not found in router") + + def test_response_schema_has_content_wrapper(self): + """The 200 response must use 'content' wrapper, not bare properties.""" + from litellm.proxy.spend_tracking.spend_management_endpoints import router + + for route in router.routes: + if hasattr(route, "path") and route.path == "/spend/calculate": + responses = route.responses or {} + response_200 = responses.get(200, {}) + # Must NOT have 'cost' as a top-level key (invalid OpenAPI) + assert "cost" not in response_200, ( + "/spend/calculate 200 response must not have 'cost' as a " + "top-level property - use 'content' wrapper instead" + ) + # Must have 'content' wrapper + assert "content" in response_200, ( + "/spend/calculate 200 response must have a 'content' field" + ) + content = response_200["content"] + assert "application/json" in content + assert "schema" in content["application/json"] + break + else: + pytest.fail("/spend/calculate route not found in router") + + +class TestCredentialEndpointsOpenAPISchema: + """Test /credentials endpoints have correct path parameters.""" + + def test_by_name_and_by_model_are_separate_handlers(self): + """ + /credentials/by_name/{credential_name} and /credentials/by_model/{model_id} + must be separate handler functions so each only declares its own path params. + """ + from litellm.proxy.credential_endpoints.endpoints import router + + by_name_routes = [] + by_model_routes = [] + for route in router.routes: + if not hasattr(route, "path"): + continue + if "by_name" in route.path: + by_name_routes.append(route) + elif "by_model" in route.path: + by_model_routes.append(route) + + assert len(by_name_routes) == 1, "Expected exactly one by_name route" + assert len(by_model_routes) == 1, "Expected exactly one by_model route" + + # They must be different endpoint functions + by_name_endpoint = by_name_routes[0].endpoint + by_model_endpoint = by_model_routes[0].endpoint + assert by_name_endpoint is not by_model_endpoint, ( + "by_name and by_model must be separate handler functions " + "to avoid path parameter conflicts in OpenAPI spec" + ) + + def test_by_model_route_does_not_require_credential_name(self): + """ + The /credentials/by_model/{model_id} route must NOT have + credential_name as a parameter. + """ + import inspect + from litellm.proxy.credential_endpoints.endpoints import ( + get_credential_by_model, + ) + + sig = inspect.signature(get_credential_by_model) + param_names = list(sig.parameters.keys()) + assert "credential_name" not in param_names, ( + "get_credential_by_model must not have a credential_name parameter" + ) + + def test_by_name_route_does_not_require_model_id(self): + """ + The /credentials/by_name/{credential_name} route must NOT have + model_id as a parameter. + """ + import inspect + from litellm.proxy.credential_endpoints.endpoints import ( + get_credential_by_name, + ) + + sig = inspect.signature(get_credential_by_name) + param_names = list(sig.parameters.keys()) + assert "model_id" not in param_names, ( + "get_credential_by_name must not have a model_id parameter" + ) + + def test_by_model_has_model_id_path_param(self): + """The by_model handler must accept model_id as a path parameter.""" + import inspect + from litellm.proxy.credential_endpoints.endpoints import ( + get_credential_by_model, + ) + + sig = inspect.signature(get_credential_by_model) + assert "model_id" in sig.parameters, ( + "get_credential_by_model must have a model_id parameter" + ) + + def test_by_name_has_credential_name_path_param(self): + """The by_name handler must accept credential_name as a path parameter.""" + import inspect + from litellm.proxy.credential_endpoints.endpoints import ( + get_credential_by_name, + ) + + sig = inspect.signature(get_credential_by_name) + assert "credential_name" in sig.parameters, ( + "get_credential_by_name must have a credential_name parameter" + ) From 518cd3ef60e5809947dbf2d262c7edd47782a9ba Mon Sep 17 00:00:00 2001 From: Dibyo Mukherjee Date: Thu, 5 Feb 2026 19:40:41 -0500 Subject: [PATCH 011/273] feat(ui): add key creation deep-links with SSO return URL support Enables deep-linking directly to the key creation modal with prefilled form data via URL parameters, including support for preserving these deep-links through SSO authentication flows. Key Creation Deep-links: - Auto-open key creation modal via ?create=true parameter - Prefill form fields from URL parameters (team_id, key_alias, models, etc.) - Role-based access control for auto-open (requires write access) - Race condition protection for redirect handling Example: /ui?create=true&team_id=abc&key_alias=my-key&models=gpt-4,claude-3 SSO Return URL Preservation: - Cookie-based return URL storage (works across ports for SSO flows) - URL validation to prevent open redirect attacks - Support for both dev and production environments Co-Authored-By: Claude Opus 4.5 --- .../(dashboard)/hooks/useAuthorized.test.ts | 12 +- .../app/(dashboard)/hooks/useAuthorized.ts | 45 +- .../src/app/login/LoginPage.tsx | 25 +- ui/litellm-dashboard/src/app/page.tsx | 136 +++++-- .../organisms/create_key_button.test.tsx | 367 ++++++++++++++--- .../organisms/create_key_button.tsx | 93 ++++- .../src/components/user_dashboard.tsx | 8 +- .../src/utils/returnUrlUtils.test.ts | 383 ++++++++++++++++++ .../src/utils/returnUrlUtils.ts | 304 ++++++++++++++ ui/litellm-dashboard/src/utils/roles.ts | 29 ++ .../tests/CreateKeyPage.expiredToken.test.tsx | 55 ++- 11 files changed, 1315 insertions(+), 142 deletions(-) create mode 100644 ui/litellm-dashboard/src/utils/returnUrlUtils.test.ts create mode 100644 ui/litellm-dashboard/src/utils/returnUrlUtils.ts diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.test.ts index 76a3129d6d7..5178aca0790 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.test.ts @@ -8,13 +8,14 @@ import useAuthorized from "./useAuthorized"; // Unmock useAuthorized to test the actual implementation vi.unmock("@/app/(dashboard)/hooks/useAuthorized"); -const { replaceMock, clearTokenCookiesMock, getProxyBaseUrlMock, getUiConfigMock, decodeTokenMock, checkTokenValidityMock } = vi.hoisted(() => ({ +const { replaceMock, clearTokenCookiesMock, getProxyBaseUrlMock, getUiConfigMock, decodeTokenMock, checkTokenValidityMock, buildLoginUrlWithReturnMock } = vi.hoisted(() => ({ replaceMock: vi.fn(), clearTokenCookiesMock: vi.fn(), getProxyBaseUrlMock: vi.fn(() => "http://proxy.example"), getUiConfigMock: vi.fn(), decodeTokenMock: vi.fn(), checkTokenValidityMock: vi.fn(), + buildLoginUrlWithReturnMock: vi.fn((baseUrl: string) => baseUrl), })); vi.mock("next/navigation", () => ({ @@ -49,6 +50,14 @@ vi.mock("@/utils/jwtUtils", async (importOriginal) => { }; }); +vi.mock("@/utils/returnUrlUtils", async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + buildLoginUrlWithReturn: buildLoginUrlWithReturnMock, + storeReturnUrl: vi.fn(), + }; +}); const createQueryClient = () => new QueryClient({ defaultOptions: { @@ -81,6 +90,7 @@ describe("useAuthorized", () => { getUiConfigMock.mockReset(); decodeTokenMock.mockReset(); checkTokenValidityMock.mockReset(); + buildLoginUrlWithReturnMock.mockClear(); clearCookie(); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.ts index 0b60971c1eb..8f8c403a4e9 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/useAuthorized.ts @@ -3,39 +3,12 @@ import { getProxyBaseUrl } from "@/components/networking"; import { clearTokenCookies, getCookie } from "@/utils/cookieUtils"; import { checkTokenValidity, decodeToken } from "@/utils/jwtUtils"; +import { buildLoginUrlWithReturn, storeReturnUrl } from "@/utils/returnUrlUtils"; import { useRouter } from "next/navigation"; -import { useEffect, useMemo } from "react"; +import { useCallback, useEffect, useMemo } from "react"; +import { formatUserRole } from "@/utils/roles"; import { useUIConfig } from "./uiConfig/useUIConfig"; -function formatUserRole(userRole: string) { - if (!userRole) { - return "Undefined Role"; - } - switch (userRole.toLowerCase()) { - case "app_owner": - return "App Owner"; - case "demo_app_owner": - return "App Owner"; - case "app_admin": - return "Admin"; - case "proxy_admin": - return "Admin"; - case "proxy_admin_viewer": - return "Admin Viewer"; - case "org_admin": - return "Org Admin"; - case "internal_user": - return "Internal User"; - case "internal_user_viewer": - case "internal_viewer": // TODO:remove if deprecated - return "Internal Viewer"; - case "app_user": - return "App User"; - default: - return "Unknown Role"; - } -} - const useAuthorized = () => { const router = useRouter(); const { data: uiConfig, isLoading: isUIConfigLoading } = useUIConfig(); @@ -47,6 +20,14 @@ const useAuthorized = () => { const isLoading = isUIConfigLoading; const isAuthorized = isTokenValid && !uiConfig?.admin_ui_disabled; + // Helper function to redirect to login while preserving the current URL + const redirectToLogin = useCallback(() => { + storeReturnUrl(); + const baseLoginUrl = `${getProxyBaseUrl()}/ui/login`; + const loginUrlWithReturn = buildLoginUrlWithReturn(baseLoginUrl); + router.replace(loginUrlWithReturn); + }, [router]); + // Single useEffect for all redirect logic useEffect(() => { if (isLoading) return; @@ -55,9 +36,9 @@ const useAuthorized = () => { if (token) { clearTokenCookies(); } - router.replace(`${getProxyBaseUrl()}/ui/login`); + redirectToLogin(); } - }, [isLoading, isAuthorized, token, router]); + }, [isLoading, isAuthorized, token, redirectToLogin]); return { isLoading, diff --git a/ui/litellm-dashboard/src/app/login/LoginPage.tsx b/ui/litellm-dashboard/src/app/login/LoginPage.tsx index a05fa4e214e..80372fcddca 100644 --- a/ui/litellm-dashboard/src/app/login/LoginPage.tsx +++ b/ui/litellm-dashboard/src/app/login/LoginPage.tsx @@ -6,6 +6,7 @@ import LoadingScreen from "@/components/common_components/LoadingScreen"; import { getProxyBaseUrl } from "@/components/networking"; import { getCookie } from "@/utils/cookieUtils"; import { isJwtExpired } from "@/utils/jwtUtils"; +import { consumeReturnUrl, getReturnUrl, isValidReturnUrl } from "@/utils/returnUrlUtils"; import { InfoCircleOutlined } from "@ant-design/icons"; import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; import { Alert, Button, Card, Form, Input, Popover, Space, Typography } from "antd"; @@ -33,12 +34,24 @@ function LoginPageContent() { const rawToken = getCookie("token"); if (rawToken && !isJwtExpired(rawToken)) { - router.replace(`${getProxyBaseUrl()}/ui`); + // User already logged in - redirect to return URL or default + const returnUrl = consumeReturnUrl(); + if (returnUrl) { + router.replace(returnUrl); + } else { + router.replace(`${getProxyBaseUrl()}/ui`); + } return; } if (uiConfig && uiConfig.auto_redirect_to_sso) { - router.push(`${getProxyBaseUrl()}/sso/key/generate`); + // For SSO, pass the return URL to the SSO endpoint + const returnUrl = getReturnUrl(); + let ssoUrl = `${getProxyBaseUrl()}/sso/key/generate`; + if (returnUrl && isValidReturnUrl(returnUrl)) { + ssoUrl += `?redirect_to=${encodeURIComponent(returnUrl)}`; + } + router.push(ssoUrl); return; } @@ -50,7 +63,13 @@ function LoginPageContent() { { username, password }, { onSuccess: (data) => { - router.push(data.redirect_url); + // Check if we have a return URL to use instead of the default redirect + const returnUrl = consumeReturnUrl(); + if (returnUrl) { + router.push(returnUrl); + } else { + router.push(data.redirect_url); + } }, }, ); diff --git a/ui/litellm-dashboard/src/app/page.tsx b/ui/litellm-dashboard/src/app/page.tsx index 258c2ccb0e0..26aebf3e3d2 100644 --- a/ui/litellm-dashboard/src/app/page.tsx +++ b/ui/litellm-dashboard/src/app/page.tsx @@ -23,7 +23,7 @@ import Navbar from "@/components/navbar"; import { getUiConfig, Organization, proxyBaseUrl, setGlobalLitellmHeaderName, getInProductNudgesCall } from "@/components/networking"; import NewUsagePage from "@/components/UsagePage/components/UsagePageView"; import OldTeams from "@/components/OldTeams"; -import { fetchUserModels } from "@/components/organisms/create_key_button"; +import { fetchUserModels, CreateKeyPrefillData } from "@/components/organisms/create_key_button"; import Organizations, { fetchOrganizations } from "@/components/organizations"; import PassThroughSettings from "@/components/pass_through_settings"; import PromptsPanel from "@/components/prompts"; @@ -43,11 +43,12 @@ import SpendLogsTable from "@/components/view_logs"; import ViewUserDashboard from "@/components/view_users"; import { ThemeProvider } from "@/contexts/ThemeContext"; import { isJwtExpired } from "@/utils/jwtUtils"; -import { isAdminRole } from "@/utils/roles"; +import { buildLoginUrlWithReturn, consumeReturnUrl, normalizeUrlForCompare, storeReturnUrl } from "@/utils/returnUrlUtils"; +import { formatUserRole, isAdminRole } from "@/utils/roles"; import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; import { jwtDecode } from "jwt-decode"; import { useSearchParams } from "next/navigation"; -import { Suspense, useEffect, useState } from "react"; +import { Suspense, useEffect, useMemo, useRef, useState } from "react"; import { ConfigProvider, theme } from "antd"; function getCookie(name: string) { @@ -67,35 +68,6 @@ function deleteCookie(name: string, path = "/") { document.cookie = `${name}=; Max-Age=0; Path=${path}`; } -function formatUserRole(userRole: string) { - if (!userRole) { - return "Undefined Role"; - } - switch (userRole.toLowerCase()) { - case "app_owner": - return "App Owner"; - case "demo_app_owner": - return "App Owner"; - case "app_admin": - return "Admin"; - case "proxy_admin": - return "Admin"; - case "proxy_admin_viewer": - return "Admin Viewer"; - case "org_admin": - return "Org Admin"; - case "internal_user": - return "Internal User"; - case "internal_user_viewer": - case "internal_viewer": // TODO:remove if deprecated - return "Internal Viewer"; - case "app_user": - return "App User"; - default: - return "Unknown Role"; - } -} - interface ProxySettings { PROXY_BASE_URL: string; PROXY_LOGOUT_URL: string; @@ -143,6 +115,58 @@ function CreateKeyPageContent() { const invitation_id = searchParams.get("invitation_id"); + // Parse URL query parameters for pre-filling the create key form + // Includes validation to prevent injection and DoS attacks + const autoOpenCreate = searchParams.get("create") === "true"; + const prefillData: CreateKeyPrefillData | undefined = useMemo(() => { + if (!autoOpenCreate) return undefined; + + const ownedBy = searchParams.get("owned_by"); + const teamId = searchParams.get("team_id"); + const keyAlias = searchParams.get("key_alias"); + const modelsParam = searchParams.get("models"); + const keyType = searchParams.get("key_type"); + + // Only return prefill data if at least one field is provided + if (!ownedBy && !teamId && !keyAlias && !modelsParam && !keyType) { + return undefined; + } + + // Validate owned_by against allowed values + const validOwnedByValues = ["you", "service_account", "another_user"]; + const validatedOwnedBy = ownedBy && validOwnedByValues.includes(ownedBy) + ? (ownedBy as CreateKeyPrefillData["owned_by"]) + : undefined; + + // Validate key_type against allowed values + const validKeyTypes = ["default", "llm_api", "management"]; + const validatedKeyType = keyType && validKeyTypes.includes(keyType) + ? (keyType as CreateKeyPrefillData["key_type"]) + : undefined; + + // Sanitize key_alias (limit length, trim whitespace) + const sanitizedKeyAlias = keyAlias + ? keyAlias.trim().slice(0, 256) // Reasonable max length + : undefined; + + // Sanitize models (limit array size and individual model name length) + const sanitizedModels = modelsParam + ? modelsParam + .split(",") + .slice(0, 100) // Limit number of models to prevent DoS + .map(m => m.trim().slice(0, 256)) // Limit individual model name length + .filter(m => m.length > 0) // Remove empty strings + : undefined; + + return { + owned_by: validatedOwnedBy, + team_id: teamId?.trim() || undefined, + key_alias: sanitizedKeyAlias, + models: sanitizedModels && sanitizedModels.length > 0 ? sanitizedModels : undefined, + key_type: validatedKeyType, + }; + }, [searchParams, autoOpenCreate]); + // Get page from URL, default to 'api-keys' if not present const [page, setPage] = useState(() => { return searchParams.get("page") || "api-keys"; @@ -163,6 +187,9 @@ function CreateKeyPageContent() { const [accessToken, setAccessToken] = useState(null); const [sidebarCollapsed, setSidebarCollapsed] = useState(false); + // Track if we've already attempted a return URL redirect to prevent race conditions + const hasAttemptedReturnRedirectRef = useRef(false); + const toggleSidebar = () => { setSidebarCollapsed(!sidebarCollapsed); }; @@ -207,12 +234,48 @@ function CreateKeyPageContent() { useEffect(() => { if (redirectToLogin) { + // Store the current URL so we can redirect back after login + storeReturnUrl(); + // Build login URL with return URL parameter + const baseLoginUrl = (proxyBaseUrl || "") + "/ui/login"; + const dest = buildLoginUrlWithReturn(baseLoginUrl); // Replace instead of assigning to avoid back-button loops - const dest = (proxyBaseUrl || "") + "/ui/login"; window.location.replace(dest); } }, [redirectToLogin]); + // Check for a stored return URL after successful authentication + // This handles the case where user comes back from SSO and we need to redirect to the original URL + useEffect(() => { + // Skip if still loading, no token, or we've already attempted a redirect + if (authLoading || !token || hasAttemptedReturnRedirectRef.current) { + return; + } + + // Mark that we've attempted the redirect to prevent race conditions + // This prevents duplicate redirects if token changes (e.g., refresh) + hasAttemptedReturnRedirectRef.current = true; + + // Check for a stored return URL + const returnUrl = consumeReturnUrl(); + if (returnUrl) { + const currentUrl = window.location.href; + const normalizedReturnUrl = normalizeUrlForCompare(returnUrl); + const normalizedCurrentUrl = normalizeUrlForCompare(currentUrl); + // Only redirect if the return URL is different from the current URL + // This prevents infinite redirect loops + if (normalizedReturnUrl !== normalizedCurrentUrl) { + window.location.replace(returnUrl); + } + } + }, [authLoading, token]); + + useEffect(() => { + if (!token) { + hasAttemptedReturnRedirectRef.current = false; + } + }, [token]); + useEffect(() => { if (!token) { return; @@ -410,9 +473,8 @@ function CreateKeyPageContent() { />
- -
- + +
{page == "api-keys" ? ( ) : page == "models" ? ( { - const fn = vi.fn().mockResolvedValue({ +const { formMock, setFieldsValueMock, radioGroupValueRef, formStateRef, mockKeyCreateCall } = vi.hoisted(() => { + const formStateRef = { current: {} as Record }; + const mockKeyCreateCall = vi.fn().mockResolvedValue({ key: "test-api-key", soft_budget: null, }); - return { mockKeyCreateCall: fn }; + const formMock = { + setFieldsValue: vi.fn((values: Record) => { + Object.assign(formStateRef.current, values); + }), + setFieldValue: vi.fn((name: string, value: any) => { + formStateRef.current[name] = value; + }), + getFieldValue: vi.fn((name: string) => formStateRef.current[name]), + resetFields: vi.fn(() => { + formStateRef.current = {}; + }), + }; + const radioGroupValueRef = { current: null as string | null }; + return { + formMock, + setFieldsValueMock: formMock.setFieldsValue, + radioGroupValueRef, + formStateRef, + mockKeyCreateCall, + }; +}); + +const defaultAuthorizedState = { + accessToken: "test-token", + userId: "test-user-id", + userRole: "Admin", + premiumUser: false, +}; + +let authorizedState = { ...defaultAuthorizedState }; + +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: () => authorizedState, +})); + +vi.mock("@/app/(dashboard)/hooks/keys/useKeys", () => ({ + keyKeys: { + lists: () => ["keys"], + }, +})); + +vi.mock("@ant-design/icons", () => ({ + InfoCircleOutlined: () => null, +})); + +vi.mock("react-copy-to-clipboard", () => ({ + CopyToClipboard: ({ children }: { children: any }) => children, +})); + +vi.mock("@tremor/react", () => { + const React = require("react"); + const Stub = ({ children }: { children?: any }) => React.createElement("div", null, children); + const Button = ({ children, ...props }: { children?: any }) => + React.createElement("button", props, children); + const TextInput = (props: any) => React.createElement("input", props); + + return { + Accordion: Stub, + AccordionBody: Stub, + AccordionHeader: Stub, + Button, + Col: Stub, + Grid: Stub, + Text: Stub, + TextInput, + Title: Stub, + }; +}); + +vi.mock("antd", () => { + const React = require("react"); + + const getValueFromEvent = (event: any) => { + if (event?.target) { + if (event.target.type === "checkbox") { + return event.target.checked; + } + return event.target.value; + } + return event; + }; + + const Form = ({ children, onFinish, ...props }: { children?: any; onFinish?: (values: Record) => void }) => + React.createElement( + "form", + { + ...props, + onSubmit: (event: Event) => { + event.preventDefault(); + onFinish?.({ ...formStateRef.current }); + }, + }, + children, + ); + + Form.Item = ({ children, name }: { children?: any; name?: string }) => { + if (!name || !React.isValidElement(children)) { + return React.createElement(React.Fragment, null, children); + } + + return React.cloneElement(children, { + value: formStateRef.current[name], + onChange: (event: any) => { + formStateRef.current[name] = getValueFromEvent(event); + }, + }); + }; + + Form.useForm = () => [formMock]; + + const Select = ({ children, onChange, ...props }: { children?: any; onChange?: (value: string) => void }) => + React.createElement( + "select", + { + ...props, + onChange: (event: any) => onChange?.(event.target.value), + }, + children, + ); + + Select.Option = ({ children, ...props }: { children?: any }) => + React.createElement("option", props, children); + + const Input = (props: any) => React.createElement("input", props); + Input.Password = (props: any) => React.createElement("input", { ...props, type: "password" }); + Input.TextArea = (props: any) => React.createElement("textarea", props); + + const Modal = ({ children, open }: { children?: any; open?: boolean }) => + open ? React.createElement("div", null, children) : null; + + const Radio = ({ children, ...props }: { children?: any }) => + React.createElement("div", props, children); + + Radio.Group = ({ children, value }: { children?: any; value?: string }) => { + radioGroupValueRef.current = value ?? null; + return React.createElement("div", null, children); + }; + + const Switch = (props: any) => React.createElement("input", { ...props, type: "checkbox" }); + const Tag = ({ children }: { children?: any }) => React.createElement("span", null, children); + const Tooltip = ({ children }: { children?: any }) => React.createElement(React.Fragment, null, children); + + const Button = ({ children, htmlType, ...props }: { children?: any; htmlType?: string }) => + React.createElement("button", { ...props, type: htmlType ?? props.type }, children); + + return { + Button, + Form, + Input, + message: { + success: vi.fn(), + error: vi.fn(), + warning: vi.fn(), + info: vi.fn(), + }, + Modal, + Radio, + Select, + Switch, + Tag, + Tooltip, + }; }); vi.mock("../networking", () => ({ keyCreateCall: mockKeyCreateCall, - modelAvailableCall: vi.fn().mockResolvedValue({ data: [{ id: "gpt-4" }, { id: "gpt-3.5-turbo" }] }), + modelAvailableCall: vi.fn().mockResolvedValue({ data: [{ id: "gpt-4" }] }), getGuardrailsList: vi.fn().mockResolvedValue({ guardrails: [] }), + getPoliciesList: vi.fn().mockResolvedValue({ policies: [] }), getPromptsList: vi.fn().mockResolvedValue({ prompts: [] }), proxyBaseUrl: "http://localhost:4000", getPossibleUserRoles: vi.fn().mockResolvedValue({ @@ -41,12 +204,31 @@ vi.mock("../molecules/notifications_manager", () => ({ }, })); +vi.mock("../agent_management/AgentSelector", () => ({ default: () => null })); +vi.mock("../common_components/budget_duration_dropdown", () => ({ default: () => null })); +vi.mock("../common_components/check_openapi_schema", () => ({ default: () => null })); +vi.mock("../common_components/KeyLifecycleSettings", () => ({ default: () => null })); +vi.mock("../common_components/ModelAliasManager", () => ({ default: () => null })); +vi.mock("../common_components/PassThroughRoutesSelector", () => ({ default: () => null })); +vi.mock("../common_components/PremiumLoggingSettings", () => ({ default: () => null })); +vi.mock("../common_components/RateLimitTypeFormItem", () => ({ default: () => null })); +vi.mock("../common_components/RouterSettingsAccordion", () => ({ default: () => null })); +vi.mock("../common_components/team_dropdown", () => ({ default: () => null })); +vi.mock("../CreateUserButton", () => ({ CreateUserButton: () => null })); +vi.mock("../mcp_server_management/MCPServerSelector", () => ({ default: () => null })); +vi.mock("../mcp_server_management/MCPToolPermissions", () => ({ default: () => null })); +vi.mock("../shared/numerical_input", () => ({ default: () => null })); +vi.mock("../vector_store_management/VectorStoreSelector", () => ({ default: () => null })); +vi.mock("../key_team_helpers/fetch_available_models_team_key", () => ({ + getModelDisplayName: (model: string) => model, +})); + vi.mock("../common_components/AccessGroupSelector", () => ({ default: ({ value = [], onChange }: { value?: string[]; onChange?: (v: string[]) => void }) => ( onChange?.(e.target.value ? e.target.value.split(",").map((s) => s.trim()) : [])} + onChange={(event) => onChange?.(event.target.value ? event.target.value.split(",").map((v) => v.trim()) : [])} /> ), })); @@ -54,14 +236,19 @@ vi.mock("../common_components/AccessGroupSelector", () => ({ describe("CreateKey", () => { const defaultProps = { team: null, - data: [], teams: [], + data: [], addKey: vi.fn(), }; beforeEach(() => { vi.clearAllMocks(); - localStorage.clear(); + if (typeof window !== "undefined" && window.localStorage && typeof window.localStorage.clear === "function") { + window.localStorage.clear(); + } + authorizedState = { ...defaultAuthorizedState }; + radioGroupValueRef.current = null; + formStateRef.current = {}; mockKeyCreateCall.mockResolvedValue({ key: "test-api-key", soft_budget: null, @@ -81,26 +268,8 @@ describe("CreateKey", () => { }); await waitFor(() => { - expect(screen.getByText("Key Type")).toBeInTheDocument(); - }); - - // Open the Key Type dropdown - const keyTypeSection = screen.getByText("Key Type").closest(".ant-form-item")!; - const selectElement = keyTypeSection.querySelector(".ant-select-selector")!; - act(() => { - fireEvent.mouseDown(selectElement); - }); - - await waitFor(() => { - // Verify "AI APIs" appears as an option - const options = document.querySelectorAll(".ant-select-item-option"); - const optionTexts = Array.from(options).map((el) => el.textContent); - const hasAIAPIs = optionTexts.some((text) => text?.includes("AI APIs")); - expect(hasAIAPIs).toBe(true); - - // Verify old "LLM API" label does NOT appear - const hasLLMAPI = optionTexts.some((text) => text?.includes("LLM API")); - expect(hasLLMAPI).toBe(false); + expect(screen.getByText("AI APIs")).toBeInTheDocument(); + expect(screen.queryByText("LLM API")).not.toBeInTheDocument(); }); }); @@ -111,46 +280,118 @@ describe("CreateKey", () => { fireEvent.click(screen.getByRole("button", { name: /create new key/i })); }); - await waitFor(() => { - expect(screen.getByLabelText(/key name/i)).toBeInTheDocument(); - }); - - fireEvent.change(screen.getByLabelText(/key name/i), { target: { value: "Test Key" } }); - - const optionalSettingsAccordion = screen.getByText("Optional Settings"); - act(() => { - fireEvent.click(optionalSettingsAccordion); - }); - await waitFor(() => { expect(screen.getByTestId("access-group-selector")).toBeInTheDocument(); }); - fireEvent.change(screen.getByTestId("access-group-selector"), { target: { value: "ag-1,ag-2" } }); + act(() => { + fireEvent.change(screen.getByTestId("access-group-selector"), { target: { value: "ag-1,ag-2" } }); + formMock.setFieldValue("key_alias", "Test Key"); + }); - const modelsCombobox = screen.getAllByRole("combobox").find((el) => el.closest('[class*="ant-form-item"]')?.textContent?.includes("Models")) || - screen.getAllByRole("combobox")[1]; - if (modelsCombobox) { - act(() => fireEvent.mouseDown(modelsCombobox)); - await waitFor(() => { - const allTeamModels = [...document.body.querySelectorAll(".ant-select-item")].find( - (el) => el.textContent?.includes("All Team Models"), - ); - if (allTeamModels) fireEvent.click(allTeamModels); - }); - } + act(() => { + fireEvent.click(screen.getByRole("button", { name: /create key/i })); + }); - const createButton = screen.getByRole("button", { name: /create key/i }); - act(() => fireEvent.click(createButton)); + await waitFor(() => { + expect(mockKeyCreateCall).toHaveBeenCalled(); + const formValues = mockKeyCreateCall.mock.calls[0][2]; + expect(formValues).toHaveProperty("access_group_ids"); + expect(formValues.access_group_ids).toEqual(["ag-1", "ag-2"]); + }); + }); - await waitFor( - () => { - expect(mockKeyCreateCall).toHaveBeenCalled(); - const formValues = mockKeyCreateCall.mock.calls[0][2]; - expect(formValues).toHaveProperty("access_group_ids"); - expect(formValues.access_group_ids).toEqual(["ag-1", "ag-2"]); - }, - { timeout: 15000 }, + it("should prefill models when provided without team_id", async () => { + renderWithProviders( + , ); - }, { timeout: 30000 }); + + await waitFor(() => { + expect(setFieldsValueMock).toHaveBeenCalledWith({ models: ["gpt-4"] }); + }); + }); + + it("should prefill team_id when it exists in teams", async () => { + renderWithProviders( + , + ); + + await waitFor(() => { + expect(setFieldsValueMock).toHaveBeenCalledWith({ team_id: "team-1" }); + }); + }); + + it("should ignore team_id when it does not exist in teams", async () => { + renderWithProviders( + , + ); + + await waitFor(() => { + expect(setFieldsValueMock).toHaveBeenCalledWith({ key_alias: "example-key" }); + }); + + expect(setFieldsValueMock).not.toHaveBeenCalledWith({ team_id: "team-404" }); + }); + + it('should fall back to "you" when owned_by is another_user for non-admin', async () => { + authorizedState = { ...defaultAuthorizedState, userRole: "Internal User" }; + + renderWithProviders( + , + ); + + await waitFor(() => { + expect(setFieldsValueMock).toHaveBeenCalledWith({ key_alias: "example-key" }); + }); + + expect(radioGroupValueRef.current).toBe("you"); + }); + + it("should apply owned_by another_user for admin", async () => { + renderWithProviders( + , + ); + + await waitFor(() => { + expect(radioGroupValueRef.current).toBe("another_user"); + }); + }); + + it("should prefill key_type when provided", async () => { + renderWithProviders( + , + ); + + await waitFor(() => { + expect(setFieldsValueMock).toHaveBeenCalledWith({ key_type: "management" }); + }); + }); }); diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx index 961a5d4d460..d071c4a4a3b 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx @@ -46,11 +46,24 @@ import { simplifyKeyGenerateError } from "./utils"; const { Option } = Select; +/** + * Interface for pre-filling the create key form from URL parameters + */ +export interface CreateKeyPrefillData { + owned_by?: "you" | "service_account" | "another_user"; + team_id?: string; + key_alias?: string; + models?: string[]; + key_type?: "default" | "llm_api" | "management"; +} + interface CreateKeyProps { team: Team | null; data: any[] | null; teams: Team[] | null; addKey: (data: any) => void; + autoOpenCreate?: boolean; + prefillData?: CreateKeyPrefillData; } interface User { @@ -141,7 +154,7 @@ export const fetchUserModels = async ( * Please contribute to the new refactor. * ───────────────────────────────────────────────────────────────────────── */ -const CreateKey: React.FC = ({ team, teams, data, addKey }) => { +const CreateKey: React.FC = ({ team, teams, data, addKey, autoOpenCreate, prefillData }) => { const { accessToken, userId: userID, userRole, premiumUser } = useAuthorized(); const queryClient = useQueryClient(); const [form] = Form.useForm(); @@ -152,6 +165,8 @@ const CreateKey: React.FC = ({ team, teams, data, addKey }) => { const [modelsToPick, setModelsToPick] = useState([]); const [keyOwner, setKeyOwner] = useState("you"); const [predefinedTags, setPredefinedTags] = useState(getPredefinedTags(data)); + const [hasPrefilled, setHasPrefilled] = useState(false); + const [pendingPrefillModels, setPendingPrefillModels] = useState(null); const [guardrailsList, setGuardrailsList] = useState([]); const [policiesList, setPoliciesList] = useState([]); const [promptsList, setPromptsList] = useState([]); @@ -274,6 +289,55 @@ const CreateKey: React.FC = ({ team, teams, data, addKey }) => { fetchPossibleRoles(); }, [accessToken]); + // Auto-open modal and prefill form from URL params (deep link). + // Guarded by write access so we don't open for read-only users. + useEffect(() => { + if (autoOpenCreate && !hasPrefilled && teams && userRole && rolesWithWriteAccess.includes(userRole)) { + // Open the modal + setIsModalVisible(true); + setHasPrefilled(true); + + // Apply prefill data if provided + if (prefillData) { + // Set key owner (owned_by) - validate that "another_user" is only allowed for Admin + if (prefillData.owned_by) { + if (prefillData.owned_by === "another_user" && userRole !== "Admin") { + // Ignore invalid owned_by for non-admin users, fall back to default + setKeyOwner("you"); + } else { + setKeyOwner(prefillData.owned_by); + } + } + + // Set team - find the team by ID and set it (only if team exists in user's teams) + if (prefillData.team_id) { + const selectedTeam = teams?.find((t) => t.team_id === prefillData.team_id) || null; + if (selectedTeam) { + setSelectedCreateKeyTeam(selectedTeam); + form.setFieldsValue({ team_id: prefillData.team_id }); + } + // Silently ignore invalid team_id - don't prefill with a team user doesn't have access to + } + + // Set key alias + if (prefillData.key_alias) { + form.setFieldsValue({ key_alias: prefillData.key_alias }); + } + + // Defer model selection until we load the allowed model list. + if (prefillData.models && prefillData.models.length > 0) { + setPendingPrefillModels(prefillData.models); + } + + // Set key type + if (prefillData.key_type) { + setKeyType(prefillData.key_type); + form.setFieldsValue({ key_type: prefillData.key_type }); + } + } + } + }, [autoOpenCreate, prefillData, teams, hasPrefilled, form, userRole]); + // Check if team selection is required const isTeamSelectionRequired = modelsToPick.includes("no-default-models"); const isFormDisabled = isTeamSelectionRequired && !selectedCreateKeyTeam; @@ -467,6 +531,9 @@ const CreateKey: React.FC = ({ team, teams, data, addKey }) => { NotificationsManager.success("Virtual Key copied to clipboard"); }; + // Fetch available models when team or auth changes. + // Note: Model prefill from URL params is handled by the useEffect below, which + // watches for pendingPrefillModels + modelsToPick to both be populated. useEffect(() => { if (userID && userRole && accessToken) { fetchTeamModels(userID, userRole, accessToken, selectedCreateKeyTeam?.team_id ?? null).then((models) => { @@ -474,8 +541,28 @@ const CreateKey: React.FC = ({ team, teams, data, addKey }) => { setModelsToPick(allModels); }); } - form.setFieldValue("models", []); - }, [selectedCreateKeyTeam, accessToken, userID, userRole]); + // Only clear models if we don't have pending prefill models + if (!pendingPrefillModels) { + form.setFieldValue("models", []); + } + }, [selectedCreateKeyTeam, accessToken, userID, userRole, form]); + + // Apply deferred model prefill once the available model list arrives. + // This handles timing where prefill data arrives before or after models are fetched. + useEffect(() => { + if (!pendingPrefillModels || pendingPrefillModels.length === 0) { + return; + } + if (!modelsToPick || modelsToPick.length === 0) { + return; + } + + const validModels = pendingPrefillModels.filter((model) => modelsToPick.includes(model)); + if (validModels.length > 0) { + form.setFieldsValue({ models: validModels }); + } + setPendingPrefillModels(null); + }, [pendingPrefillModels, modelsToPick, form]); // Add a callback function to handle user creation const handleUserCreated = (userId: string) => { diff --git a/ui/litellm-dashboard/src/components/user_dashboard.tsx b/ui/litellm-dashboard/src/components/user_dashboard.tsx index ec6de82fdf9..ecb17027548 100644 --- a/ui/litellm-dashboard/src/components/user_dashboard.tsx +++ b/ui/litellm-dashboard/src/components/user_dashboard.tsx @@ -16,7 +16,7 @@ import { Organization, userInfoCall, } from "./networking"; -import CreateKey from "./organisms/create_key_button"; +import CreateKey, { CreateKeyPrefillData } from "./organisms/create_key_button"; import { VirtualKeysTable } from "./VirtualKeysPage/VirtualKeysTable"; export interface ProxySettings { @@ -55,6 +55,8 @@ interface UserDashboardProps { organizations: Organization[] | null; addKey: (data: any) => void; createClicked: boolean; + autoOpenCreate?: boolean; + prefillData?: CreateKeyPrefillData; } type TeamInterface = { @@ -77,6 +79,8 @@ const UserDashboard: React.FC = ({ organizations, addKey, createClicked, + autoOpenCreate, + prefillData, }) => { const [userSpendData, setUserSpendData] = useState(null); const [currentOrg, setCurrentOrg] = useState(null); @@ -350,6 +354,8 @@ const UserDashboard: React.FC = ({ teams={teams as Team[]} data={keys} addKey={addKey} + autoOpenCreate={autoOpenCreate} + prefillData={prefillData} /> diff --git a/ui/litellm-dashboard/src/utils/returnUrlUtils.test.ts b/ui/litellm-dashboard/src/utils/returnUrlUtils.test.ts new file mode 100644 index 00000000000..3c09e550145 --- /dev/null +++ b/ui/litellm-dashboard/src/utils/returnUrlUtils.test.ts @@ -0,0 +1,383 @@ +import { + buildLoginUrlWithReturn, + clearStoredReturnUrl, + consumeReturnUrl, + getCurrentUrl, + getReturnUrl, + getReturnUrlFromParams, + getStoredReturnUrl, + isValidReturnUrl, + storeReturnUrl, +} from "./returnUrlUtils"; + +describe("returnUrlUtils", () => { + const originalLocation = window.location; + + beforeEach(() => { + // Clear cookies before each test + document.cookie.split(";").forEach((c) => { + document.cookie = c + .replace(/^ +/, "") + .replace(/=.*/, "=;expires=" + new Date().toUTCString() + ";path=/"); + }); + + // Reset location mock + Object.defineProperty(window, "location", { + value: { + href: "http://localhost:3000/ui?page=api-keys", + origin: "http://localhost:3000", + hostname: "localhost", + pathname: "/ui", + search: "?page=api-keys", + }, + writable: true, + }); + }); + + afterEach(() => { + // Restore original location + Object.defineProperty(window, "location", { + value: originalLocation, + writable: true, + }); + }); + + describe("getCurrentUrl", () => { + it("should return the current URL", () => { + const url = getCurrentUrl(); + expect(url).toBe("http://localhost:3000/ui?page=api-keys"); + }); + }); + + describe("storeReturnUrl and getStoredReturnUrl", () => { + it("should store and retrieve the return URL from cookie", () => { + storeReturnUrl(); + const storedUrl = getStoredReturnUrl(); + expect(storedUrl).toBe("http://localhost:3000/ui?page=api-keys"); + }); + + it("should return null if no URL is stored", () => { + const storedUrl = getStoredReturnUrl(); + expect(storedUrl).toBeNull(); + }); + }); + + describe("clearStoredReturnUrl", () => { + it("should clear the stored return URL", () => { + storeReturnUrl(); + expect(getStoredReturnUrl()).not.toBeNull(); + + clearStoredReturnUrl(); + expect(getStoredReturnUrl()).toBeNull(); + }); + }); + + describe("getReturnUrlFromParams", () => { + it("should return the redirect_to parameter from URL", () => { + Object.defineProperty(window, "location", { + value: { + ...window.location, + search: "?redirect_to=http%3A%2F%2Flocalhost%3A3000%2Fui%3Fcreate%3Dtrue", + }, + writable: true, + }); + + const returnUrl = getReturnUrlFromParams(); + expect(returnUrl).toBe("http://localhost:3000/ui?create=true"); + }); + + it("should return null if redirect_to parameter is not present", () => { + Object.defineProperty(window, "location", { + value: { + ...window.location, + search: "?page=api-keys", + }, + writable: true, + }); + + const returnUrl = getReturnUrlFromParams(); + expect(returnUrl).toBeNull(); + }); + }); + + describe("buildLoginUrlWithReturn", () => { + it("should build login URL with return URL parameter", () => { + Object.defineProperty(window, "location", { + value: { + ...window.location, + href: "http://localhost:3000/ui?create=true&team_id=123", + }, + writable: true, + }); + + const loginUrl = buildLoginUrlWithReturn("/ui/login"); + expect(loginUrl).toBe( + "/ui/login?redirect_to=http%3A%2F%2Flocalhost%3A3000%2Fui%3Fcreate%3Dtrue%26team_id%3D123" + ); + }); + + it("should not add return URL if already on login page", () => { + Object.defineProperty(window, "location", { + value: { + ...window.location, + href: "http://localhost:3000/ui/login", + }, + writable: true, + }); + + const loginUrl = buildLoginUrlWithReturn("/ui/login"); + expect(loginUrl).toBe("/ui/login"); + }); + + it("should handle login URL with existing query parameters", () => { + Object.defineProperty(window, "location", { + value: { + ...window.location, + href: "http://localhost:3000/ui?page=api-keys", + }, + writable: true, + }); + + const loginUrl = buildLoginUrlWithReturn("/ui/login?foo=bar"); + expect(loginUrl).toContain("&redirect_to="); + }); + }); + + describe("getReturnUrl", () => { + it("should prefer URL params over cookie", () => { + // Store a URL in cookie + storeReturnUrl(); + + // Set a different URL in the params + Object.defineProperty(window, "location", { + value: { + ...window.location, + search: "?redirect_to=http%3A%2F%2Flocalhost%3A3000%2Fui%3Fpage%3Dteams", + }, + writable: true, + }); + + const returnUrl = getReturnUrl(); + expect(returnUrl).toBe("http://localhost:3000/ui?page=teams"); + }); + + it("should fall back to cookie if no URL param", () => { + // Store a URL in cookie first + Object.defineProperty(window, "location", { + value: { + href: "http://localhost:3000/ui?create=true", + origin: "http://localhost:3000", + hostname: "localhost", + pathname: "/ui", + search: "?create=true", + }, + writable: true, + }); + storeReturnUrl(); + + // Clear the URL params + Object.defineProperty(window, "location", { + value: { + ...window.location, + search: "", + }, + writable: true, + }); + + const returnUrl = getReturnUrl(); + expect(returnUrl).toBe("http://localhost:3000/ui?create=true"); + }); + + it("should return null if no return URL found", () => { + Object.defineProperty(window, "location", { + value: { + ...window.location, + search: "", + }, + writable: true, + }); + + const returnUrl = getReturnUrl(); + expect(returnUrl).toBeNull(); + }); + }); + + describe("isValidReturnUrl", () => { + it("should validate relative URLs starting with /", () => { + expect(isValidReturnUrl("/ui?page=api-keys")).toBe(true); + expect(isValidReturnUrl("/ui/teams")).toBe(true); + }); + + it("should reject protocol-relative URLs", () => { + expect(isValidReturnUrl("//evil.com")).toBe(false); + }); + + it("should validate same-hostname URLs (even with different ports) in dev", () => { + // Same hostname, same port + expect(isValidReturnUrl("http://localhost:3000/ui?page=teams")).toBe(true); + // Same hostname, different port (important for dev environments) + expect(isValidReturnUrl("http://localhost:4000/ui?page=teams")).toBe(true); + }); + + it("should reject different-hostname URLs", () => { + expect(isValidReturnUrl("http://evil.com/ui")).toBe(false); + expect(isValidReturnUrl("https://google.com")).toBe(false); + }); + + it("should reject empty URLs", () => { + expect(isValidReturnUrl("")).toBe(false); + }); + + it("should reject invalid URLs", () => { + expect(isValidReturnUrl("not-a-url")).toBe(false); + }); + + it("should reject XSS attempts with javascript: protocol", () => { + expect(isValidReturnUrl('javascript:alert("xss")')).toBe(false); + expect(isValidReturnUrl("javascript:void(0)")).toBe(false); + }); + + it("should reject data: URLs", () => { + expect(isValidReturnUrl("data:text/html,")).toBe(false); + }); + + it("should allow 127.x.x.x addresses in dev environment", () => { + Object.defineProperty(window, "location", { + value: { + href: "http://127.0.0.1:3000/ui", + origin: "http://127.0.0.1:3000", + hostname: "127.0.0.1", + protocol: "http:", + pathname: "/ui", + search: "", + }, + writable: true, + }); + + expect(isValidReturnUrl("http://127.0.0.1:4000/ui")).toBe(true); + }); + + it("should allow .local domains in dev environment", () => { + Object.defineProperty(window, "location", { + value: { + href: "http://myapp.local:3000/ui", + origin: "http://myapp.local:3000", + hostname: "myapp.local", + protocol: "http:", + pathname: "/ui", + search: "", + }, + writable: true, + }); + + // Same hostname with different port should be allowed in dev + expect(isValidReturnUrl("http://myapp.local:4000/ui")).toBe(true); + }); + + it("should reject cross-port redirects in production environment", () => { + // Simulate production environment + Object.defineProperty(window, "location", { + value: { + href: "https://app.example.com/ui", + origin: "https://app.example.com", + hostname: "app.example.com", + protocol: "https:", + pathname: "/ui", + search: "", + }, + writable: true, + }); + + // Same origin should work + expect(isValidReturnUrl("https://app.example.com/ui?page=teams")).toBe(true); + // Different port should be rejected in production + expect(isValidReturnUrl("https://app.example.com:8080/ui")).toBe(false); + // Different hostname should be rejected + expect(isValidReturnUrl("https://evil.com/ui")).toBe(false); + }); + }); + + describe("consumeReturnUrl", () => { + it("should return and clear the stored return URL", () => { + Object.defineProperty(window, "location", { + value: { + href: "http://localhost:3000/ui?create=true", + origin: "http://localhost:3000", + hostname: "localhost", + pathname: "/ui", + search: "?create=true", + }, + writable: true, + }); + storeReturnUrl(); + + // Clear the URL params for the consume call + Object.defineProperty(window, "location", { + value: { + href: "http://localhost:3000/ui/login", + origin: "http://localhost:3000", + hostname: "localhost", + pathname: "/ui/login", + search: "", + }, + writable: true, + }); + + const returnUrl = consumeReturnUrl(); + expect(returnUrl).toBe("http://localhost:3000/ui?create=true"); + expect(getStoredReturnUrl()).toBeNull(); + }); + + it("should return null for invalid return URLs (different hostname)", () => { + // Manually set an invalid URL in cookie + document.cookie = "litellm_return_url=" + encodeURIComponent("http://evil.com/phishing") + "; path=/"; + + const returnUrl = consumeReturnUrl(); + expect(returnUrl).toBeNull(); + }); + + it("should allow URLs with different ports on same hostname", () => { + // Store URL with port 3000 + Object.defineProperty(window, "location", { + value: { + href: "http://localhost:3000/ui?create=true", + origin: "http://localhost:3000", + hostname: "localhost", + pathname: "/ui", + search: "?create=true", + }, + writable: true, + }); + storeReturnUrl(); + + // Now we're on port 4000 + Object.defineProperty(window, "location", { + value: { + href: "http://localhost:4000/ui", + origin: "http://localhost:4000", + hostname: "localhost", + pathname: "/ui", + search: "", + }, + writable: true, + }); + + const returnUrl = consumeReturnUrl(); + // Should be valid because same hostname (localhost) + expect(returnUrl).toBe("http://localhost:3000/ui?create=true"); + }); + + it("should return null if no return URL found", () => { + Object.defineProperty(window, "location", { + value: { + ...window.location, + search: "", + }, + writable: true, + }); + + const returnUrl = consumeReturnUrl(); + expect(returnUrl).toBeNull(); + }); + }); +}); diff --git a/ui/litellm-dashboard/src/utils/returnUrlUtils.ts b/ui/litellm-dashboard/src/utils/returnUrlUtils.ts new file mode 100644 index 00000000000..76562a0122a --- /dev/null +++ b/ui/litellm-dashboard/src/utils/returnUrlUtils.ts @@ -0,0 +1,304 @@ +/** + * Utility functions for managing return URLs during authentication flows. + * + * When a user is redirected to login, we store the original URL so they can be + * redirected back after successful authentication. + * + * NOTE: We use cookies instead of sessionStorage because the SSO flow may cross + * different ports (e.g., localhost:3000 -> localhost:4000), and sessionStorage + * is not shared across different origins. Cookies on the same hostname are shared + * across different ports. + */ + +const RETURN_URL_COOKIE_NAME = "litellm_return_url"; +const RETURN_URL_PARAM = "redirect_to"; + +/** + * Gets the current URL with all query parameters. + * Returns null if running on server-side. + */ +export function getCurrentUrl(): string | null { + if (typeof window === "undefined") { + return null; + } + return window.location.href; +} + +/** + * Sets a cookie with the given name and value. + * Automatically adds Secure flag when running over HTTPS. + */ +function setCookie(name: string, value: string, maxAgeSeconds: number = 300): void { + if (typeof document === "undefined") { + return; + } + // Set cookie with path=/ so it's available across all paths + // Use SameSite=Lax to allow the cookie to be sent on navigation from external sites (SSO redirect) + // Add Secure flag when running over HTTPS to prevent cookie from being sent over unencrypted connections + const isSecure = typeof window !== "undefined" && window.location.protocol === "https:"; + const secureFlag = isSecure ? "; Secure" : ""; + document.cookie = `${name}=${encodeURIComponent(value)}; path=/; max-age=${maxAgeSeconds}; SameSite=Lax${secureFlag}`; +} + +/** + * Gets a cookie value by name. + */ +function getCookie(name: string): string | null { + if (typeof document === "undefined") { + return null; + } + const match = document.cookie.match(new RegExp(`(^| )${name}=([^;]+)`)); + if (match) { + try { + return decodeURIComponent(match[2]); + } catch { + return match[2]; + } + } + return null; +} + +/** + * Deletes a cookie by name. + */ +function deleteCookie(name: string): void { + if (typeof document === "undefined") { + return; + } + document.cookie = `${name}=; path=/; max-age=0`; +} + +/** + * Stores the current URL in a cookie before redirecting to login. + * This allows us to redirect the user back to their original destination after login. + * Cookie expires in 5 minutes (300 seconds). + */ +export function storeReturnUrl(): void { + if (typeof window === "undefined") { + return; + } + + const currentUrl = getCurrentUrl(); + if (currentUrl) { + setCookie(RETURN_URL_COOKIE_NAME, currentUrl, 300); + } +} + +/** + * Retrieves the stored return URL from the cookie. + * Returns null if no return URL is stored or if running on server-side. + */ +export function getStoredReturnUrl(): string | null { + if (typeof window === "undefined") { + return null; + } + return getCookie(RETURN_URL_COOKIE_NAME); +} + +/** + * Clears the stored return URL from the cookie. + * Should be called after redirecting to the return URL. + */ +export function clearStoredReturnUrl(): void { + if (typeof window === "undefined") { + return; + } + + try { + deleteCookie(RETURN_URL_COOKIE_NAME); + } catch (error) { + console.error("Failed to clear return URL cookie:", error); + } +} + +/** + * Gets the return URL from URL query parameters. + * Used when the return URL is passed via query string to the login page. + */ +export function getReturnUrlFromParams(): string | null { + if (typeof window === "undefined") { + return null; + } + + const searchParams = new URLSearchParams(window.location.search); + return searchParams.get(RETURN_URL_PARAM); +} + +/** + * Builds a login URL with the return URL as a query parameter. + * + * @param baseLoginUrl - The base login URL (e.g., "/ui/login") + * @param returnUrl - The URL to redirect to after login (defaults to current URL) + */ +export function buildLoginUrlWithReturn(baseLoginUrl: string, returnUrl?: string): string { + const url = returnUrl || getCurrentUrl(); + + if (!url) { + return baseLoginUrl; + } + + // Don't add return URL if we're already on the login page + if (url.includes("/login")) { + return baseLoginUrl; + } + + const separator = baseLoginUrl.includes("?") ? "&" : "?"; + return `${baseLoginUrl}${separator}${RETURN_URL_PARAM}=${encodeURIComponent(url)}`; +} + +/** + * Gets the best return URL to use after login. + * Priority: + * 1. URL query parameter (redirect_to) + * 2. Cookie + * 3. null (caller should use default) + */ +export function getReturnUrl(): string | null { + // First check URL params + const paramUrl = getReturnUrlFromParams(); + if (paramUrl) { + return paramUrl; + } + + // Then check cookie + const storedUrl = getStoredReturnUrl(); + if (storedUrl) { + return storedUrl; + } + + return null; +} + +/** + * Checks if we're running in a development environment. + * Returns true for localhost, 127.0.0.1, IPv6 localhost, or .local domains. + * This determines whether cross-port redirects are allowed (dev only). + */ +function isDevEnvironment(): boolean { + if (typeof window === "undefined") { + return false; + } + const hostname = window.location.hostname; + return ( + hostname === "localhost" || + hostname === "127.0.0.1" || + hostname === "::1" || + hostname.startsWith("127.") || // Full IPv4 loopback range (127.0.0.0/8) + hostname.endsWith(".local") // Common dev domain suffix + ); +} + +/** + * Validates a return URL to prevent open redirect attacks. + * - Always allows relative URLs (starting with / but not //) + * - In dev (localhost): allows same hostname with any port + * - In production: requires exact origin match (protocol + hostname + port) + * + * @param url - The URL to validate + * @returns true if the URL is safe to redirect to + */ +export function isValidReturnUrl(url: string): boolean { + if (!url) { + return false; + } + + // Allow relative URLs + if (url.startsWith("/") && !url.startsWith("//")) { + return true; + } + + // For absolute URLs, validate against current origin + if (typeof window === "undefined") { + return false; + } + + try { + const returnUrlObj = new URL(url); + const currentHostname = window.location.hostname; + + // Hostname must always match + if (returnUrlObj.hostname !== currentHostname) { + return false; + } + + // In dev environments (localhost), allow any port on the same hostname + // This supports SSO flows that cross ports (e.g., localhost:3000 -> localhost:4000) + if (isDevEnvironment()) { + return true; + } + + // In production, require exact origin match (protocol + hostname + port) + return returnUrlObj.origin === window.location.origin; + } catch { + // Invalid URL + return false; + } +} + +export function normalizeUrlForCompare(url: string): string { + if (typeof window === "undefined") { + return url; + } + + try { + const parsed = new URL(url, window.location.origin); + let pathname = parsed.pathname; + if (pathname.length > 1 && pathname.endsWith("/")) { + pathname = pathname.slice(0, -1); + } + + const params = new URLSearchParams(parsed.search); + const sortedParams = new URLSearchParams(); + Array.from(params.entries()) + .sort(([a], [b]) => a.localeCompare(b)) + .forEach(([key, value]) => { + sortedParams.append(key, value); + }); + + const search = sortedParams.toString(); + const hash = parsed.hash || ""; + return `${parsed.origin}${pathname}${search ? `?${search}` : ""}${hash}`; + } catch { + return url; + } +} + +/** + * Gets and clears the return URL in one operation. + * Returns the validated return URL or null if invalid/not found. + * + * Priority: + * 1. If redirect_to param is valid, use it and clear cookie + * 2. If redirect_to param is invalid/missing, check cookie + * 3. Only clear cookie when we have a valid URL to return + */ +export function consumeReturnUrl(): string | null { + // Check URL param first + const paramUrl = getReturnUrlFromParams(); + if (paramUrl) { + if (isValidReturnUrl(paramUrl)) { + clearStoredReturnUrl(); + return paramUrl; + } + // Log rejected URLs in development for debugging + if (isDevEnvironment()) { + console.warn("[returnUrlUtils] Invalid return URL in params rejected:", paramUrl); + } + } + + // Fall back to cookie + const storedUrl = getStoredReturnUrl(); + if (storedUrl) { + if (isValidReturnUrl(storedUrl)) { + clearStoredReturnUrl(); + return storedUrl; + } + // Log rejected URLs in development for debugging + if (isDevEnvironment()) { + console.warn("[returnUrlUtils] Invalid return URL in cookie rejected:", storedUrl); + } + } + + // No valid URL found - don't clear cookie (nothing to clear or already invalid) + return null; +} diff --git a/ui/litellm-dashboard/src/utils/roles.ts b/ui/litellm-dashboard/src/utils/roles.ts index 580b4568c53..608a54ae143 100644 --- a/ui/litellm-dashboard/src/utils/roles.ts +++ b/ui/litellm-dashboard/src/utils/roles.ts @@ -31,3 +31,32 @@ export const isUserTeamAdminForSingleTeam = (teamMemberWithRoles: Member[] | nul } return teamMemberWithRoles.some((member) => member.user_id === userID && member.role === "admin"); }; + +export const formatUserRole = (userRole: string): string => { + if (!userRole) { + return "Undefined Role"; + } + switch (userRole.toLowerCase()) { + case "app_owner": + return "App Owner"; + case "demo_app_owner": + return "App Owner"; + case "app_admin": + return "Admin"; + case "proxy_admin": + return "Admin"; + case "proxy_admin_viewer": + return "Admin Viewer"; + case "org_admin": + return "Org Admin"; + case "internal_user": + return "Internal User"; + case "internal_user_viewer": + case "internal_viewer": // TODO:remove if deprecated + return "Internal Viewer"; + case "app_user": + return "App User"; + default: + return "Unknown Role"; + } +}; diff --git a/ui/litellm-dashboard/tests/CreateKeyPage.expiredToken.test.tsx b/ui/litellm-dashboard/tests/CreateKeyPage.expiredToken.test.tsx index c3c5ae59237..8b05def9ba3 100644 --- a/ui/litellm-dashboard/tests/CreateKeyPage.expiredToken.test.tsx +++ b/ui/litellm-dashboard/tests/CreateKeyPage.expiredToken.test.tsx @@ -5,12 +5,13 @@ import { vi, describe, it, beforeEach, afterEach, expect } from "vitest"; /** ---------------------------- * Hoisted helpers for mocks (required by Vitest) * --------------------------- */ -const { stub, jwtDecodeMock } = vi.hoisted(() => { +const { stub, jwtDecodeMock, consumeReturnUrlMock } = vi.hoisted(() => { const React = require("react"); const stub = (name: string) => () => React.createElement("div", { "data-testid": name }); return { stub, jwtDecodeMock: vi.fn(), + consumeReturnUrlMock: vi.fn(), }; }); @@ -84,6 +85,14 @@ vi.mock("jwt-decode", () => ({ jwtDecode: (token: string) => jwtDecodeMock(token), })); +vi.mock("@/utils/returnUrlUtils", async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + consumeReturnUrl: consumeReturnUrlMock, + }; +}); + // Super-light stubs for all heavy components so rendering doesn't explode vi.mock("@/components/navbar", () => ({ default: stub("navbar") })); vi.mock("@/components/user_dashboard", () => ({ default: stub("user-dashboard") })); @@ -152,6 +161,7 @@ beforeEach(() => { // Fresh module state & DOM vi.clearAllMocks(); clearAllCookies(); + consumeReturnUrlMock.mockReturnValue(null); // Make location.replace spy-able to validate redirect delete (window as any).location; @@ -191,9 +201,11 @@ describe("CreateKeyPage auth behavior", () => { // Act render(); - // Assert: we eventually redirect to SSO login (single replace, not assign/href) + // Assert: we eventually redirect to SSO login with return URL (single replace, not assign/href) await waitFor(() => { - expect(window.location.replace).toHaveBeenCalledWith("https://example.com/ui/login"); + expect(window.location.replace).toHaveBeenCalledWith( + expect.stringContaining("https://example.com/ui/login?redirect_to=") + ); }); // And we attempted to clear the cookie (defensive deletion) @@ -235,4 +247,41 @@ describe("CreateKeyPage auth behavior", () => { expect(screen.getByTestId("navbar")).toBeInTheDocument(); }); }); + + it("should not redirect when return URL only differs by query order", async () => { + setCookie("token=validtoken"); + + jwtDecodeMock.mockImplementation((tok: string) => { + expect(tok).toBe("validtoken"); + return { + exp: Math.floor(Date.now() / 1000) + 60 * 60, + key: "accessKey-123", + user_role: "app_user", + user_email: "user@example.com", + login_method: "username_password", + premium_user: false, + auth_header_name: "x-litellm-auth", + user_id: "u_123", + }; + }); + + // Current URL has params in a different order + delete (window as any).location; + (window as any).location = { + ...originalLocation, + href: "http://localhost/ui?b=2&a=1", + origin: "http://localhost", + assign: vi.fn(), + replace: vi.fn(), + }; + + // Return URL has the same params in a different order + consumeReturnUrlMock.mockReturnValue("http://localhost/ui?a=1&b=2"); + + render(); + + await waitFor(() => { + expect(window.location.replace).not.toHaveBeenCalled(); + }); + }); }); From f1c563d2b2550d553f7adc368d5097dbbf2f92a7 Mon Sep 17 00:00:00 2001 From: Alejandro Tapia Date: Fri, 27 Feb 2026 14:58:17 -0800 Subject: [PATCH 012/273] org-exclusive-add-member --- .../internal_user_endpoints.py | 62 ++++++++- .../test_internal_user_endpoints.py | 127 +++++++++++++++++- 2 files changed, 179 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index e535ccaaa46..f5488ce865d 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -30,6 +30,7 @@ from litellm.proxy.management_endpoints.common_daily_activity import ( get_daily_activity, get_daily_activity_aggregated, ) +from litellm.proxy.auth.auth_checks import get_user_object from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view from litellm.proxy.management_endpoints.key_management_endpoints import ( generate_key_helper_fn, @@ -1830,7 +1831,11 @@ async def ui_view_users( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ - [PROXY-ADMIN ONLY]Filter users based on partial match of user_id or email with pagination. + Filter users based on partial match of user_id or email with pagination. + + - Proxy admins: receive all matching users. + - Organization admins: receive only users in their own organization(s). + - Other roles: access denied (403). Args: user_id (Optional[str]): Partial user ID to search for @@ -1840,19 +1845,60 @@ async def ui_view_users( user_api_key_dict (UserAPIKeyAuth): User authentication information Returns: - List[LiteLLM_SpendLogs]: Paginated list of matching user records + List of matching user records (LiteLLM_UserTableFiltered), scoped by org for org admins. """ - from litellm.proxy.proxy_server import prisma_client + from litellm.proxy.proxy_server import ( + prisma_client, + proxy_logging_obj, + user_api_key_cache, + ) if prisma_client is None: raise HTTPException(status_code=500, detail={"error": "No db connected"}) try: + # Restrict by caller role: proxy admin sees all; org admin sees only their org(s); others 403 + is_proxy_admin = user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN + if not is_proxy_admin: + if user_api_key_dict.user_id is None: + raise HTTPException( + status_code=403, + detail={ + "error": "Only proxy admins and organization admins can search users." + }, + ) + caller_user = await get_user_object( + user_id=user_api_key_dict.user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + user_id_upsert=False, + proxy_logging_obj=proxy_logging_obj, + ) + if caller_user is None: + raise HTTPException( + status_code=403, + detail={ + "error": "Only proxy admins and organization admins can search users." + }, + ) + org_admin_org_ids = [ + m.organization_id + for m in (caller_user.organization_memberships or []) + if m.user_role == LitellmUserRoles.ORG_ADMIN.value + ] + if not org_admin_org_ids: + raise HTTPException( + status_code=403, + detail={ + "error": "Only proxy admins and organization admins can search users." + }, + ) + # Calculate offset for pagination skip = (page - 1) * page_size # Build where conditions based on provided parameters - where_conditions = {} + where_conditions: Dict[str, Any] = {} if user_id: where_conditions["user_id"] = { @@ -1866,6 +1912,12 @@ async def ui_view_users( "mode": "insensitive", # Case-insensitive search } + # Org admins: only users in their org(s) + if not is_proxy_admin and org_admin_org_ids: + where_conditions["organization_memberships"] = { + "some": {"organization_id": {"in": org_admin_org_ids}} + } + # Query users with pagination and filters users: Optional[List[BaseModel]] = ( await prisma_client.db.litellm_usertable.find_many( @@ -1881,6 +1933,8 @@ async def ui_view_users( return [LiteLLM_UserTableFiltered(**user.model_dump()) for user in users] + except HTTPException: + raise except Exception as e: verbose_proxy_logger.exception(f"Error searching users: {str(e)}") raise HTTPException(status_code=500, detail=f"Error searching users: {str(e)}") diff --git a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py index 839885bc752..16b5feb108a 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py @@ -34,7 +34,8 @@ client = TestClient(app) @pytest.mark.asyncio async def test_ui_view_users_with_null_email(mocker, caplog): """ - Test that /user/filter/ui endpoint returns users even when they have null email fields + Test that /user/filter/ui endpoint returns users even when they have null email fields. + Uses proxy admin so no org filtering is applied. """ # Mock the prisma client mock_prisma_client = mocker.MagicMock() @@ -48,19 +49,18 @@ async def test_ui_view_users_with_null_email(mocker, caplog): "created_at": "2024-01-01T00:00:00Z", } - # Setup the mock find_many response - # Setup the mock find_many response as an async function async def mock_find_many(*args, **kwargs): return [mock_user] mock_prisma_client.db.litellm_usertable.find_many = mock_find_many - # Patch the prisma client import in the endpoint mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) - # Call ui_view_users function directly + # Proxy admin: no org filter, no get_user_object call response = await ui_view_users( - user_api_key_dict=UserAPIKeyAuth(user_id="test_user"), + user_api_key_dict=UserAPIKeyAuth( + user_id="test_user", user_role=LitellmUserRoles.PROXY_ADMIN + ), user_id="test_user", user_email=None, page=1, @@ -72,6 +72,121 @@ async def test_ui_view_users_with_null_email(mocker, caplog): ] +@pytest.mark.asyncio +async def test_ui_view_users_proxy_admin_no_org_filter(mocker): + """ + Proxy admin: find_many is called without organization_memberships in where. + """ + mock_prisma_client = mocker.MagicMock() + async def mock_find_many(*args, **kwargs): + assert "organization_memberships" not in (kwargs.get("where") or {}) + return [] + + mock_prisma_client.db.litellm_usertable.find_many = mock_find_many + mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + + await ui_view_users( + user_api_key_dict=UserAPIKeyAuth( + user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN + ), + user_id=None, + user_email="foo", + page=1, + page_size=50, + ) + + +@pytest.mark.asyncio +async def test_ui_view_users_org_admin_filtered_by_org(mocker): + """ + Org admin: find_many is called with organization_memberships filter so only users + in the caller's org(s) are returned. + """ + from litellm.proxy._types import LiteLLM_OrganizationMembershipTable + + mock_prisma_client = mocker.MagicMock() + org_id = "org-123" + + async def mock_find_many(*args, **kwargs): + where = kwargs.get("where") or {} + assert "organization_memberships" in where + assert where["organization_memberships"] == { + "some": {"organization_id": {"in": [org_id]}} + } + return [] + + mock_prisma_client.db.litellm_usertable.find_many = mock_find_many + mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", mocker.MagicMock()) + mocker.patch("litellm.proxy.proxy_server.proxy_logging_obj", mocker.MagicMock()) + + caller_user = mocker.MagicMock() + caller_user.organization_memberships = [ + LiteLLM_OrganizationMembershipTable( + user_id="org-admin", + organization_id=org_id, + user_role=LitellmUserRoles.ORG_ADMIN.value, + created_at=datetime.now(timezone.utc), + updated_at=datetime.now(timezone.utc), + ) + ] + + async def mock_get_user_object(*args, **kwargs): + return caller_user + + mocker.patch( + "litellm.proxy.management_endpoints.internal_user_endpoints.get_user_object", + side_effect=mock_get_user_object, + ) + + response = await ui_view_users( + user_api_key_dict=UserAPIKeyAuth(user_id="org-admin", user_role=None), + user_id=None, + user_email="u", + page=1, + page_size=50, + ) + + assert response == [] + + +@pytest.mark.asyncio +async def test_ui_view_users_non_org_admin_returns_403(mocker): + """ + Caller is not proxy admin and not org admin: endpoint returns 403. + """ + from fastapi import HTTPException + + mock_prisma_client = mocker.MagicMock() + mocker.patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + mocker.patch("litellm.proxy.proxy_server.user_api_key_cache", mocker.MagicMock()) + mocker.patch("litellm.proxy.proxy_server.proxy_logging_obj", mocker.MagicMock()) + + # Caller has no org admin membership + caller_user = mocker.MagicMock() + caller_user.organization_memberships = [] # not an org admin + + async def mock_get_user_object(*args, **kwargs): + return caller_user + + mocker.patch( + "litellm.proxy.management_endpoints.internal_user_endpoints.get_user_object", + side_effect=mock_get_user_object, + ) + + with pytest.raises(HTTPException) as exc_info: + await ui_view_users( + user_api_key_dict=UserAPIKeyAuth(user_id="internal_user", user_role=None), + user_id=None, + user_email="u", + page=1, + page_size=50, + ) + + assert exc_info.value.status_code == 403 + assert "Only proxy admins and organization admins" in str(exc_info.value.detail) + + def test_user_daily_activity_types(): """ Assert all fiels in SpendMetrics are reported in DailySpendMetadata as "total_" From dcfd25e1f1e5a7707ac54fcc168b64ed0d732493 Mon Sep 17 00:00:00 2001 From: David Velarde Date: Sat, 28 Feb 2026 10:56:38 +0100 Subject: [PATCH 013/273] [Feature] Add Gemini 3.1 Flash Image Preview pricing details --- model_prices_and_context_window.json | 33 ++++++++++++++++++++++++++++ 1 file changed, 33 insertions(+) diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index f52288ea72a..5a43447e2c2 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -16421,6 +16421,39 @@ "supports_vision": true, "supports_web_search": true }, + "gemini/gemini-3.1-flash-image-preview": { + "input_cost_per_image": 0.0001375, + "input_cost_per_token": 2.5e-07, + "litellm_provider": "gemini", + "max_input_tokens": 65536, + "max_output_tokens": 32768, + "max_tokens": 32768, + "mode": "image_generation", + "output_cost_per_image": 0.045, + "output_cost_per_image_token": 6e-05, + "output_cost_per_token": 1.5e-06, + "rpm": 1000, + "tpm": 4000000, + "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-3.1-flash-image-preview", + "supported_endpoints": [ + "/v1/chat/completions", + "/v1/completions" + ], + "supported_modalities": [ + "text", + "image" + ], + "supported_output_modalities": [ + "text", + "image" + ], + "supports_function_calling": false, + "supports_prompt_caching": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_vision": true, + "supports_web_search": true + }, "gemini/deep-research-pro-preview-12-2025": { "input_cost_per_image": 0.0011, "input_cost_per_token": 2e-06, From 29d1d0479f3ef7d897fbd7cb707b0744f727101b Mon Sep 17 00:00:00 2001 From: David Velarde Date: Sat, 28 Feb 2026 11:09:38 +0100 Subject: [PATCH 014/273] [Feature] Add Gemini 3.1 Flash Image Preview input and output cost details --- model_prices_and_context_window.json | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 5a43447e2c2..f785fbbbb6e 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -16422,8 +16422,8 @@ "supports_web_search": true }, "gemini/gemini-3.1-flash-image-preview": { - "input_cost_per_image": 0.0001375, "input_cost_per_token": 2.5e-07, + "input_cost_per_token_batches": 1.25e-07, "litellm_provider": "gemini", "max_input_tokens": 65536, "max_output_tokens": 32768, @@ -16431,13 +16431,16 @@ "mode": "image_generation", "output_cost_per_image": 0.045, "output_cost_per_image_token": 6e-05, + "output_cost_per_image_token_batches": 3e-05, "output_cost_per_token": 1.5e-06, + "output_cost_per_token_batches": 7.5e-07, "rpm": 1000, "tpm": 4000000, "source": "https://ai.google.dev/gemini-api/docs/pricing#gemini-3.1-flash-image-preview", "supported_endpoints": [ "/v1/chat/completions", - "/v1/completions" + "/v1/completions", + "/v1/batch" ], "supported_modalities": [ "text", From 8c8d1debee7f91078998f7eac0f15dae57db167c Mon Sep 17 00:00:00 2001 From: Kerem Turgutlu Date: Tue, 3 Mar 2026 08:51:06 +0300 Subject: [PATCH 015/273] fix: preserve usage/cached_tokens in Responses API streaming bridge (#22194) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The response.completed handler in the completion→responses streaming bridge was discarding the usage object, causing prompt_tokens_details (and cached_tokens) to always be None when streaming with models that use the Responses API (e.g. gpt-5.2-codex, gpt-5.3-codex). Extract usage from the response.completed event and translate it via the existing _transform_response_api_usage_to_chat_usage helper. Fixes #22192 --- .../transformation.py | 9 +++- ...responses_transformation_transformation.py | 51 +++++++++++++++++++ 2 files changed, 59 insertions(+), 1 deletion(-) diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index 1704861686e..413c19bfc25 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -1088,6 +1088,12 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): finish_reason = "tool_calls" if has_function_calls else "stop" + usage = None + if response_data.get("usage"): + from litellm.responses.utils import ResponseAPILoggingUtils + usage = ResponseAPILoggingUtils._transform_response_api_usage_to_chat_usage( + response_data.get("usage") + ) return ModelResponseStream( choices=[ StreamingChoices( @@ -1095,7 +1101,8 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): delta=Delta(content=""), finish_reason=finish_reason, ) - ] + ], + usage=usage ) else: pass diff --git a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py index 3021fff9a22..cdafe247990 100644 --- a/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py +++ b/tests/test_litellm/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py @@ -738,6 +738,57 @@ def test_response_completed_with_message_only_emits_stop_finish_reason(): ) + +def test_response_completed_preserves_usage_with_cached_tokens(): + """ + Test that response.completed correctly translates Responses API usage + (input_tokens_details) to chat completion usage (prompt_tokens_details). + + This is a regression test for an issue where streaming with models that + use the Responses API bridge (e.g. gpt-5.2-codex) would drop + prompt_tokens_details, causing cached_tokens to always be None. + """ + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + OpenAiResponsesToChatCompletionStreamIterator, + ) + + iterator = OpenAiResponsesToChatCompletionStreamIterator(streaming_response=None, sync_stream=True) + + chunk = { + "type": "response.completed", + "response": { + "id": "resp_789", + "status": "completed", + "output": [ + { + "type": "message", + "id": "msg_abc", + "role": "assistant", + "content": [{"type": "output_text", "text": "Six"}], + "status": "completed", + } + ], + "usage": { + "input_tokens": 1226, + "output_tokens": 5, + "total_tokens": 1231, + "input_tokens_details": {"cached_tokens": 1024}, + "output_tokens_details": {"reasoning_tokens": 0}, + }, + }, + } + + result = iterator.chunk_parser(chunk) + + assert result.usage is not None, "usage should be set on response.completed chunk" + assert result.usage.prompt_tokens == 1226, "prompt_tokens should map from input_tokens" + assert result.usage.completion_tokens == 5, "completion_tokens should map from output_tokens" + assert result.usage.prompt_tokens_details is not None, "prompt_tokens_details should be set" + assert result.usage.prompt_tokens_details.cached_tokens == 1024, ( + "cached_tokens should be preserved from input_tokens_details" + ) + + def test_function_call_done_emits_is_finished(): """ Test that OUTPUT_ITEM_DONE for a function_call still emits is_finished=True. From 239f044721a34901bfaa1216bf0079149adf0fb3 Mon Sep 17 00:00:00 2001 From: pnookala-godaddy <93624827+pnookala-godaddy@users.noreply.github.com> Date: Mon, 2 Mar 2026 21:52:32 -0800 Subject: [PATCH 016/273] fix(caching): inject default_in_memory_ttl in DualCache async_set_cache and async_set_cache_pipeline (#22241) DualCache.async_set_cache and async_set_cache_pipeline were missing the default_in_memory_ttl injection that the sync set_cache method has. This caused InMemoryCache to fall back to its own default_ttl (600s) instead of using DualCache's configured default_in_memory_ttl (typically 60s). This is particularly impactful for end-user budget enforcement in the proxy, where cached spend values could remain stale for 10 minutes instead of 1 minute, allowing users to exceed their budgets. --- litellm/caching/dual_cache.py | 4 + tests/test_litellm/caching/test_dual_cache.py | 103 ++++++++++++++++++ 2 files changed, 107 insertions(+) diff --git a/litellm/caching/dual_cache.py b/litellm/caching/dual_cache.py index 6df570c72b9..48f4d8b8d3d 100644 --- a/litellm/caching/dual_cache.py +++ b/litellm/caching/dual_cache.py @@ -346,6 +346,8 @@ class DualCache(BaseCache): ) try: if self.in_memory_cache is not None: + if "ttl" not in kwargs and self.default_in_memory_ttl is not None: + kwargs["ttl"] = self.default_in_memory_ttl await self.in_memory_cache.async_set_cache(key, value, **kwargs) if self.redis_cache is not None and local_only is False: @@ -367,6 +369,8 @@ class DualCache(BaseCache): ) try: if self.in_memory_cache is not None: + if "ttl" not in kwargs and self.default_in_memory_ttl is not None: + kwargs["ttl"] = self.default_in_memory_ttl await self.in_memory_cache.async_set_cache_pipeline( cache_list=cache_list, **kwargs ) diff --git a/tests/test_litellm/caching/test_dual_cache.py b/tests/test_litellm/caching/test_dual_cache.py index 9974c23e4b4..606f25ddf44 100644 --- a/tests/test_litellm/caching/test_dual_cache.py +++ b/tests/test_litellm/caching/test_dual_cache.py @@ -1,9 +1,11 @@ import asyncio +import time from unittest.mock import AsyncMock, MagicMock, patch import pytest from litellm.caching.dual_cache import DualCache +from litellm.caching.in_memory_cache import InMemoryCache from litellm.caching.redis_cache import RedisCache @@ -56,3 +58,104 @@ async def test_dual_cache_async_batch_get_cache_rolls_back_redis_reservation_on_ assert mock_async_batch_get_cache.call_count == 2 assert "shared_a" not in dual_cache.last_redis_batch_access_time assert "shared_b" not in dual_cache.last_redis_batch_access_time + + +@pytest.mark.asyncio +async def test_dual_cache_async_set_cache_injects_default_in_memory_ttl(): + """ + Test that async_set_cache injects default_in_memory_ttl into kwargs + when no explicit ttl is provided, matching the sync set_cache behavior. + + Regression test for: async_set_cache was missing the TTL injection that + sync set_cache has, causing InMemoryCache to use its own default_ttl (600s) + instead of DualCache's default_in_memory_ttl. + """ + in_memory_cache = InMemoryCache(default_ttl=600) + dual_cache = DualCache( + in_memory_cache=in_memory_cache, + default_in_memory_ttl=60, + ) + + before = time.time() + await dual_cache.async_set_cache(key="test_key", value="test_value") + after = time.time() + + # The TTL stored should reflect default_in_memory_ttl (60s), not + # InMemoryCache's default_ttl (600s) + expiry = in_memory_cache.ttl_dict["test_key"] + assert expiry >= before + 60 + assert expiry <= after + 60 + + +@pytest.mark.asyncio +async def test_dual_cache_async_set_cache_respects_explicit_ttl(): + """ + Test that async_set_cache does NOT override an explicitly provided ttl. + """ + in_memory_cache = InMemoryCache(default_ttl=600) + dual_cache = DualCache( + in_memory_cache=in_memory_cache, + default_in_memory_ttl=60, + ) + + before = time.time() + await dual_cache.async_set_cache(key="test_key", value="test_value", ttl=30) + after = time.time() + + # The explicit ttl=30 should be used, not default_in_memory_ttl (60) + expiry = in_memory_cache.ttl_dict["test_key"] + assert expiry >= before + 30 + assert expiry <= after + 30 + + +@pytest.mark.asyncio +async def test_dual_cache_async_set_cache_pipeline_injects_default_in_memory_ttl(): + """ + Test that async_set_cache_pipeline injects default_in_memory_ttl into kwargs + when no explicit ttl is provided. + """ + in_memory_cache = InMemoryCache(default_ttl=600) + dual_cache = DualCache( + in_memory_cache=in_memory_cache, + default_in_memory_ttl=60, + ) + + cache_list = [("key_a", "value_a"), ("key_b", "value_b")] + + before = time.time() + await dual_cache.async_set_cache_pipeline(cache_list=cache_list) + after = time.time() + + for key in ["key_a", "key_b"]: + expiry = in_memory_cache.ttl_dict[key] + assert expiry >= before + 60 + assert expiry <= after + 60 + + +@pytest.mark.asyncio +async def test_dual_cache_sync_and_async_set_cache_use_same_ttl(): + """ + Test that sync set_cache and async async_set_cache produce the same TTL + when no explicit ttl is provided, ensuring parity between the two paths. + """ + in_memory_sync = InMemoryCache(default_ttl=600) + dual_cache_sync = DualCache( + in_memory_cache=in_memory_sync, + default_in_memory_ttl=60, + ) + + in_memory_async = InMemoryCache(default_ttl=600) + dual_cache_async = DualCache( + in_memory_cache=in_memory_async, + default_in_memory_ttl=60, + ) + + dual_cache_sync.set_cache(key="test_key", value="test_value") + await dual_cache_async.async_set_cache(key="test_key", value="test_value") + + sync_expiry = in_memory_sync.ttl_dict["test_key"] + async_expiry = in_memory_async.ttl_dict["test_key"] + + # Both should use default_in_memory_ttl=60, so their expiry times + # should be within a small tolerance of each other + assert abs(sync_expiry - async_expiry) < 1.0 From 52c5f2af6bb0649e9e3eee85935742fb47e9facc Mon Sep 17 00:00:00 2001 From: Umut Polat <52835619+umut-polat@users.noreply.github.com> Date: Tue, 3 Mar 2026 08:56:37 +0300 Subject: [PATCH 017/273] fix: apply server root path to mapped passthrough route matching (#22310) mapped passthrough routes (vertex_ai, bedrock, etc) were compared against the raw request path without prepending SERVER_ROOT_PATH. db-registered routes already used _build_full_path_with_root for this but the mapped routes branch was missed. fixes #22272 --- .../pass_through_endpoints.py | 3 +- .../test_pass_through_endpoints.py | 39 +++++++++++++++++++ 2 files changed, 41 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 356807415de..4d95fda0a44 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -2062,7 +2062,8 @@ class InitPassThroughEndpointHelpers: """ ## CHECK IF MAPPED PASS THROUGH ENDPOINT for mapped_route in LiteLLMRoutes.mapped_pass_through_routes.value: - if route.startswith(mapped_route): + full_mapped_route = InitPassThroughEndpointHelpers._build_full_path_with_root(mapped_route) + if route.startswith(full_mapped_route): return True # Fast path: check if any registered route key contains this path diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py index 7ec97ddc185..71420c23ad1 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -2369,3 +2369,42 @@ def test_get_registered_pass_through_route_with_custom_root(): # Clean up _registered_pass_through_routes.clear() + + +def test_mapped_pass_through_routes_with_server_root_path(): + """ + Mapped passthrough routes (vertex_ai, bedrock, etc) should match + even when SERVER_ROOT_PATH is set and the incoming route is prefixed. + + Regression test for https://github.com/BerriAI/litellm/issues/22272 + """ + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( + InitPassThroughEndpointHelpers, + ) + + with patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_server_root_path" + ) as mock_get_root: + mock_get_root.return_value = "/litellm" + + # prefixed route should match mapped routes like /vertex_ai + assert ( + InitPassThroughEndpointHelpers.is_registered_pass_through_route( + "/litellm/vertex_ai/v1/projects/foo" + ) + is True + ) + assert ( + InitPassThroughEndpointHelpers.is_registered_pass_through_route( + "/litellm/bedrock/model/invoke" + ) + is True + ) + + # bare route without prefix should not match when root is set + assert ( + InitPassThroughEndpointHelpers.is_registered_pass_through_route( + "/vertex_ai/v1/projects/foo" + ) + is False + ) From 2e362327b630ce2ce93751ecb020e787fbae7b0a Mon Sep 17 00:00:00 2001 From: Alejandro Tapia <67175024+atapia27@users.noreply.github.com> Date: Tue, 3 Mar 2026 16:04:51 -0800 Subject: [PATCH 018/273] Update litellm/proxy/management_endpoints/internal_user_endpoints.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- litellm/proxy/management_endpoints/internal_user_endpoints.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index f5488ce865d..d88d2810193 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -1858,7 +1858,7 @@ async def ui_view_users( try: # Restrict by caller role: proxy admin sees all; org admin sees only their org(s); others 403 - is_proxy_admin = user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN + is_proxy_admin = _user_has_admin_view(user_api_key_dict) if not is_proxy_admin: if user_api_key_dict.user_id is None: raise HTTPException( From 1c04016d7bced99f9747debe2d6e255c7860292d Mon Sep 17 00:00:00 2001 From: Alejandro Tapia Date: Tue, 3 Mar 2026 16:07:18 -0800 Subject: [PATCH 019/273] Fix: get_user_object raises on missing user, never returns None --- .../internal_user_endpoints.py | 23 +++++++++++++------ 1 file changed, 16 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index d88d2810193..3d799c5731c 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -1867,13 +1867,22 @@ async def ui_view_users( "error": "Only proxy admins and organization admins can search users." }, ) - caller_user = await get_user_object( - user_id=user_api_key_dict.user_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - user_id_upsert=False, - proxy_logging_obj=proxy_logging_obj, - ) + try: + caller_user = await get_user_object( + user_id=user_api_key_dict.user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + user_id_upsert=False, + proxy_logging_obj=proxy_logging_obj, + ) + except ValueError: + # get_user_object raises ValueError when user not found (user_id_upsert=False) + raise HTTPException( + status_code=403, + detail={ + "error": "Only proxy admins and organization admins can search users." + }, + ) if caller_user is None: raise HTTPException( status_code=403, From cb07c75201d5f926361b19de28da302a43cfc15e Mon Sep 17 00:00:00 2001 From: Alejandro Tapia <67175024+atapia27@users.noreply.github.com> Date: Tue, 3 Mar 2026 16:11:31 -0800 Subject: [PATCH 020/273] Update litellm/proxy/management_endpoints/internal_user_endpoints.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- litellm/proxy/management_endpoints/internal_user_endpoints.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 3d799c5731c..70a1801fb1e 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -1922,7 +1922,7 @@ async def ui_view_users( } # Org admins: only users in their org(s) - if not is_proxy_admin and org_admin_org_ids: + if not is_proxy_admin: where_conditions["organization_memberships"] = { "some": {"organization_id": {"in": org_admin_org_ids}} } From 36999b23ee976726631035054ca2f7df3196c62a Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Wed, 4 Mar 2026 13:07:25 +0530 Subject: [PATCH 021/273] [Chore] update mcp documentation for header forwarding --- docs/my-website/docs/mcp.md | 57 +++++++++++++++++++ docs/my-website/docs/mcp_control.md | 8 +-- .../src/components/mcp_tools/mcp_connect.tsx | 2 +- 3 files changed, 62 insertions(+), 5 deletions(-) diff --git a/docs/my-website/docs/mcp.md b/docs/my-website/docs/mcp.md index fcbb31c07d3..c7789201579 100644 --- a/docs/my-website/docs/mcp.md +++ b/docs/my-website/docs/mcp.md @@ -870,6 +870,63 @@ asyncio.run(main()) [Learn more about customer management →](./proxy/customers) +## Calling the Proxy's /v1/responses Endpoint + +When calling your LiteLLM Proxy's `/v1/responses` endpoint to use MCP tools, **always use `server_url: "litellm_proxy"`** in the tools array. This tells the proxy to use its configured MCP servers. + +:::important Do not use the full proxy URL +Using `server_url: "https://your-proxy.com/mcp"` is incorrect when the request is already going to the proxy. The proxy needs the literal value `litellm_proxy` to route to its configured MCP servers. +::: + +```bash title="Correct: Using litellm_proxy" showLineNumbers +curl --location 'https://your-proxy.com/v1/responses' \ +--header 'Content-Type: application/json' \ +--header "Authorization: Bearer $LITELLM_API_KEY" \ +--data '{ + "model": "gpt-4", + "tools": [ + { + "type": "mcp", + "server_label": "litellm", + "server_url": "litellm_proxy", + "require_approval": "never" + } + ], + "input": "Run available tools", + "tool_choice": "required" +}' +``` + +### Sending Custom Headers to MCP Servers + +To pass custom headers (e.g., API keys, auth tokens) to specific MCP servers, use either: + +**Option 1: Request headers** – Add `x-mcp-{server_alias}-{header_name}` to your request headers. The proxy forwards these to the matching MCP server. + +```bash +# Send Authorization header to the "weather2" MCP server +--header 'x-mcp-weather2-authorization: Bearer your-token' + +# Send custom header to the "github" MCP server +--header 'x-mcp-github-x-api-key: your-api-key' +``` + +**Option 2: Headers in tool config** – Include a `headers` object in the tool definition. These are merged with request headers. + +```json +{ + "type": "mcp", + "server_label": "litellm", + "server_url": "litellm_proxy", + "require_approval": "never", + "headers": { + "x-litellm-api-key": "Bearer YOUR_LITELLM_API_KEY", + "x-mcp-servers": "Zapier_MCP,dev-group", + "x-mcp-weather2-authorization": "Bearer your-weather-api-token" + } +} +``` + ## Using your MCP with client side credentials Use this if you want to pass a client side authentication token to LiteLLM to then pass to your MCP to auth to your MCP. diff --git a/docs/my-website/docs/mcp_control.md b/docs/my-website/docs/mcp_control.md index 96c71ef9278..ccaa37f9497 100644 --- a/docs/my-website/docs/mcp_control.md +++ b/docs/my-website/docs/mcp_control.md @@ -323,7 +323,7 @@ curl --location '/v1/responses' \ { "type": "mcp", "server_label": "litellm", - "server_url": "/dev_group/mcp", + "server_url": "litellm_proxy", "require_approval": "never", "headers": { "x-litellm-api-key": "Bearer YOUR_LITELLM_API_KEY" @@ -335,7 +335,7 @@ curl --location '/v1/responses' \ }' ``` -This example uses URL namespacing to access all servers in the "dev_group" access group. +This example uses the `x-mcp-servers` header to access all servers in the "dev_group" access group. Use `server_url: "litellm_proxy"` when calling the proxy's `/v1/responses` endpoint—do not use the full proxy URL. @@ -423,7 +423,7 @@ curl --location '/v1/responses' \ { "type": "mcp", "server_label": "litellm", - "server_url": "/mcp/", + "server_url": "litellm_proxy", "require_approval": "never", "headers": { "x-litellm-api-key": "Bearer YOUR_LITELLM_API_KEY", @@ -436,7 +436,7 @@ curl --location '/v1/responses' \ }' ``` -This configuration restricts the request to only use tools from the specified MCP servers. +This configuration restricts the request to only use tools from the specified MCP servers. Use `server_url: "litellm_proxy"` when calling the proxy's `/v1/responses` endpoint. diff --git a/ui/litellm-dashboard/src/components/mcp_tools/mcp_connect.tsx b/ui/litellm-dashboard/src/components/mcp_tools/mcp_connect.tsx index c48b9a755b7..1c82859e062 100644 --- a/ui/litellm-dashboard/src/components/mcp_tools/mcp_connect.tsx +++ b/ui/litellm-dashboard/src/components/mcp_tools/mcp_connect.tsx @@ -256,7 +256,7 @@ const MCPConnect: React.FC = ({ currentServerAccessGroups = [] { "type": "mcp", "server_label": "litellm", - "server_url": "${proxyBaseUrl}/mcp", + "server_url": "litellm_proxy", "require_approval": "never", "headers": { "x-litellm-api-key": "Bearer YOUR_LITELLM_VIRTUAL_KEY", From 32b387468470b65561798e8f385eea7106aab051 Mon Sep 17 00:00:00 2001 From: Milan Date: Wed, 4 Mar 2026 16:24:09 +0200 Subject: [PATCH 022/273] fix: update Okta SSO docs and custom SSO handler example 1. Okta SSO docs (admin_ui_sso.md): - Rewrite Step 3 to document both Org Auth Server (free) and Custom Auth Server (paid SKU) as tabbed options - Add Step 4 for GENERIC_CLIENT_STATE and PKCE configuration (moved from troubleshooting into the main guide) - Clarify no_matching_policy error only applies to Custom Auth Server - Deduplicate troubleshooting section to reference Step 4 2. Custom SSO handler (custom_sso.py + custom_sso.md): - Replace broken user_info() call with prisma_client.get_data() - user_info() is a FastAPI route handler requiring Request and UserAPIKeyAuth params, cannot be called directly - Keep new_user/add_new_member as commented-out import references in docs for customers who need them --- docs/my-website/docs/proxy/admin_ui_sso.md | 72 +++++++++++++--------- docs/my-website/docs/proxy/custom_sso.md | 16 ++--- litellm/proxy/custom_sso.py | 7 ++- 3 files changed, 54 insertions(+), 41 deletions(-) diff --git a/docs/my-website/docs/proxy/admin_ui_sso.md b/docs/my-website/docs/proxy/admin_ui_sso.md index f88d3480446..2bd4cf24b49 100644 --- a/docs/my-website/docs/proxy/admin_ui_sso.md +++ b/docs/my-website/docs/proxy/admin_ui_sso.md @@ -41,12 +41,38 @@ After creating the app, copy your **Client ID** and **Client Secret** from the a Ensure users are assigned to the app in the **Assignments** tab. If Federation Broker Mode is enabled, you may need to disable it to assign users manually. -#### Step 3: Configure Authorization Server Access Policy +#### Step 3: Set Environment Variables -:::warning Important -This step is required. Without an Access Policy for your app, users will get a `no_matching_policy` error when attempting to log in. +Set the following environment variables. The only difference between the two Okta authorization servers is the endpoint URLs: + +**Org Authorization Server** (available on all Okta plans, no additional SKU required): +```bash +GENERIC_CLIENT_ID="" +GENERIC_CLIENT_SECRET="" +GENERIC_AUTHORIZATION_ENDPOINT="https:///oauth2/v1/authorize" +GENERIC_TOKEN_ENDPOINT="https:///oauth2/v1/token" +GENERIC_USERINFO_ENDPOINT="https:///oauth2/v1/userinfo" +PROXY_BASE_URL="https://" +``` + +**Custom Authorization Server** (requires the Okta API Access Management SKU): +```bash +GENERIC_CLIENT_ID="" +GENERIC_CLIENT_SECRET="" +GENERIC_AUTHORIZATION_ENDPOINT="https:///oauth2/default/v1/authorize" +GENERIC_TOKEN_ENDPOINT="https:///oauth2/default/v1/token" +GENERIC_USERINFO_ENDPOINT="https:///oauth2/default/v1/userinfo" +PROXY_BASE_URL="https://" +``` + +:::tip +You can find all OAuth endpoints at `https:///.well-known/openid-configuration` ::: +#### Step 3a: Configure Access Policy (Custom Authorization Server only) + +If you are using the Custom Authorization Server, you must configure an Access Policy. Without it, users will get a `no_matching_policy` error. Skip this step if you are using the Org Authorization Server. + 1. Go to **Security** → **API** @@ -62,21 +88,21 @@ This step is required. Without an Access Policy for your app, users will get a ` See [Okta's Access Policy documentation](https://help.okta.com/en-us/content/topics/security/api-access-management/access-policies.htm) for more details. -#### Step 4: Configure LiteLLM Environment Variables +#### Step 4: Configure Okta Security Settings + +**GENERIC_CLIENT_STATE** is recommended for Okta to prevent CSRF attacks: ```bash -GENERIC_CLIENT_ID="" -GENERIC_CLIENT_SECRET="" -GENERIC_AUTHORIZATION_ENDPOINT="https:///oauth2/default/v1/authorize" -GENERIC_TOKEN_ENDPOINT="https:///oauth2/default/v1/token" -GENERIC_USERINFO_ENDPOINT="https:///oauth2/default/v1/userinfo" GENERIC_CLIENT_STATE="random-string" -PROXY_BASE_URL="https://" ``` -:::tip -You can find all OAuth endpoints at `https:///.well-known/openid-configuration` -::: +**PKCE (Proof Key for Code Exchange)** — If your Okta application is configured to require PKCE, enable it by setting: + +```bash +GENERIC_CLIENT_USE_PKCE="true" +``` + +LiteLLM will automatically handle PKCE parameter generation and verification during the OAuth flow. #### Step 5: Test the SSO Flow @@ -91,7 +117,7 @@ You can find all OAuth endpoints at `https:///.well-known/open |-------|-------|----------| | `redirect_uri` error | Redirect URI not configured | Add `/sso/callback` to Sign-in redirect URIs in Okta | | `access_denied` | User not assigned to app | Assign the user in the Assignments tab | -| `no_matching_policy` | Missing Access Policy | Create an Access Policy in the Authorization Server (see Step 3) | +| `no_matching_policy` | Missing Access Policy (Custom Authorization Server only) | Create an Access Policy in the Authorization Server (see Step 3a) | @@ -456,23 +482,9 @@ PROXY_BASE_URL=http://litellm.platform.com PROXY_BASE_URL=litellm.platform.com ``` -**2. For Okta specifically, ensure GENERIC_CLIENT_STATE is set** +**2. For Okta specifically, ensure `GENERIC_CLIENT_STATE` is set and PKCE is configured if required** -Okta requires the `GENERIC_CLIENT_STATE` parameter: - -```bash -GENERIC_CLIENT_STATE="random-string" # Required for Okta -``` - -### Okta PKCE - -If your Okta application is configured to require PKCE (Proof Key for Code Exchange), enable it by setting: - -```bash -GENERIC_CLIENT_USE_PKCE="true" -``` - -This is required when your Okta app settings enforce PKCE for enhanced security. LiteLLM will automatically handle PKCE parameter generation and verification during the OAuth flow. +See [Okta SSO — Step 4: Configure Okta Security Settings](#step-4-configure-okta-security-settings) for details on `GENERIC_CLIENT_STATE` and PKCE configuration. ### Common Configuration Issues diff --git a/docs/my-website/docs/proxy/custom_sso.md b/docs/my-website/docs/proxy/custom_sso.md index 8b7adeb0c5a..41ecde6e369 100644 --- a/docs/my-website/docs/proxy/custom_sso.md +++ b/docs/my-website/docs/proxy/custom_sso.md @@ -121,15 +121,14 @@ Use this if you want to run your own code **after** a user signs on to the LiteL Make sure the response type follows the `SSOUserDefinedValues` pydantic object. This is used for logging the user into the Admin UI: ```python -from fastapi import Request from fastapi_sso.sso.base import OpenID from litellm.proxy._types import LitellmUserRoles, SSOUserDefinedValues -from litellm.proxy.management_endpoints.internal_user_endpoints import ( - new_user, - user_info, -) -from litellm.proxy.management_endpoints.team_endpoints import add_new_member +from litellm.proxy import proxy_server + +# These imports are available if you need to create users or manage team membership: +# from litellm.proxy.management_endpoints.internal_user_endpoints import new_user +# from litellm.proxy.management_endpoints.team_endpoints import add_new_member async def custom_sso_handler(userIDPInfo: OpenID) -> SSOUserDefinedValues: @@ -158,8 +157,9 @@ async def custom_sso_handler(userIDPInfo: OpenID) -> SSOUserDefinedValues: ################################################# # Run your custom code / logic here # check if user exists in litellm proxy DB - _user_info = await user_info(user_id=userIDPInfo.id) - print("_user_info from litellm DB ", _user_info) # noqa + if proxy_server.prisma_client is not None: + _user_info = await proxy_server.prisma_client.get_data(user_id=userIDPInfo.id) + print("_user_info from litellm DB ", _user_info) # noqa ################################################# return SSOUserDefinedValues( diff --git a/litellm/proxy/custom_sso.py b/litellm/proxy/custom_sso.py index b2b028dfbe3..1419b5551c8 100644 --- a/litellm/proxy/custom_sso.py +++ b/litellm/proxy/custom_sso.py @@ -15,7 +15,7 @@ Flow: from fastapi_sso.sso.base import OpenID from litellm.proxy._types import LitellmUserRoles, SSOUserDefinedValues -from litellm.proxy.management_endpoints.internal_user_endpoints import user_info +from litellm.proxy import proxy_server async def custom_sso_handler(userIDPInfo: OpenID) -> SSOUserDefinedValues: @@ -32,8 +32,9 @@ async def custom_sso_handler(userIDPInfo: OpenID) -> SSOUserDefinedValues: # user_groups = extra_fields.get("group", []) # check if user exists in litellm proxy DB - _user_info = await user_info(user_id=userIDPInfo.id) - print("_user_info from litellm DB ", _user_info) # noqa + if proxy_server.prisma_client is not None: + _user_info = await proxy_server.prisma_client.get_data(user_id=userIDPInfo.id) + print("_user_info from litellm DB ", _user_info) # noqa return SSOUserDefinedValues( models=[], From 96b75be03d1db6e4957183061fb20e97163318ee Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Wed, 4 Mar 2026 20:13:14 -0800 Subject: [PATCH 023/273] [Feature] RBAC for Vector Stores and Agents Add proxy-admin-configurable toggles to restrict internal users (and optionally team admins) from accessing agent and vector store management features. Co-Authored-By: Claude Sonnet 4.6 --- litellm/proxy/agent_endpoints/endpoints.py | 20 ++- litellm/proxy/common_utils/rbac_utils.py | 126 ++++++++++++++ .../proxy_setting_endpoints.py | 60 +++++-- .../management_endpoints.py | 11 ++ .../proxy/agent_endpoints/test_agent_rbac.py | 84 ++++++++++ .../proxy/common_utils/test_rbac_utils.py | 156 ++++++++++++++++++ .../test_vector_store_rbac.py | 121 ++++++++++++++ .../components/SidebarProvider.tsx | 12 ++ .../AdminSettings/UISettings/UISettings.tsx | 136 +++++++++++++++ .../src/components/leftnav.tsx | 8 +- 10 files changed, 720 insertions(+), 14 deletions(-) create mode 100644 litellm/proxy/common_utils/rbac_utils.py create mode 100644 tests/litellm/proxy/agent_endpoints/test_agent_rbac.py create mode 100644 tests/litellm/proxy/common_utils/test_rbac_utils.py create mode 100644 tests/litellm/proxy/vector_store_endpoints/test_vector_store_rbac.py diff --git a/litellm/proxy/agent_endpoints/endpoints.py b/litellm/proxy/agent_endpoints/endpoints.py index 65674d01be7..80c55f634f7 100644 --- a/litellm/proxy/agent_endpoints/endpoints.py +++ b/litellm/proxy/agent_endpoints/endpoints.py @@ -16,6 +16,7 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.rbac_utils import check_feature_access_for_user from litellm.proxy.management_endpoints.common_daily_activity import get_daily_activity from litellm.types.agents import ( AgentConfig, @@ -69,6 +70,8 @@ async def get_agents( Returns: List[AgentResponse] """ + await check_feature_access_for_user(user_api_key_dict, "agents") + from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( AgentRequestHandler, @@ -179,6 +182,8 @@ async def create_agent( }' ``` """ + await check_feature_access_for_user(user_api_key_dict, "agents") + from litellm.proxy.proxy_server import prisma_client _check_agent_management_permission(user_api_key_dict) @@ -233,7 +238,10 @@ async def create_agent( dependencies=[Depends(user_api_key_auth)], response_model=AgentResponse, ) -async def get_agent_by_id(agent_id: str): +async def get_agent_by_id( + agent_id: str, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): """ Get a specific agent by ID @@ -243,6 +251,8 @@ async def get_agent_by_id(agent_id: str): -H "Authorization: Bearer " ``` """ + await check_feature_access_for_user(user_api_key_dict, "agents") + from litellm.proxy.proxy_server import prisma_client if prisma_client is None: @@ -319,6 +329,8 @@ async def update_agent( }' ``` """ + await check_feature_access_for_user(user_api_key_dict, "agents") + from litellm.proxy.proxy_server import prisma_client _check_agent_management_permission(user_api_key_dict) @@ -410,6 +422,8 @@ async def patch_agent( }' ``` """ + await check_feature_access_for_user(user_api_key_dict, "agents") + from litellm.proxy.proxy_server import prisma_client _check_agent_management_permission(user_api_key_dict) @@ -484,6 +498,8 @@ async def delete_agent( } ``` """ + await check_feature_access_for_user(user_api_key_dict, "agents") + from litellm.proxy.proxy_server import prisma_client _check_agent_management_permission(user_api_key_dict) @@ -763,6 +779,8 @@ async def get_agent_daily_activity( """ Get daily activity for specific agents or all accessible agents. """ + await check_feature_access_for_user(user_api_key_dict, "agents") + from litellm.proxy.proxy_server import prisma_client if prisma_client is None: diff --git a/litellm/proxy/common_utils/rbac_utils.py b/litellm/proxy/common_utils/rbac_utils.py new file mode 100644 index 00000000000..2b187d18065 --- /dev/null +++ b/litellm/proxy/common_utils/rbac_utils.py @@ -0,0 +1,126 @@ +""" +RBAC utility helpers for feature-level access control. + +These helpers are used by agent and vector store endpoints to enforce +proxy-admin-configurable toggles that restrict access for internal users. +""" + +from typing import TYPE_CHECKING + +from fastapi import HTTPException + +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import LiteLLM_TeamTable, LitellmUserRoles, UserAPIKeyAuth + +if TYPE_CHECKING: + pass + + +def _is_user_team_admin_for_any_team( + user_api_key_dict: UserAPIKeyAuth, + teams: list, +) -> bool: + """ + Return True if the user is an admin member in at least one of the given teams. + + Args: + user_api_key_dict: The authenticated user. + teams: List of Prisma team records (from litellm_teamtable.find_many). + """ + for team in teams: + team_obj = LiteLLM_TeamTable(**team.model_dump()) + for member in team_obj.members_with_roles: + if ( + member.user_id is not None + and member.user_id == user_api_key_dict.user_id + and member.role == "admin" + ): + return True + return False + + +async def check_feature_access_for_user( + user_api_key_dict: UserAPIKeyAuth, + feature_name: str, +) -> None: + """ + Raise HTTP 403 if the user's role is blocked from accessing the given feature + by the UI settings stored in general_settings. + + Args: + user_api_key_dict: The authenticated user. + feature_name: Either "agents" or "vector_stores". + """ + # Proxy admins (and view-only admins) are never blocked. + if user_api_key_dict.user_role in ( + LitellmUserRoles.PROXY_ADMIN, + LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, + LitellmUserRoles.PROXY_ADMIN.value, + LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value, + ): + return + + from litellm.proxy.proxy_server import general_settings + + disable_flag = f"disable_{feature_name}_for_internal_users" + allow_team_admins_flag = f"allow_{feature_name}_for_team_admins" + + if not general_settings.get(disable_flag, False): + # Feature is not disabled — allow all authenticated users. + return + + # Feature is disabled. Check if team admins are exempted. + if general_settings.get(allow_team_admins_flag, False): + is_team_admin = await _check_if_team_admin(user_api_key_dict) + if is_team_admin: + return + + raise HTTPException( + status_code=403, + detail={ + "error": f"Access to {feature_name} is disabled for your role. Contact your proxy admin." + }, + ) + + +async def _check_if_team_admin(user_api_key_dict: UserAPIKeyAuth) -> bool: + """ + Return True if the user is a team admin in any team. + Mirrors the logic in management_endpoints/common_utils._user_has_admin_privileges + but scoped to team-admin check only. + """ + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None or user_api_key_dict.user_id is None: + return False + + from litellm.caching import DualCache + from litellm.proxy.auth.auth_checks import get_user_object + + try: + user_obj = await get_user_object( + user_id=user_api_key_dict.user_id, + prisma_client=prisma_client, + user_api_key_cache=DualCache(), + user_id_upsert=False, + proxy_logging_obj=None, + ) + + if user_obj is None: + return False + + if user_obj.teams is None or len(user_obj.teams) == 0: + return False + + teams = await prisma_client.db.litellm_teamtable.find_many( + where={"team_id": {"in": user_obj.teams}} + ) + + return _is_user_team_admin_for_any_team(user_api_key_dict, teams) + + except Exception as e: + verbose_proxy_logger.debug( + f"rbac_utils: error checking team admin status for user " + f"{user_api_key_dict.user_id}: {e}" + ) + return False diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index ceda08d520a..8991dc5fd5c 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -104,6 +104,26 @@ class UISettings(BaseModel): description="If enabled, shows the Projects feature in the UI sidebar and the project field in key management.", ) + disable_agents_for_internal_users: bool = Field( + default=False, + description="If true, internal users cannot access agent management endpoints or the Agents page in the UI.", + ) + + allow_agents_for_team_admins: bool = Field( + default=False, + description="If true, team admins are exempt from the agents disable restriction (only takes effect when disable_agents_for_internal_users is true).", + ) + + disable_vector_stores_for_internal_users: bool = Field( + default=False, + description="If true, internal users cannot access vector store management endpoints or the Vector Stores page in the UI.", + ) + + allow_vector_stores_for_team_admins: bool = Field( + default=False, + description="If true, team admins are exempt from the vector stores disable restriction (only takes effect when disable_vector_stores_for_internal_users is true).", + ) + class UISettingsResponse(SettingsResponse): """Response model for UI settings""" @@ -119,6 +139,10 @@ ALLOWED_UI_SETTINGS_FIELDS = { "require_auth_for_public_ai_hub", "forward_client_headers_to_llm_api", "enable_projects_ui", + "disable_agents_for_internal_users", + "allow_agents_for_team_admins", + "disable_vector_stores_for_internal_users", + "allow_vector_stores_for_team_admins", } @@ -976,14 +1000,20 @@ async def get_ui_settings(): k: v for k, v in ui_settings.items() if k in ALLOWED_UI_SETTINGS_FIELDS } - # Sync forward_client_headers_to_llm_api into general_settings so the proxy - # picks it up at runtime (covers server restart scenarios). - if "forward_client_headers_to_llm_api" in ui_settings: + # Sync runtime flags into general_settings so the proxy picks them up + # at runtime (covers server restart scenarios). + _runtime_flags = [ + "forward_client_headers_to_llm_api", + "disable_agents_for_internal_users", + "allow_agents_for_team_admins", + "disable_vector_stores_for_internal_users", + "allow_vector_stores_for_team_admins", + ] + _flags_to_sync = {k: ui_settings[k] for k in _runtime_flags if k in ui_settings} + if _flags_to_sync: from litellm.proxy.proxy_server import general_settings - general_settings["forward_client_headers_to_llm_api"] = ui_settings[ - "forward_client_headers_to_llm_api" - ] + general_settings.update(_flags_to_sync) # Build config-like object for schema helper config: Dict[str, Any] = {"litellm_settings": {"ui_settings": ui_settings}} @@ -1048,14 +1078,20 @@ async def update_ui_settings( }, ) - # Sync forward_client_headers_to_llm_api to general_settings so the proxy - # picks it up at runtime (general_settings is checked in pre-call utils). - if "forward_client_headers_to_llm_api" in ui_settings: + # Sync runtime flags to general_settings so the proxy picks them up + # at runtime (general_settings is checked in pre-call utils). + _runtime_flags = [ + "forward_client_headers_to_llm_api", + "disable_agents_for_internal_users", + "allow_agents_for_team_admins", + "disable_vector_stores_for_internal_users", + "allow_vector_stores_for_team_admins", + ] + _flags_to_sync = {k: ui_settings[k] for k in _runtime_flags if k in ui_settings} + if _flags_to_sync: from litellm.proxy.proxy_server import general_settings - general_settings["forward_client_headers_to_llm_api"] = ui_settings[ - "forward_client_headers_to_llm_api" - ] + general_settings.update(_flags_to_sync) return { "message": "UI settings updated successfully", diff --git a/litellm/proxy/vector_store_endpoints/management_endpoints.py b/litellm/proxy/vector_store_endpoints/management_endpoints.py index cccbb51f47b..068f4217e0f 100644 --- a/litellm/proxy/vector_store_endpoints/management_endpoints.py +++ b/litellm/proxy/vector_store_endpoints/management_endpoints.py @@ -24,6 +24,7 @@ from litellm.proxy._types import ( ) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper +from litellm.proxy.common_utils.rbac_utils import check_feature_access_for_user from litellm.secret_managers.main import get_secret from litellm.types.vector_stores import ( LiteLLM_ManagedVectorStore, @@ -439,6 +440,8 @@ async def new_vector_store( - vector_store_description: Optional[str] - Description of the vector store - vector_store_metadata: Optional[Dict] - Additional metadata for the vector store """ + await check_feature_access_for_user(user_api_key_dict, "vector_stores") + from litellm.proxy.proxy_server import prisma_client try: @@ -506,6 +509,8 @@ async def list_vector_stores( - page: int - Page number for pagination (default: 1) - page_size: int - Number of items per page (default: 100) """ + await check_feature_access_for_user(user_api_key_dict, "vector_stores") + from litellm.proxy.proxy_server import prisma_client vector_store_map: Dict[str, LiteLLM_ManagedVectorStore] = {} @@ -605,6 +610,8 @@ async def delete_vector_store( Parameters: - vector_store_id: str - ID of the vector store to delete """ + await check_feature_access_for_user(user_api_key_dict, "vector_stores") + from litellm.proxy.proxy_server import prisma_client if prisma_client is None: @@ -687,6 +694,8 @@ async def get_vector_store_info( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """Return a single vector store's details""" + await check_feature_access_for_user(user_api_key_dict, "vector_stores") + from litellm.proxy.proxy_server import prisma_client if prisma_client is None: @@ -770,6 +779,8 @@ async def update_vector_store( Update vector store details in both database and in-memory registry. The updated data is immediately synchronized to the in-memory registry. """ + await check_feature_access_for_user(user_api_key_dict, "vector_stores") + from litellm.proxy.proxy_server import prisma_client from litellm.types.router import GenericLiteLLMParams diff --git a/tests/litellm/proxy/agent_endpoints/test_agent_rbac.py b/tests/litellm/proxy/agent_endpoints/test_agent_rbac.py new file mode 100644 index 00000000000..a863201ddb5 --- /dev/null +++ b/tests/litellm/proxy/agent_endpoints/test_agent_rbac.py @@ -0,0 +1,84 @@ +""" +Tests for RBAC enforcement on agent endpoints. + +Verifies that check_feature_access_for_user is called and that a 403 is +raised when agents are disabled for internal users. +""" + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from fastapi import HTTPException + +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + + +def _make_internal_user(user_id: str = "user-1") -> UserAPIKeyAuth: + return UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER.value, + user_id=user_id, + ) + + +def _make_admin_user(user_id: str = "admin-1") -> UserAPIKeyAuth: + return UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN.value, + user_id=user_id, + ) + + +# --------------------------------------------------------------------------- +# get_agents +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_get_agents_blocked_for_internal_user_when_disabled(): + """get_agents should raise 403 when agents are disabled for internal users.""" + from litellm.proxy.agent_endpoints.endpoints import get_agents + + user = _make_internal_user() + gs = {"disable_agents_for_internal_users": True, "allow_agents_for_team_admins": False} + + request_mock = MagicMock() + with patch.dict("litellm.proxy.proxy_server.general_settings", gs, clear=True): + with pytest.raises(HTTPException) as exc_info: + await get_agents(request=request_mock, user_api_key_dict=user) + assert exc_info.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_get_agents_allowed_when_not_disabled(): + """get_agents should not raise RBAC 403 when agents are not disabled.""" + from litellm.proxy.agent_endpoints.endpoints import get_agents + + user = _make_internal_user() + request_mock = MagicMock() + + with patch.dict("litellm.proxy.proxy_server.general_settings", {}, clear=True): + with patch( + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry", + MagicMock(get_agent_list=MagicMock(return_value=[])), + ): + with patch( + "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.get_allowed_agents", + new=AsyncMock(return_value=[]), + ): + result = await get_agents(request=request_mock, user_api_key_dict=user) + assert result == [] + + +# --------------------------------------------------------------------------- +# get_agent_daily_activity +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_get_agent_daily_activity_blocked_when_disabled(): + from litellm.proxy.agent_endpoints.endpoints import get_agent_daily_activity + + user = _make_internal_user() + gs = {"disable_agents_for_internal_users": True, "allow_agents_for_team_admins": False} + + with patch.dict("litellm.proxy.proxy_server.general_settings", gs, clear=True): + with pytest.raises(HTTPException) as exc_info: + await get_agent_daily_activity(user_api_key_dict=user) + assert exc_info.value.status_code == 403 diff --git a/tests/litellm/proxy/common_utils/test_rbac_utils.py b/tests/litellm/proxy/common_utils/test_rbac_utils.py new file mode 100644 index 00000000000..997a2e19b77 --- /dev/null +++ b/tests/litellm/proxy/common_utils/test_rbac_utils.py @@ -0,0 +1,156 @@ +""" +Tests for litellm/proxy/common_utils/rbac_utils.py + +Covers check_feature_access_for_user for agents and vector_stores features. +""" + +from unittest.mock import AsyncMock, patch + +import pytest +from fastapi import HTTPException + +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.common_utils.rbac_utils import check_feature_access_for_user + + +def _make_user(role: str, user_id: str = "user-1") -> UserAPIKeyAuth: + return UserAPIKeyAuth(user_role=role, user_id=user_id) + + +# general_settings is imported from litellm.proxy.proxy_server inside the +# function, so we patch it via patch.dict on the original dict. +_GS_PATH = "litellm.proxy.proxy_server.general_settings" + + +# --------------------------------------------------------------------------- +# Proxy admin is always allowed +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_proxy_admin_always_allowed(): + user = _make_user(LitellmUserRoles.PROXY_ADMIN.value) + with patch.dict(_GS_PATH, {"disable_agents_for_internal_users": True}): + await check_feature_access_for_user(user, "agents") + + +@pytest.mark.asyncio +async def test_proxy_admin_view_only_always_allowed(): + user = _make_user(LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value) + with patch.dict(_GS_PATH, {"disable_agents_for_internal_users": True}): + await check_feature_access_for_user(user, "agents") + + +# --------------------------------------------------------------------------- +# Feature not disabled — everyone allowed +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_feature_not_disabled_allows_internal_user(): + user = _make_user(LitellmUserRoles.INTERNAL_USER.value) + with patch.dict(_GS_PATH, {}, clear=True): + await check_feature_access_for_user(user, "agents") + + +@pytest.mark.asyncio +async def test_feature_not_disabled_allows_vector_stores(): + user = _make_user(LitellmUserRoles.INTERNAL_USER.value) + with patch.dict(_GS_PATH, {"disable_vector_stores_for_internal_users": False}, clear=True): + await check_feature_access_for_user(user, "vector_stores") + + +# --------------------------------------------------------------------------- +# Feature disabled, team-admin exemption OFF — internal user blocked +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_agents_disabled_blocks_internal_user(): + user = _make_user(LitellmUserRoles.INTERNAL_USER.value) + with patch.dict( + _GS_PATH, + {"disable_agents_for_internal_users": True, "allow_agents_for_team_admins": False}, + clear=True, + ): + with pytest.raises(HTTPException) as exc_info: + await check_feature_access_for_user(user, "agents") + assert exc_info.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_vector_stores_disabled_blocks_internal_user(): + user = _make_user(LitellmUserRoles.INTERNAL_USER.value) + with patch.dict( + _GS_PATH, + {"disable_vector_stores_for_internal_users": True, "allow_vector_stores_for_team_admins": False}, + clear=True, + ): + with pytest.raises(HTTPException) as exc_info: + await check_feature_access_for_user(user, "vector_stores") + assert exc_info.value.status_code == 403 + + +# --------------------------------------------------------------------------- +# Feature disabled, allow_team_admins ON — team admin allowed, non-admin blocked +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_agents_disabled_team_admin_allowed(): + user = _make_user(LitellmUserRoles.INTERNAL_USER.value, user_id="team-admin-user") + with patch.dict( + _GS_PATH, + {"disable_agents_for_internal_users": True, "allow_agents_for_team_admins": True}, + clear=True, + ): + with patch( + "litellm.proxy.common_utils.rbac_utils._check_if_team_admin", + new=AsyncMock(return_value=True), + ): + await check_feature_access_for_user(user, "agents") + + +@pytest.mark.asyncio +async def test_agents_disabled_non_team_admin_blocked(): + user = _make_user(LitellmUserRoles.INTERNAL_USER.value, user_id="regular-user") + with patch.dict( + _GS_PATH, + {"disable_agents_for_internal_users": True, "allow_agents_for_team_admins": True}, + clear=True, + ): + with patch( + "litellm.proxy.common_utils.rbac_utils._check_if_team_admin", + new=AsyncMock(return_value=False), + ): + with pytest.raises(HTTPException) as exc_info: + await check_feature_access_for_user(user, "agents") + assert exc_info.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_vector_stores_disabled_team_admin_allowed(): + user = _make_user(LitellmUserRoles.INTERNAL_USER.value, user_id="team-admin-user") + with patch.dict( + _GS_PATH, + {"disable_vector_stores_for_internal_users": True, "allow_vector_stores_for_team_admins": True}, + clear=True, + ): + with patch( + "litellm.proxy.common_utils.rbac_utils._check_if_team_admin", + new=AsyncMock(return_value=True), + ): + await check_feature_access_for_user(user, "vector_stores") + + +@pytest.mark.asyncio +async def test_vector_stores_disabled_non_team_admin_blocked(): + user = _make_user(LitellmUserRoles.INTERNAL_USER.value, user_id="regular-user") + with patch.dict( + _GS_PATH, + {"disable_vector_stores_for_internal_users": True, "allow_vector_stores_for_team_admins": True}, + clear=True, + ): + with patch( + "litellm.proxy.common_utils.rbac_utils._check_if_team_admin", + new=AsyncMock(return_value=False), + ): + with pytest.raises(HTTPException) as exc_info: + await check_feature_access_for_user(user, "vector_stores") + assert exc_info.value.status_code == 403 diff --git a/tests/litellm/proxy/vector_store_endpoints/test_vector_store_rbac.py b/tests/litellm/proxy/vector_store_endpoints/test_vector_store_rbac.py new file mode 100644 index 00000000000..3eb49bdf114 --- /dev/null +++ b/tests/litellm/proxy/vector_store_endpoints/test_vector_store_rbac.py @@ -0,0 +1,121 @@ +""" +Tests for RBAC enforcement on vector store management endpoints. + +Verifies that check_feature_access_for_user is called and that a 403 is +raised when vector stores are disabled for internal users. +""" + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from fastapi import HTTPException + +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + + +def _make_internal_user(user_id: str = "user-1") -> UserAPIKeyAuth: + return UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER.value, + user_id=user_id, + ) + + +_DISABLED_GS = { + "disable_vector_stores_for_internal_users": True, + "allow_vector_stores_for_team_admins": False, +} + +_ENABLED_GS: dict = {} + + +# --------------------------------------------------------------------------- +# list_vector_stores +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_list_vector_stores_blocked_when_disabled(): + from litellm.proxy.vector_store_endpoints.management_endpoints import list_vector_stores + + user = _make_internal_user() + with patch.dict("litellm.proxy.proxy_server.general_settings", _DISABLED_GS, clear=True): + with pytest.raises(HTTPException) as exc_info: + await list_vector_stores(user_api_key_dict=user) + assert exc_info.value.status_code == 403 + + +@pytest.mark.asyncio +async def test_list_vector_stores_allowed_when_not_disabled(): + """list_vector_stores should not raise 403 when vector stores are not disabled.""" + from litellm.proxy.vector_store_endpoints.management_endpoints import list_vector_stores + + import litellm + user = _make_internal_user() + mock_prisma = MagicMock() + mock_prisma.db.litellm_managedvectorstorestable.find_many = AsyncMock(return_value=[]) + + raised_403 = False + with patch.dict("litellm.proxy.proxy_server.general_settings", _ENABLED_GS, clear=True): + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma): + with patch.object(litellm, "vector_store_registry", None): + with patch( + "litellm.proxy.vector_store_endpoints.management_endpoints.VectorStoreRegistry._get_vector_stores_from_db", + new=AsyncMock(return_value=[]), + ): + try: + await list_vector_stores(user_api_key_dict=user) + except HTTPException as e: + if e.status_code == 403: + raised_403 = True + assert not raised_403, "Should not raise 403 when vector stores are not disabled" + + +# --------------------------------------------------------------------------- +# new_vector_store +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_new_vector_store_blocked_when_disabled(): + from litellm.proxy.vector_store_endpoints.management_endpoints import new_vector_store + from litellm.types.vector_stores import LiteLLM_ManagedVectorStore + + user = _make_internal_user() + vs = LiteLLM_ManagedVectorStore(vector_store_id="vs-1", custom_llm_provider="openai") # type: ignore[call-arg] + + with patch.dict("litellm.proxy.proxy_server.general_settings", _DISABLED_GS, clear=True): + with pytest.raises(HTTPException) as exc_info: + await new_vector_store(vector_store=vs, user_api_key_dict=user) + assert exc_info.value.status_code == 403 + + +# --------------------------------------------------------------------------- +# Admin user is never blocked +# --------------------------------------------------------------------------- + +@pytest.mark.asyncio +async def test_list_vector_stores_admin_not_blocked(): + """Proxy admin should never be blocked, even when vector stores are disabled.""" + from litellm.proxy.vector_store_endpoints.management_endpoints import list_vector_stores + + import litellm + admin = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN.value, + user_id="admin-1", + ) + + mock_prisma = MagicMock() + mock_prisma.db.litellm_managedvectorstorestable.find_many = AsyncMock(return_value=[]) + + raised_403 = False + with patch.dict("litellm.proxy.proxy_server.general_settings", _DISABLED_GS, clear=True): + with patch("litellm.proxy.proxy_server.prisma_client", mock_prisma): + with patch.object(litellm, "vector_store_registry", None): + with patch( + "litellm.proxy.vector_store_endpoints.management_endpoints.VectorStoreRegistry._get_vector_stores_from_db", + new=AsyncMock(return_value=[]), + ): + try: + await list_vector_stores(user_api_key_dict=admin) + except HTTPException as e: + if e.status_code == 403: + raised_403 = True + assert not raised_403, "Admin should not be blocked even when vector stores are disabled" diff --git a/ui/litellm-dashboard/src/app/(dashboard)/components/SidebarProvider.tsx b/ui/litellm-dashboard/src/app/(dashboard)/components/SidebarProvider.tsx index 17f62a20f7d..7dcc3fa8a1a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/components/SidebarProvider.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/components/SidebarProvider.tsx @@ -15,6 +15,8 @@ const SidebarProvider = ({ setPage, defaultSelectedKey, sidebarCollapsed }: Side const { accessToken } = useAuthorized(); const [enabledPagesInternalUsers, setEnabledPagesInternalUsers] = useState(null); const [enableProjectsUI, setEnableProjectsUI] = useState(false); + const [disableAgentsForInternalUsers, setDisableAgentsForInternalUsers] = useState(false); + const [disableVectorStoresForInternalUsers, setDisableVectorStoresForInternalUsers] = useState(false); useEffect(() => { const fetchUISettings = async () => { @@ -39,6 +41,14 @@ const SidebarProvider = ({ setPage, defaultSelectedKey, sidebarCollapsed }: Side if (settings?.values?.enable_projects_ui !== undefined) { setEnableProjectsUI(Boolean(settings.values.enable_projects_ui)); } + + if (settings?.values?.disable_agents_for_internal_users !== undefined) { + setDisableAgentsForInternalUsers(Boolean(settings.values.disable_agents_for_internal_users)); + } + + if (settings?.values?.disable_vector_stores_for_internal_users !== undefined) { + setDisableVectorStoresForInternalUsers(Boolean(settings.values.disable_vector_stores_for_internal_users)); + } } catch (error) { console.error("[SidebarProvider] Failed to fetch UI settings:", error); } @@ -54,6 +64,8 @@ const SidebarProvider = ({ setPage, defaultSelectedKey, sidebarCollapsed }: Side collapsed={sidebarCollapsed} enabledPagesInternalUsers={enabledPagesInternalUsers} enableProjectsUI={enableProjectsUI} + disableAgentsForInternalUsers={disableAgentsForInternalUsers} + disableVectorStoresForInternalUsers={disableVectorStoresForInternalUsers} /> ); }; diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.tsx index 5d99dd2969d..dfc66d3484d 100644 --- a/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.tsx +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/UISettings/UISettings.tsx @@ -19,9 +19,15 @@ export default function UISettings() { const forwardClientHeadersProperty = schema?.properties?.forward_client_headers_to_llm_api; const enableProjectsUIProperty = schema?.properties?.enable_projects_ui; const enabledPagesProperty = schema?.properties?.enabled_ui_pages_internal_users; + const disableAgentsProperty = schema?.properties?.disable_agents_for_internal_users; + const allowAgentsTeamAdminsProperty = schema?.properties?.allow_agents_for_team_admins; + const disableVectorStoresProperty = schema?.properties?.disable_vector_stores_for_internal_users; + const allowVectorStoresTeamAdminsProperty = schema?.properties?.allow_vector_stores_for_team_admins; const values = data?.values ?? {}; const isDisabledForInternalUsers = Boolean(values.disable_model_add_for_internal_users); const isDisabledTeamAdminDeleteTeamUser = Boolean(values.disable_team_admin_delete_team_user); + const isAgentsDisabled = Boolean(values.disable_agents_for_internal_users); + const isVectorStoresDisabled = Boolean(values.disable_vector_stores_for_internal_users); const handleToggle = (checked: boolean) => { updateSettings( @@ -105,6 +111,62 @@ export default function UISettings() { ); }; + const handleToggleDisableAgents = (checked: boolean) => { + updateSettings( + { disable_agents_for_internal_users: checked }, + { + onSuccess: () => { + NotificationManager.success("UI settings updated successfully"); + }, + onError: (error) => { + NotificationManager.fromBackend(error); + }, + }, + ); + }; + + const handleToggleAllowAgentsTeamAdmins = (checked: boolean) => { + updateSettings( + { allow_agents_for_team_admins: checked }, + { + onSuccess: () => { + NotificationManager.success("UI settings updated successfully"); + }, + onError: (error) => { + NotificationManager.fromBackend(error); + }, + }, + ); + }; + + const handleToggleDisableVectorStores = (checked: boolean) => { + updateSettings( + { disable_vector_stores_for_internal_users: checked }, + { + onSuccess: () => { + NotificationManager.success("UI settings updated successfully"); + }, + onError: (error) => { + NotificationManager.fromBackend(error); + }, + }, + ); + }; + + const handleToggleAllowVectorStoresTeamAdmins = (checked: boolean) => { + updateSettings( + { allow_vector_stores_for_team_admins: checked }, + { + onSuccess: () => { + NotificationManager.success("UI settings updated successfully"); + }, + onError: (error) => { + NotificationManager.fromBackend(error); + }, + }, + ); + }; + return ( {isLoading ? ( @@ -211,6 +273,80 @@ export default function UISettings() { + {/* Agents access control */} + + + + Disable agents for internal users + {disableAgentsProperty?.description && ( + {disableAgentsProperty.description} + )} + + + + + + + + Allow agents for team admins + + {allowAgentsTeamAdminsProperty?.description && ( + {allowAgentsTeamAdminsProperty.description} + )} + + + + + + {/* Vector Stores access control */} + + + + Disable vector stores for internal users + {disableVectorStoresProperty?.description && ( + {disableVectorStoresProperty.description} + )} + + + + + + + + Allow vector stores for team admins + + {allowVectorStoresTeamAdminsProperty?.description && ( + {allowVectorStoresTeamAdminsProperty.description} + )} + + + + + {/* Page Visibility for Internal Users */} = ({ setPage, defaultSelectedKey, collapsed = false, enabledPagesInternalUsers, enableProjectsUI }) => { +const Sidebar: React.FC = ({ setPage, defaultSelectedKey, collapsed = false, enabledPagesInternalUsers, enableProjectsUI, disableAgentsForInternalUsers, disableVectorStoresForInternalUsers }) => { const { userId, accessToken, userRole } = useAuthorized(); const { data: organizations } = useOrganizations(); @@ -450,6 +452,10 @@ const Sidebar: React.FC = ({ setPage, defaultSelectedKey, collapse // Hide Projects page if enableProjectsUI is not enabled if (item.key === "projects" && !enableProjectsUI) return false; + // Hide agents and vector-stores pages for non-admin users when disabled + if (!isAdmin && item.key === "agents" && disableAgentsForInternalUsers) return false; + if (!isAdmin && item.key === "vector-stores" && disableVectorStoresForInternalUsers) return false; + // Existing role check if (item.roles && !item.roles.includes(userRole)) return false; From 7c47609f7af6f32cc5777e04bfe19763fa8dd1f2 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Thu, 5 Mar 2026 00:33:17 -0500 Subject: [PATCH 024/273] fix(provider): register bedrock_mantle in model_list and models_by_provider Adds bedrock_mantle_models to the model_list union and models_by_provider dict so models are discoverable via litellm.model_list and litellm.models_by_provider["bedrock_mantle"]. Co-Authored-By: Claude Sonnet 4.6 --- litellm/__init__.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/litellm/__init__.py b/litellm/__init__.py index 57e9cb25f43..ff7ef55c50c 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -962,6 +962,7 @@ model_list = list( | ovhcloud_models | lemonade_models | docker_model_runner_models + | bedrock_mantle_models | set(clarifai_models) ) @@ -1065,6 +1066,7 @@ models_by_provider: dict = { "aws_polly": aws_polly_models, "gigachat": gigachat_models, "llamagate": llamagate_models, + "bedrock_mantle": bedrock_mantle_models } # mapping for those models which have larger equivalents From f1b86366d38d0c090f483db6cd34d98f4452c013 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Thu, 5 Mar 2026 00:33:44 -0500 Subject: [PATCH 025/273] Revert "fix(provider): register bedrock_mantle in model_list and models_by_provider" This reverts commit 7c47609f7af6f32cc5777e04bfe19763fa8dd1f2. --- litellm/__init__.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/litellm/__init__.py b/litellm/__init__.py index ff7ef55c50c..57e9cb25f43 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -962,7 +962,6 @@ model_list = list( | ovhcloud_models | lemonade_models | docker_model_runner_models - | bedrock_mantle_models | set(clarifai_models) ) @@ -1066,7 +1065,6 @@ models_by_provider: dict = { "aws_polly": aws_polly_models, "gigachat": gigachat_models, "llamagate": llamagate_models, - "bedrock_mantle": bedrock_mantle_models } # mapping for those models which have larger equivalents From 028e6871dd5f8611f84c1e2dc853f44e506e5a92 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Thu, 5 Mar 2026 14:27:51 +0530 Subject: [PATCH 026/273] feat(agents): add static_headers and extra_headers fields to schema and types Add two new fields to LiteLLM_AgentsTable: - static_headers (Json): admin-configured headers always sent to the backend agent - extra_headers (String[]): header names to extract from the client request and forward Extend AgentConfig, PatchAgentRequest, and AgentResponse with the same fields. Also remove duplicate spec_path field from LiteLLM_MCPServerTable. Co-Authored-By: Claude Sonnet 4.6 --- litellm/proxy/schema.prisma | 3 ++- litellm/types/agents.py | 6 ++++++ 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 43972724ecc..6f4ef0c24b6 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -63,6 +63,8 @@ model LiteLLM_AgentsTable { agent_name String @unique litellm_params Json? agent_card_params Json + static_headers Json? @default("{}") + extra_headers String[] @default([]) agent_access_groups String[] @default([]) object_permission_id String? object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) @@ -305,7 +307,6 @@ model LiteLLM_MCPServerTable { registration_url String? allow_all_keys Boolean @default(false) available_on_public_internet Boolean @default(true) - spec_path String? is_byok Boolean @default(false) byok_description String[] @default([]) byok_api_key_help_url String? diff --git a/litellm/types/agents.py b/litellm/types/agents.py index 3ad898b1935..7879cae9ff6 100644 --- a/litellm/types/agents.py +++ b/litellm/types/agents.py @@ -179,6 +179,8 @@ class AgentConfig(TypedDict, total=False): agent_card_params: Required[AgentCard] litellm_params: Dict[str, Any] # allow for any future litellm params object_permission: AgentObjectPermission + static_headers: Optional[Dict[str, str]] + extra_headers: Optional[List[str]] class PatchAgentRequest(TypedDict, total=False): @@ -186,6 +188,8 @@ class PatchAgentRequest(TypedDict, total=False): agent_card_params: AgentCard litellm_params: Dict[str, Any] object_permission: AgentObjectPermission + static_headers: Optional[Dict[str, str]] + extra_headers: Optional[List[str]] # Request/Response models for CRUD endpoints @@ -197,6 +201,8 @@ class AgentResponse(BaseModel): litellm_params: Optional[Dict[str, Any]] = None agent_card_params: Dict[str, Any] object_permission: Optional[Dict[str, Any]] = None + static_headers: Optional[Dict[str, str]] = None + extra_headers: Optional[List[str]] = None created_at: Optional[datetime] = None updated_at: Optional[datetime] = None created_by: Optional[str] = None From 07ee1e9886f54b773f4d9de7e3c8181e90d30d6e Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Thu, 5 Mar 2026 14:28:01 +0530 Subject: [PATCH 027/273] feat(agents): persist static_headers and extra_headers in agent registry Update add_agent_to_db, patch_agent_in_db, and update_agent_in_db to read and write the two new header fields when creating or updating agents. Co-Authored-By: Claude Sonnet 4.6 --- .../proxy/agent_endpoints/agent_registry.py | 31 +++++++++++++++++++ 1 file changed, 31 insertions(+) diff --git a/litellm/proxy/agent_endpoints/agent_registry.py b/litellm/proxy/agent_endpoints/agent_registry.py index 159c9fb93d9..550182f966f 100644 --- a/litellm/proxy/agent_endpoints/agent_registry.py +++ b/litellm/proxy/agent_endpoints/agent_registry.py @@ -128,6 +128,14 @@ class AgentRegistry: agent_copy, None, prisma_client ) + # Serialize static_headers + static_headers_obj = agent.get("static_headers") + static_headers_val: Optional[str] = ( + safe_dumps(dict(static_headers_obj)) if static_headers_obj else None + ) + + extra_headers_val: Optional[List[str]] = agent.get("extra_headers") + create_data: Dict[str, Any] = { "agent_name": agent_name, "litellm_params": litellm_params, @@ -137,6 +145,10 @@ class AgentRegistry: "created_at": datetime.now(timezone.utc), "updated_at": datetime.now(timezone.utc), } + if static_headers_val is not None: + create_data["static_headers"] = static_headers_val + if extra_headers_val is not None: + create_data["extra_headers"] = extra_headers_val if object_permission_id is not None: create_data["object_permission_id"] = object_permission_id @@ -214,6 +226,12 @@ class AgentRegistry: update_data["agent_card_params"] = safe_dumps( augment_agent.get("agent_card_params") ) + if agent.get("static_headers") is not None: + update_data["static_headers"] = safe_dumps( + dict(agent.get("static_headers")) # type: ignore + ) + if agent.get("extra_headers") is not None: + update_data["extra_headers"] = agent.get("extra_headers") if agent.get("object_permission") is not None: agent_copy = dict(augment_agent) existing_object_permission_id = existing_agent.get( @@ -281,6 +299,15 @@ class AgentRegistry: ) agent_card_params: str = safe_dumps(agent_card_params_dict) + # Serialize static_headers for update + static_headers_obj_u = agent.get("static_headers") + static_headers_val_u: Optional[str] = ( + safe_dumps(dict(static_headers_obj_u)) + if static_headers_obj_u is not None + else None + ) + extra_headers_val_u: Optional[List[str]] = agent.get("extra_headers") + update_data: Dict[str, Any] = { "agent_name": agent_name, "litellm_params": litellm_params, @@ -288,6 +315,10 @@ class AgentRegistry: "updated_by": updated_by, "updated_at": datetime.now(timezone.utc), } + if static_headers_val_u is not None: + update_data["static_headers"] = static_headers_val_u + if extra_headers_val_u is not None: + update_data["extra_headers"] = extra_headers_val_u if agent.get("object_permission") is not None: existing_agent = await prisma_client.db.litellm_agentstable.find_unique( where={"agent_id": agent_id} From 16a30b55f5493bbf0754aac0dc4ea4c54b681804 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Thu, 5 Mar 2026 14:28:11 +0530 Subject: [PATCH 028/273] feat(agents): add merge_agent_headers utility Mirrors merge_mcp_headers from the MCP server utils. Dynamic headers come first; static (admin-configured) headers overlay and win on conflict. Co-Authored-By: Claude Sonnet 4.6 --- litellm/proxy/agent_endpoints/utils.py | 27 ++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) create mode 100644 litellm/proxy/agent_endpoints/utils.py diff --git a/litellm/proxy/agent_endpoints/utils.py b/litellm/proxy/agent_endpoints/utils.py new file mode 100644 index 00000000000..2b968de54be --- /dev/null +++ b/litellm/proxy/agent_endpoints/utils.py @@ -0,0 +1,27 @@ +"""Utility helpers for A2A agent endpoints.""" + +from typing import Dict, Mapping, Optional + + +def merge_agent_headers( + *, + dynamic_headers: Optional[Mapping[str, str]] = None, + static_headers: Optional[Mapping[str, str]] = None, +) -> Optional[Dict[str, str]]: + """Merge outbound HTTP headers for A2A agent calls. + + Merge rules: + - Start with ``dynamic_headers`` (values extracted from the incoming client request). + - Overlay ``static_headers`` (admin-configured per agent). + + If both contain the same key, ``static_headers`` wins. + """ + merged: Dict[str, str] = {} + + if dynamic_headers: + merged.update({str(k): str(v) for k, v in dynamic_headers.items()}) + + if static_headers: + merged.update({str(k): str(v) for k, v in static_headers.items()}) + + return merged or None From 20a4eea27e71cfc5933670b73747fb46d66dd41d Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Thu, 5 Mar 2026 14:28:28 +0530 Subject: [PATCH 029/273] feat(agents): forward custom headers to backend A2A agents In invoke_agent_a2a: - Extract admin-configured extra_headers from client request by name - Extract convention-based headers (x-a2a-{agent_id/name}-{header}) from client request - Merge with static_headers (static wins on conflict) - Pass merged headers down to asend_message and _handle_stream_message In asend_message / asend_message_streaming: - Accept agent_extra_headers kwarg - Overlay onto LiteLLM internal headers before creating the httpx client Co-Authored-By: Claude Sonnet 4.6 --- litellm/a2a_protocol/main.py | 14 ++++++- .../proxy/agent_endpoints/a2a_endpoints.py | 37 ++++++++++++++++++- 2 files changed, 48 insertions(+), 3 deletions(-) diff --git a/litellm/a2a_protocol/main.py b/litellm/a2a_protocol/main.py index 485b57e311b..6ac88d3a430 100644 --- a/litellm/a2a_protocol/main.py +++ b/litellm/a2a_protocol/main.py @@ -169,6 +169,7 @@ async def asend_message( api_base: Optional[str] = None, litellm_params: Optional[Dict[str, Any]] = None, agent_id: Optional[str] = None, + agent_extra_headers: Optional[Dict[str, str]] = None, **kwargs: Any, ) -> LiteLLMSendMessageResponse: """ @@ -250,9 +251,12 @@ async def asend_message( "Either a2a_client or api_base is required for standard A2A flow" ) trace_id = trace_id or str(uuid.uuid4()) - extra_headers = {"X-LiteLLM-Trace-Id": trace_id} + extra_headers: Dict[str, str] = {"X-LiteLLM-Trace-Id": trace_id} if agent_id: extra_headers["X-LiteLLM-Agent-Id"] = agent_id + # Overlay agent-level headers (agent headers take precedence over LiteLLM internal ones) + if agent_extra_headers: + extra_headers.update(agent_extra_headers) a2a_client = await create_a2a_client( base_url=api_base, extra_headers=extra_headers ) @@ -426,6 +430,7 @@ async def asend_message_streaming( agent_id: Optional[str] = None, metadata: Optional[Dict[str, Any]] = None, proxy_server_request: Optional[Dict[str, Any]] = None, + agent_extra_headers: Optional[Dict[str, str]] = None, ) -> AsyncIterator[Any]: """ Async: Send a streaming message to an A2A agent. @@ -507,7 +512,12 @@ async def asend_message_streaming( raise ValueError( "Either a2a_client or api_base is required for standard A2A flow" ) - a2a_client = await create_a2a_client(base_url=api_base) + streaming_extra_headers: Optional[Dict[str, str]] = None + if agent_extra_headers: + streaming_extra_headers = dict(agent_extra_headers) + a2a_client = await create_a2a_client( + base_url=api_base, extra_headers=streaming_extra_headers + ) # Type assertion: a2a_client is guaranteed to be non-None here assert a2a_client is not None diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index 6bcee14f29e..344070d17fc 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -6,13 +6,14 @@ The A2A SDK can point to LiteLLM's URL and invoke agents registered with LiteLLM """ import json -from typing import Any, Optional +from typing import Any, Dict, Optional from fastapi import APIRouter, Depends, HTTPException, Request, Response from fastapi.responses import JSONResponse, StreamingResponse from litellm._logging import verbose_proxy_logger from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.agent_endpoints.utils import merge_agent_headers from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.types.utils import all_litellm_params @@ -55,6 +56,7 @@ async def _handle_stream_message( metadata: Optional[dict] = None, proxy_server_request: Optional[dict] = None, *, + agent_extra_headers: Optional[Dict[str, str]] = None, user_api_key_dict: Optional[UserAPIKeyAuth] = None, request_data: Optional[dict] = None, proxy_logging_obj: Optional[Any] = None, @@ -105,6 +107,7 @@ async def _handle_stream_message( agent_id=agent_id, metadata=metadata, proxy_server_request=proxy_server_request, + agent_extra_headers=agent_extra_headers, ) if ( @@ -385,6 +388,36 @@ async def invoke_agent_a2a( version=version, ) + # Build merged headers for the backend agent + static_headers: Dict[str, str] = dict(agent.static_headers or {}) + + raw_headers = dict(request.headers) + normalized = {k.lower(): v for k, v in raw_headers.items()} + + dynamic_headers: Dict[str, str] = {} + + # 1. Admin-configured extra_headers: forward named headers from client request + if agent.extra_headers: + for header_name in agent.extra_headers: + val = normalized.get(header_name.lower()) + if val is not None: + dynamic_headers[header_name] = val + + # 2. Convention-based forwarding: x-a2a-{agent_id_or_name}-{header_name} + # Matches both agent_id (UUID) and agent_name (alias), case-insensitive. + for alias in (agent.agent_id.lower(), agent.agent_name.lower()): + prefix = f"x-a2a-{alias}-" + for key, val in normalized.items(): + if key.startswith(prefix): + header_name = key[len(prefix) :] + if header_name: + dynamic_headers[header_name] = val + + agent_extra_headers = merge_agent_headers( + dynamic_headers=dynamic_headers or None, + static_headers=static_headers or None, + ) + # Route through SDK functions if method == "message/send": from a2a.types import MessageSendParams, SendMessageRequest @@ -401,6 +434,7 @@ async def invoke_agent_a2a( metadata=data.get("metadata", {}), proxy_server_request=data.get("proxy_server_request"), litellm_logging_obj=logging_obj, + agent_extra_headers=agent_extra_headers, ) response = await proxy_logging_obj.post_call_success_hook( @@ -425,6 +459,7 @@ async def invoke_agent_a2a( agent_id=agent.agent_id, metadata=data.get("metadata", {}), proxy_server_request=data.get("proxy_server_request"), + agent_extra_headers=agent_extra_headers, user_api_key_dict=user_api_key_dict, request_data=data, proxy_logging_obj=proxy_logging_obj, From 6e9c7c4a8dd8ddce1b911d77e2009aac3de5f9d3 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Thu, 5 Mar 2026 14:28:36 +0530 Subject: [PATCH 030/273] feat(agents): add Prisma migration for agent header columns ALTER TABLE LiteLLM_AgentsTable to add: - static_headers JSONB DEFAULT '{}' - extra_headers TEXT[] DEFAULT ARRAY[]::TEXT[] Co-Authored-By: Claude Sonnet 4.6 --- .../20260305000000_add_agent_headers/migration.sql | 5 +++++ 1 file changed, 5 insertions(+) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260305000000_add_agent_headers/migration.sql diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260305000000_add_agent_headers/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260305000000_add_agent_headers/migration.sql new file mode 100644 index 00000000000..acb35baba96 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260305000000_add_agent_headers/migration.sql @@ -0,0 +1,5 @@ +-- Add static_headers and extra_headers to LiteLLM_AgentsTable + +ALTER TABLE "LiteLLM_AgentsTable" + ADD COLUMN IF NOT EXISTS "static_headers" JSONB DEFAULT '{}', + ADD COLUMN IF NOT EXISTS "extra_headers" TEXT[] DEFAULT ARRAY[]::TEXT[]; From fd53678898b71f6da4384e5984a4d1308f2ee060 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Thu, 5 Mar 2026 14:28:48 +0530 Subject: [PATCH 031/273] test(agents): add tests for A2A custom header forwarding Covers: - Static headers forwarded to backend - Dynamic headers extracted by name (extra_headers config) - Convention-based x-a2a-{agent_id/name}-{header} forwarding - Static headers win over dynamic on conflict - Unrelated x-a2a- prefixes are not forwarded - No-header case leaves existing behaviour unchanged - merge_agent_headers utility unit tests Co-Authored-By: Claude Sonnet 4.6 --- .../agent_endpoints/test_agent_headers.py | 339 ++++++++++++++++++ 1 file changed, 339 insertions(+) create mode 100644 tests/test_litellm/proxy/agent_endpoints/test_agent_headers.py diff --git a/tests/test_litellm/proxy/agent_endpoints/test_agent_headers.py b/tests/test_litellm/proxy/agent_endpoints/test_agent_headers.py new file mode 100644 index 00000000000..b52c0afb0c0 --- /dev/null +++ b/tests/test_litellm/proxy/agent_endpoints/test_agent_headers.py @@ -0,0 +1,339 @@ +""" +Unit tests for A2A agent custom header forwarding. + +Tests cover: +- Static headers forwarded to backend agent +- Dynamic headers extracted from client request and forwarded +- Static headers win over dynamic on conflict +- No headers configured — existing behavior unchanged +- merge_agent_headers utility +""" + +import sys +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + + +# --------------------------------------------------------------------------- +# Helper: build a minimal mock agent +# --------------------------------------------------------------------------- + +def _make_mock_agent( + static_headers=None, + extra_headers=None, + url="http://backend-agent:10001", +): + mock_agent = MagicMock() + mock_agent.agent_id = "agent-123" + mock_agent.agent_card_params = {"url": url, "name": "Test Agent"} + mock_agent.litellm_params = {} + mock_agent.static_headers = static_headers or {} + mock_agent.extra_headers = extra_headers or [] + return mock_agent + + +def _make_mock_request(extra_headers=None, method="message/send"): + """Build a mock FastAPI Request with configurable headers.""" + mock_request = MagicMock() + headers = {"content-type": "application/json"} + if extra_headers: + headers.update(extra_headers) + mock_request.headers = headers + mock_request.json = AsyncMock( + return_value={ + "jsonrpc": "2.0", + "id": "test-id", + "method": method, + "params": { + "message": { + "role": "user", + "parts": [{"kind": "text", "text": "Hello"}], + "messageId": "msg-123", + } + }, + } + ) + return mock_request + + +def _make_a2a_types_module(): + """Return (module, MessageSendParams, SendMessageRequest, SendStreamingMessageRequest).""" + try: + from a2a.types import ( + MessageSendParams, + SendMessageRequest, + SendStreamingMessageRequest, + ) + mock_a2a_types = MagicMock() + mock_a2a_types.MessageSendParams = MessageSendParams + mock_a2a_types.SendMessageRequest = SendMessageRequest + mock_a2a_types.SendStreamingMessageRequest = SendStreamingMessageRequest + return mock_a2a_types + except ImportError: + pass + + def _make_cls(name): + class MockCls: + def __init__(self, **kwargs): + self.__dict__.update(kwargs) + self._kwargs = kwargs + + def model_dump(self, mode="json", exclude_none=False): + result = dict(self._kwargs) + if exclude_none: + result = {k: v for k, v in result.items() if v is not None} + return result + + MockCls.__name__ = name + return MockCls + + mock_a2a_types = MagicMock() + mock_a2a_types.MessageSendParams = _make_cls("MessageSendParams") + mock_a2a_types.SendMessageRequest = _make_cls("SendMessageRequest") + mock_a2a_types.SendStreamingMessageRequest = _make_cls( + "SendStreamingMessageRequest" + ) + return mock_a2a_types + + +async def _invoke(mock_agent, mock_request, mock_asend_message): + """Run invoke_agent_a2a with standard patches applied.""" + from litellm.proxy._types import UserAPIKeyAuth + + mock_user_api_key_dict = UserAPIKeyAuth(api_key="sk-test", user_id="u1") + mock_fastapi_response = MagicMock() + mock_a2a_types = _make_a2a_types_module() + + mock_response = MagicMock() + mock_response.model_dump.return_value = { + "jsonrpc": "2.0", + "id": "test-id", + "result": {"status": "success"}, + } + + with patch( + "litellm.proxy.agent_endpoints.a2a_endpoints._get_agent", + return_value=mock_agent, + ), patch( + "litellm.proxy.agent_endpoints.auth.agent_permission_handler.AgentRequestHandler.is_agent_allowed", + new_callable=AsyncMock, + return_value=True, + ), patch( + "litellm.proxy.common_request_processing.add_litellm_data_to_request", + side_effect=lambda data, **kw: data, + ), patch( + "litellm.a2a_protocol.asend_message", + new_callable=AsyncMock, + return_value=mock_response, + ) as mock_asend, patch( + "litellm.a2a_protocol.create_a2a_client", + new_callable=AsyncMock, + ), patch( + "litellm.proxy.proxy_server.general_settings", + {}, + ), patch( + "litellm.proxy.proxy_server.proxy_config", + MagicMock(), + ), patch( + "litellm.proxy.proxy_server.version", + "1.0.0", + ), patch.dict( + sys.modules, + {"a2a": MagicMock(), "a2a.types": mock_a2a_types}, + ), patch( + "litellm.a2a_protocol.main.A2A_SDK_AVAILABLE", + True, + ): + from litellm.proxy.agent_endpoints.a2a_endpoints import invoke_agent_a2a + + await invoke_agent_a2a( + agent_id="test-agent", + request=mock_request, + fastapi_response=mock_fastapi_response, + user_api_key_dict=mock_user_api_key_dict, + ) + return mock_asend + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_static_headers_forwarded(): + """Static headers configured on the agent are passed to asend_message.""" + mock_agent = _make_mock_agent( + static_headers={"Authorization": "Bearer token123"} + ) + mock_request = _make_mock_request() + + mock_asend = await _invoke(mock_agent, mock_request, None) + + call_kwargs = mock_asend.call_args.kwargs + headers = call_kwargs.get("agent_extra_headers") + assert headers is not None, "agent_extra_headers should not be None" + assert headers.get("Authorization") == "Bearer token123" + + +@pytest.mark.asyncio +async def test_dynamic_headers_forwarded(): + """Dynamic headers listed in extra_headers are extracted from the client request.""" + mock_agent = _make_mock_agent(extra_headers=["x-api-key"]) + mock_request = _make_mock_request(extra_headers={"x-api-key": "secret"}) + + mock_asend = await _invoke(mock_agent, mock_request, None) + + call_kwargs = mock_asend.call_args.kwargs + headers = call_kwargs.get("agent_extra_headers") + assert headers is not None + assert headers.get("x-api-key") == "secret" + + +@pytest.mark.asyncio +async def test_static_overrides_dynamic(): + """When the same header appears in both static and dynamic, static wins.""" + mock_agent = _make_mock_agent( + static_headers={"Authorization": "Bearer static-token"}, + extra_headers=["Authorization"], + ) + # Client sends a different value for Authorization + mock_request = _make_mock_request( + extra_headers={"Authorization": "Bearer dynamic-token"} + ) + + mock_asend = await _invoke(mock_agent, mock_request, None) + + call_kwargs = mock_asend.call_args.kwargs + headers = call_kwargs.get("agent_extra_headers") + assert headers is not None + assert headers.get("Authorization") == "Bearer static-token" + + +@pytest.mark.asyncio +async def test_no_headers(): + """When no headers are configured, agent_extra_headers is None and behaviour is unchanged.""" + mock_agent = _make_mock_agent() # no static_headers or extra_headers + mock_request = _make_mock_request() + + mock_asend = await _invoke(mock_agent, mock_request, None) + + call_kwargs = mock_asend.call_args.kwargs + headers = call_kwargs.get("agent_extra_headers") + assert headers is None + + +# --------------------------------------------------------------------------- +# Convention-based x-a2a-{agent_id/name}-{header_name} tests +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_convention_header_by_agent_name(): + """x-a2a-{agent_name}-{header} is forwarded using the agent name alias.""" + mock_agent = _make_mock_agent() + mock_agent.agent_name = "my-agent" + mock_request = _make_mock_request( + extra_headers={"x-a2a-my-agent-authorization": "Bearer conv-token"} + ) + + mock_asend = await _invoke(mock_agent, mock_request, None) + + headers = mock_asend.call_args.kwargs.get("agent_extra_headers") + assert headers is not None + assert headers.get("authorization") == "Bearer conv-token" + + +@pytest.mark.asyncio +async def test_convention_header_by_agent_id(): + """x-a2a-{agent_id}-{header} is forwarded using the agent UUID.""" + mock_agent = _make_mock_agent() + mock_agent.agent_id = "abc-123" + mock_agent.agent_name = "other-name" + mock_request = _make_mock_request( + extra_headers={"x-a2a-abc-123-x-api-key": "id-secret"} + ) + + mock_asend = await _invoke(mock_agent, mock_request, None) + + headers = mock_asend.call_args.kwargs.get("agent_extra_headers") + assert headers is not None + assert headers.get("x-api-key") == "id-secret" + + +@pytest.mark.asyncio +async def test_convention_header_static_still_wins(): + """Static headers still override convention-based dynamic headers.""" + mock_agent = _make_mock_agent( + static_headers={"authorization": "Bearer static-wins"} + ) + mock_agent.agent_name = "my-agent" + mock_request = _make_mock_request( + extra_headers={"x-a2a-my-agent-authorization": "Bearer conv-value"} + ) + + mock_asend = await _invoke(mock_agent, mock_request, None) + + headers = mock_asend.call_args.kwargs.get("agent_extra_headers") + assert headers is not None + assert headers.get("authorization") == "Bearer static-wins" + + +@pytest.mark.asyncio +async def test_convention_unrelated_prefix_not_forwarded(): + """Headers that start with x-a2a- but target a different agent are ignored.""" + mock_agent = _make_mock_agent() + mock_agent.agent_id = "agent-abc" + mock_agent.agent_name = "my-agent" + mock_request = _make_mock_request( + extra_headers={"x-a2a-other-agent-authorization": "Bearer wrong"} + ) + + mock_asend = await _invoke(mock_agent, mock_request, None) + + headers = mock_asend.call_args.kwargs.get("agent_extra_headers") + assert headers is None + + +# --------------------------------------------------------------------------- +# Direct unit test for the merge utility +# --------------------------------------------------------------------------- + + +def test_merge_agent_headers_util_dynamic_only(): + from litellm.proxy.agent_endpoints.utils import merge_agent_headers + + result = merge_agent_headers(dynamic_headers={"x-key": "val"}) + assert result == {"x-key": "val"} + + +def test_merge_agent_headers_util_static_only(): + from litellm.proxy.agent_endpoints.utils import merge_agent_headers + + result = merge_agent_headers(static_headers={"Authorization": "Bearer tok"}) + assert result == {"Authorization": "Bearer tok"} + + +def test_merge_agent_headers_util_static_wins(): + from litellm.proxy.agent_endpoints.utils import merge_agent_headers + + result = merge_agent_headers( + dynamic_headers={"Authorization": "dynamic", "x-extra": "d"}, + static_headers={"Authorization": "static"}, + ) + assert result == {"Authorization": "static", "x-extra": "d"} + + +def test_merge_agent_headers_util_none_returns_none(): + from litellm.proxy.agent_endpoints.utils import merge_agent_headers + + result = merge_agent_headers() + assert result is None + + +def test_merge_agent_headers_util_empty_dicts_returns_none(): + from litellm.proxy.agent_endpoints.utils import merge_agent_headers + + result = merge_agent_headers(dynamic_headers={}, static_headers={}) + assert result is None From 36d279ab42c20185d435d502f18f895487249ab3 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Thu, 5 Mar 2026 14:34:11 +0530 Subject: [PATCH 032/273] feat(ui/agents): add Authentication Headers section to agent create/edit form Add a new "Authentication Headers" panel to AgentFormFields: - Static Headers: key-value Form.List (always sent to the backend agent, static wins on conflict with dynamic) - Forward Client Headers: Select[tags] of header names to extract from the client request and forward (extra_headers) Update buildAgentDataFromForm to serialize both fields for the API. Update parseAgentForForm to deserialize them back for editing. Covers both the create wizard (add_agent_form) and the edit view (agent_info). Co-Authored-By: Claude Sonnet 4.6 --- .../src/components/agents/agent_config.ts | 26 +++++++ .../components/agents/agent_form_fields.tsx | 70 ++++++++++++++++++- 2 files changed, 94 insertions(+), 2 deletions(-) diff --git a/ui/litellm-dashboard/src/components/agents/agent_config.ts b/ui/litellm-dashboard/src/components/agents/agent_config.ts index f85c4daac66..01041c5cee4 100644 --- a/ui/litellm-dashboard/src/components/agents/agent_config.ts +++ b/ui/litellm-dashboard/src/components/agents/agent_config.ts @@ -269,6 +269,23 @@ export const buildAgentDataFromForm = (values: any, existingAgent?: any) => { agentData.litellm_params = params; } + // static_headers: convert [{header, value}, ...] → {header: value, ...} + if (Array.isArray(values.static_headers) && values.static_headers.length > 0) { + const staticHeaders: Record = {}; + values.static_headers.forEach((entry: { header?: string; value?: string }) => { + const key = entry?.header?.trim(); + if (key) staticHeaders[key] = entry?.value ?? ""; + }); + if (Object.keys(staticHeaders).length > 0) { + agentData.static_headers = staticHeaders; + } + } + + // extra_headers: already an array of strings from Select tags + if (Array.isArray(values.extra_headers) && values.extra_headers.length > 0) { + agentData.extra_headers = values.extra_headers; + } + return agentData; }; @@ -302,5 +319,14 @@ export const parseAgentForForm = (agent: any) => { cost_per_query: agent.litellm_params?.cost_per_query, input_cost_per_token: agent.litellm_params?.input_cost_per_token, output_cost_per_token: agent.litellm_params?.output_cost_per_token, + // static_headers: {key: value} → [{header, value}, ...] + static_headers: agent.static_headers + ? Object.entries(agent.static_headers as Record).map(([header, value]) => ({ + header, + value, + })) + : [], + // extra_headers: already an array of strings + extra_headers: agent.extra_headers ?? [], }; }; diff --git a/ui/litellm-dashboard/src/components/agents/agent_form_fields.tsx b/ui/litellm-dashboard/src/components/agents/agent_form_fields.tsx index d5429d2a3b5..42e55b8c56f 100644 --- a/ui/litellm-dashboard/src/components/agents/agent_form_fields.tsx +++ b/ui/litellm-dashboard/src/components/agents/agent_form_fields.tsx @@ -1,7 +1,7 @@ import React from "react"; -import { Form, Input, Switch, Collapse } from "antd"; +import { Form, Input, Switch, Collapse, Select, Space, Tooltip } from "antd"; import { Button as AntButton } from "antd"; -import { PlusOutlined, MinusCircleOutlined } from "@ant-design/icons"; +import { PlusOutlined, MinusCircleOutlined, InfoCircleOutlined } from "@ant-design/icons"; import { AGENT_FORM_CONFIG, SKILL_FIELD_CONFIG } from "./agent_config"; import CostConfigFields from "./cost_config_fields"; @@ -188,6 +188,72 @@ const AgentFormFields: React.FC = ({ showAgentName = true, ))} )} + + {/* Authentication Headers */} + {shouldShow("auth_headers") && ( + + {/* Static Headers */} + + Static Headers{" "} + + + + + } + > + + {(fields, { add, remove }) => ( + <> + {fields.map(({ key, name, ...restField }) => ( + + + + + + + + remove(name)} style={{ color: "#ff4d4f" }} /> + + ))} + add()} icon={} style={{ width: "100%" }}> + Add Static Header + + + )} + + + + {/* Extra Headers (dynamic forwarding) */} + + Forward Client Headers{" "} + + + + + } + name="extra_headers" + > + + )} + + ); + }; + + return ( + + + + + } + onCancel={handleCancel} + > +
+ {FIELD_GROUPS.map((group, index) => ( +
+ {index > 0 && } + + {group.title} + + {group.subtitle && ( + + {group.subtitle} + + )} + {group.fields.map(renderField)} +
+ ))} +
+
+ ); +}; + +export default EditHashicorpVaultModal; diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/HashicorpVault/HashicorpVault.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/HashicorpVault/HashicorpVault.tsx new file mode 100644 index 00000000000..a2693903c52 --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/HashicorpVault/HashicorpVault.tsx @@ -0,0 +1,250 @@ +"use client"; + +import { useState } from "react"; +import { useHashicorpVaultConfig } from "@/app/(dashboard)/hooks/configOverrides/useHashicorpVaultConfig"; +import { useDeleteHashicorpVaultConfig } from "@/app/(dashboard)/hooks/configOverrides/useDeleteHashicorpVaultConfig"; +import { useUpdateHashicorpVaultConfig } from "@/app/(dashboard)/hooks/configOverrides/useUpdateHashicorpVaultConfig"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import DeleteResourceModal from "@/components/common_components/DeleteResourceModal"; +import NotificationManager from "@/components/molecules/notifications_manager"; +import { testHashicorpVaultConnection } from "@/components/networking"; +import { Alert, Button, Card, Descriptions, Skeleton, Space, Typography } from "antd"; +import { Edit, KeyRound, PlugZap, Trash2 } from "lucide-react"; +import { SENSITIVE_FIELDS, FIELD_LABELS } from "./constants"; +import EditHashicorpVaultModal from "./EditHashicorpVaultModal"; +import HashicorpVaultEmptyPlaceholder from "./HashicorpVaultEmptyPlaceholder"; + +const { Title, Text } = Typography; + +function detectAuthMethod(values: Record): string { + if (values.vault_token) return "Token"; + if (values.approle_role_id || values.approle_secret_id) return "AppRole"; + return "None"; +} + +const descriptionsConfig = { + column: { xxl: 1, xl: 1, lg: 1, md: 1, sm: 1, xs: 1 }, +}; + +export default function HashicorpVault() { + const { accessToken } = useAuthorized(); + const { data, isLoading, isError, error, refetch } = useHashicorpVaultConfig(); + const { mutate: deleteConfig, isPending: isDeleting } = useDeleteHashicorpVaultConfig(accessToken); + const { mutateAsync: updateConfig } = useUpdateHashicorpVaultConfig(accessToken); + + const [isEditModalVisible, setIsEditModalVisible] = useState(false); + const [isDeleteModalOpen, setIsDeleteModalOpen] = useState(false); + const [clearingField, setClearingField] = useState(null); + const [isClearingField, setIsClearingField] = useState(false); + const [isTesting, setIsTesting] = useState(false); + + const rawValues = data?.values ?? {}; + const isConfigured = Boolean(rawValues.vault_addr); + + const handleTestConnection = async () => { + if (!accessToken) return; + setIsTesting(true); + try { + const result = await testHashicorpVaultConnection(accessToken); + NotificationManager.success(result.message || "Connection to Vault successful!"); + } catch (err) { + NotificationManager.fromBackend(err); + } finally { + setIsTesting(false); + } + }; + + const handleDelete = () => { + deleteConfig(undefined, { + onSuccess: () => { + NotificationManager.success("Hashicorp Vault configuration deleted"); + setIsDeleteModalOpen(false); + }, + onError: (err) => { + NotificationManager.fromBackend(err); + }, + }); + }; + + const handleClearField = async () => { + if (!clearingField) return; + setIsClearingField(true); + try { + await updateConfig({ [clearingField]: "" }); + NotificationManager.success(`${FIELD_LABELS[clearingField] ?? clearingField} cleared`); + setClearingField(null); + refetch(); + } catch (err) { + NotificationManager.fromBackend(err); + } finally { + setIsClearingField(false); + } + }; + + const renderValue = (key: string) => { + const value = rawValues[key]; + if (!value) { + return Not configured; + } + if (SENSITIVE_FIELDS.has(key)) { + return ( +
+ {value} +
+ ); + } + return {value}; + }; + + const renderSettings = () => { + // Only show fields that have values, plus auth method + const fieldsToShow = Object.entries(rawValues).filter( + ([_, value]) => value != null && value !== "" + ); + + if (fieldsToShow.length === 0) return null; + + return ( + + + {detectAuthMethod(rawValues)} + + {fieldsToShow.map(([key]) => ( + + {renderValue(key)} + + ))} + + ); + }; + + return ( + <> + {isLoading ? ( + + + + ) : isError ? ( + + + + ) : ( + + + + {/* Header */} +
+
+ +
+ Hashicorp Vault + Manage secret manager configuration +
+
+ +
+ {isConfigured && ( + <> + + + + + )} +
+
+ + {isConfigured && ( + + vault kv put secret/SECRET_NAME key=secret_value +
+ + View documentation + + + } + /> + )} + + {isConfigured ? ( + renderSettings() + ) : ( + setIsEditModalVisible(true)} /> + )} +
+
+
+ )} + + setIsEditModalVisible(false)} + onSuccess={() => { + setIsEditModalVisible(false); + refetch(); + }} + /> + + setIsDeleteModalOpen(false)} + onOk={handleDelete} + confirmLoading={isDeleting} + /> + + setClearingField(null)} + onOk={handleClearField} + confirmLoading={isClearingField} + /> + + ); +} diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/HashicorpVault/HashicorpVaultEmptyPlaceholder.tsx b/ui/litellm-dashboard/src/components/Settings/AdminSettings/HashicorpVault/HashicorpVaultEmptyPlaceholder.tsx new file mode 100644 index 00000000000..49860fc7617 --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/HashicorpVault/HashicorpVaultEmptyPlaceholder.tsx @@ -0,0 +1,30 @@ +import { Empty, Typography, Button } from "antd"; + +const { Title, Paragraph } = Typography; + +interface HashicorpVaultEmptyPlaceholderProps { + onAdd: () => void; +} + +export default function HashicorpVaultEmptyPlaceholder({ onAdd }: HashicorpVaultEmptyPlaceholderProps) { + return ( +
+ + No Vault Configuration Found + + Configure Hashicorp Vault to securely manage provider API keys and secrets + for your LiteLLM deployment. + +
+ } + > + + + + ); +} diff --git a/ui/litellm-dashboard/src/components/Settings/AdminSettings/HashicorpVault/constants.ts b/ui/litellm-dashboard/src/components/Settings/AdminSettings/HashicorpVault/constants.ts new file mode 100644 index 00000000000..ef924f5f122 --- /dev/null +++ b/ui/litellm-dashboard/src/components/Settings/AdminSettings/HashicorpVault/constants.ts @@ -0,0 +1,20 @@ +export const SENSITIVE_FIELDS = new Set([ + "vault_token", + "approle_role_id", + "approle_secret_id", + "client_key", +]); + +export const FIELD_LABELS: Record = { + vault_addr: "Vault Address", + vault_namespace: "Namespace", + vault_mount_name: "KV Mount Name", + vault_path_prefix: "Path Prefix", + vault_token: "Token", + approle_role_id: "Role ID", + approle_secret_id: "Secret ID", + approle_mount_path: "Mount Path", + client_cert: "Client Certificate", + client_key: "Client Key", + vault_cert_role: "Certificate Role", +}; diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 0df6d813d7c..a8f6013726c 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -9659,6 +9659,95 @@ export const updateUiSettings = async (accessToken: string, settings: Record { + const proxyBaseUrl = getProxyBaseUrl(); + const url = proxyBaseUrl + ? `${proxyBaseUrl}/config_overrides/hashicorp_vault` + : `/config_overrides/hashicorp_vault`; + const response = await fetch(url, { + method: "GET", + headers: { + [globalLitellmHeaderName]: `Bearer ${accessToken}`, + }, + }); + if (!response.ok) { + const errorData = await response.json(); + const detail = errorData?.detail; + const errorMessage = + (typeof detail === "object" && detail?.error) || + (typeof detail === "string" && detail) || + deriveErrorMessage(errorData); + throw new Error(errorMessage); + } + const data = await response.json(); + return data; +}; + +export const updateHashicorpVaultConfig = async ( + accessToken: string, + config: Record, +) => { + const proxyBaseUrl = getProxyBaseUrl(); + const url = proxyBaseUrl + ? `${proxyBaseUrl}/config_overrides/hashicorp_vault` + : `/config_overrides/hashicorp_vault`; + const response = await fetch(url, { + method: "POST", + headers: { + [globalLitellmHeaderName]: `Bearer ${accessToken}`, + "Content-Type": "application/json", + }, + body: JSON.stringify(config), + }); + if (!response.ok) { + const errorData = await response.json(); + const errorMessage = deriveErrorMessage(errorData); + throw new Error(errorMessage); + } + const data = await response.json(); + return data; +}; + +export const deleteHashicorpVaultConfig = async (accessToken: string) => { + const proxyBaseUrl = getProxyBaseUrl(); + const url = proxyBaseUrl + ? `${proxyBaseUrl}/config_overrides/hashicorp_vault` + : `/config_overrides/hashicorp_vault`; + const response = await fetch(url, { + method: "DELETE", + headers: { + [globalLitellmHeaderName]: `Bearer ${accessToken}`, + }, + }); + if (!response.ok) { + const errorData = await response.json(); + const errorMessage = deriveErrorMessage(errorData); + throw new Error(errorMessage); + } + const data = await response.json(); + return data; +}; + +export const testHashicorpVaultConnection = async (accessToken: string) => { + const proxyBaseUrl = getProxyBaseUrl(); + const url = proxyBaseUrl + ? `${proxyBaseUrl}/config_overrides/hashicorp_vault/test_connection` + : `/config_overrides/hashicorp_vault/test_connection`; + const response = await fetch(url, { + method: "POST", + headers: { + [globalLitellmHeaderName]: `Bearer ${accessToken}`, + }, + }); + if (!response.ok) { + const errorData = await response.json(); + const errorMessage = deriveErrorMessage(errorData); + throw new Error(errorMessage); + } + const data = await response.json(); + return data; +}; + // ============================================================ // Claude Code Marketplace Networking Functions // ============================================================ From 21718d208d78eec53e668891956bffc7d5d7fd32 Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Thu, 5 Mar 2026 16:34:11 -0800 Subject: [PATCH 060/273] feat: Hashicorp Vault config override backend endpoints Add CRUD endpoints for managing Hashicorp Vault configuration via the proxy admin API, with background sync, env var management, and connection testing. Fix pre-existing bug where premium check ran after global state mutation, and guard DELETE against clearing non-Vault secret managers. --- .../config_override_endpoints.py | 405 ++++++++++++++++++ litellm/proxy/proxy_server.py | 69 +++ litellm/proxy/schema.prisma | 8 + .../hashicorp_secret_manager.py | 21 +- .../management_endpoints/config_overrides.py | 64 +++ .../test_config_override_endpoints.py | 251 +++++++++++ 6 files changed, 809 insertions(+), 9 deletions(-) create mode 100644 litellm/proxy/management_endpoints/config_override_endpoints.py create mode 100644 litellm/types/proxy/management_endpoints/config_overrides.py create mode 100644 tests/test_litellm/proxy/management_endpoints/test_config_override_endpoints.py diff --git a/litellm/proxy/management_endpoints/config_override_endpoints.py b/litellm/proxy/management_endpoints/config_override_endpoints.py new file mode 100644 index 00000000000..2978a523fb1 --- /dev/null +++ b/litellm/proxy/management_endpoints/config_override_endpoints.py @@ -0,0 +1,405 @@ +import json +import os +from typing import Any, Dict, Set + +from fastapi import APIRouter, Depends, HTTPException +from prisma.errors import RecordNotFoundError +from pydantic import TypeAdapter + +import litellm +from litellm._logging import verbose_proxy_logger +from litellm.llms.custom_httpx.http_handler import _get_httpx_client +from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker +from litellm.proxy._types import CommonProxyErrors, KeyManagementSystem, LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.types.proxy.management_endpoints.config_overrides import ( + ConfigOverrideSettingsResponse, + HashicorpVaultConfig, +) + +router = APIRouter() + +# --- Hashicorp Vault constants --- + +HASHICORP_ENV_VAR_MAPPING: Dict[str, str] = { + "vault_addr": "HCP_VAULT_ADDR", + "vault_token": "HCP_VAULT_TOKEN", + "approle_role_id": "HCP_VAULT_APPROLE_ROLE_ID", + "approle_secret_id": "HCP_VAULT_APPROLE_SECRET_ID", + "approle_mount_path": "HCP_VAULT_APPROLE_MOUNT_PATH", + "client_cert": "HCP_VAULT_CLIENT_CERT", + "client_key": "HCP_VAULT_CLIENT_KEY", + "vault_cert_role": "HCP_VAULT_CERT_ROLE", + "vault_namespace": "HCP_VAULT_NAMESPACE", + "vault_mount_name": "HCP_VAULT_MOUNT_NAME", + "vault_path_prefix": "HCP_VAULT_PATH_PREFIX", +} + +HASHICORP_SENSITIVE_FIELDS: Set[str] = { + "vault_token", + "approle_role_id", + "approle_secret_id", + "client_key", +} + +_sensitive_masker = SensitiveDataMasker() + + +# --- Shared helpers --- + + +def _mask_sensitive_fields( + data: Dict[str, Any], sensitive_fields: Set[str] +) -> Dict[str, Any]: + """Mask sensitive fields for API responses. Non-sensitive fields are left as-is.""" + masked = {} + for key, value in data.items(): + if value is not None and key in sensitive_fields and isinstance(value, str): + masked[key] = _sensitive_masker._mask_value(value) + else: + masked[key] = value + return masked + + +def _get_current_env_values(env_var_mapping: Dict[str, str]) -> Dict[str, Any]: + """Read current env var values as fallback when no DB record exists.""" + values = {} + for field_name, env_var_name in env_var_mapping.items(): + env_value = os.environ.get(env_var_name) + values[field_name] = env_value + return values + + +def _extract_field_type(field_info: Dict[str, Any]) -> str: + """Extract the non-null type from a Pydantic v2 JSON schema field.""" + if "type" in field_info: + return field_info["type"] + for option in field_info.get("anyOf", []): + if option.get("type") != "null": + return option.get("type", "string") + return "string" + + +def _build_field_schema(model_class: type) -> Dict[str, Any]: + """Build field_schema dict from a Pydantic model for UI rendering.""" + schema = TypeAdapter(model_class).json_schema(by_alias=True) + properties = {} + for field_name, field_info in schema.get("properties", {}).items(): + properties[field_name] = { + "description": field_info.get("description", ""), + "type": _extract_field_type(field_info), + } + return { + "description": schema.get("description", ""), + "properties": properties, + } + + +def _parse_config_value(raw: Any) -> Dict[str, Any]: + """Parse a config_value from DB (may be JSON string or dict).""" + if isinstance(raw, str): + return json.loads(raw) + return dict(raw) + + +def _set_env_vars(config_data: Dict[str, Any]) -> None: + """Set HCP_VAULT_* env vars from config data. Unsets vars for missing/None/empty fields.""" + for field_name, env_var_name in HASHICORP_ENV_VAR_MAPPING.items(): + value = config_data.get(field_name) + if value is not None and value != "": + os.environ[env_var_name] = str(value) + else: + os.environ.pop(env_var_name, None) + + +def _clear_hashicorp_vault_state(proxy_config: Any) -> None: + """Clear all Hashicorp Vault state: env vars, secret manager, and change-detection cache.""" + _set_env_vars({}) + if litellm._key_management_system == KeyManagementSystem.HASHICORP_VAULT: + litellm.secret_manager_client = None + litellm._key_management_system = None + proxy_config._last_hashicorp_vault_config = None + + +# --- Hashicorp Vault endpoints --- + + +@router.post( + "/config_overrides/hashicorp_vault", + tags=["Config Overrides"], + dependencies=[Depends(user_api_key_auth)], +) +async def update_hashicorp_vault_config( + config: HashicorpVaultConfig, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Update Hashicorp Vault secret manager configuration. + Sets environment variables, encrypts sensitive fields, and stores in DB. + Reinitializes the secret manager on this pod. + """ + from litellm.proxy.proxy_server import prisma_client, proxy_config + + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + raise HTTPException( + status_code=403, + detail="Only admin users can update config overrides", + ) + + if prisma_client is None: + raise HTTPException( + status_code=500, + detail=CommonProxyErrors.db_not_connected_error.value, + ) + + config_data = config.model_dump(exclude_none=True) + + # Merge ALL fields the user didn't send: try DB first, fall back to env vars. + # Omitted field = keep existing; empty string = clear/remove the field. + existing_record = await prisma_client.db.litellm_configoverrides.find_unique( + where={"config_type": "hashicorp_vault"} + ) + if existing_record is not None and existing_record.config_value is not None: + existing_data = _parse_config_value(existing_record.config_value) + existing_decrypted = proxy_config._decrypt_db_variables(existing_data) + for field in HASHICORP_ENV_VAR_MAPPING: + if field not in config_data and existing_decrypted.get(field): + config_data[field] = existing_decrypted[field] + else: + # No DB record yet — merge from current env vars + env_values = _get_current_env_values(HASHICORP_ENV_VAR_MAPPING) + for field in HASHICORP_ENV_VAR_MAPPING: + if field not in config_data and env_values.get(field): + config_data[field] = env_values[field] + + # Strip empty strings — they signal "clear this field" + config_data = {k: v for k, v in config_data.items() if v != ""} + + # Validate that the config has enough fields to initialize + has_vault_addr = bool(config_data.get("vault_addr")) + has_token_auth = bool(config_data.get("vault_token")) + has_approle_auth = bool( + config_data.get("approle_role_id") and config_data.get("approle_secret_id") + ) + has_tls_cert_auth = bool( + config_data.get("client_cert") and config_data.get("client_key") + ) + + if not has_vault_addr: + raise HTTPException( + status_code=400, + detail="Vault Address is required", + ) + + if not has_token_auth and not has_approle_auth and not has_tls_cert_auth: + raise HTTPException( + status_code=400, + detail="At least one authentication method is required: " + "provide a Token, both AppRole Role ID and Secret ID, " + "or both Client Certificate and Client Key", + ) + + # Snapshot current env vars so we can restore on failure + previous_env = _get_current_env_values(HASHICORP_ENV_VAR_MAPPING) + + # Set env vars and verify the secret manager can initialize before persisting + _set_env_vars(config_data) + + try: + proxy_config.initialize_secret_manager( + key_management_system="hashicorp_vault" + ) + except Exception as e: + _set_env_vars(previous_env) + verbose_proxy_logger.exception( + "Error reinitializing Hashicorp Vault secret manager: %s", str(e) + ) + raise HTTPException( + status_code=500, + detail="Failed to initialize secret manager", + ) + + # Only persist to DB after successful init + encrypted_data = proxy_config._encrypt_env_variables(config_data) + config_value = json.dumps(encrypted_data) + await prisma_client.db.litellm_configoverrides.upsert( + where={"config_type": "hashicorp_vault"}, + data={ + "create": { + "config_type": "hashicorp_vault", + "config_value": config_value, + }, + "update": { + "config_value": config_value, + }, + }, + ) + + # Update change-detection cache so the background reload doesn't redundantly re-init + proxy_config._last_hashicorp_vault_config = json.loads(config_value) + + return { + "message": "Hashicorp Vault configuration updated successfully", + "status": "success", + } + + +@router.get( + "/config_overrides/hashicorp_vault", + tags=["Config Overrides"], + dependencies=[Depends(user_api_key_auth)], + response_model=ConfigOverrideSettingsResponse, +) +async def get_hashicorp_vault_config( + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Get current Hashicorp Vault configuration. + Returns decrypted values from DB, or falls back to current env vars. + """ + from litellm.proxy.proxy_server import prisma_client, proxy_config + + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + raise HTTPException( + status_code=403, + detail="Only admin users can view config overrides", + ) + + if prisma_client is None: + raise HTTPException( + status_code=500, + detail=CommonProxyErrors.db_not_connected_error.value, + ) + + field_schema = _build_field_schema(HashicorpVaultConfig) + + # Try to load from DB + db_record = await prisma_client.db.litellm_configoverrides.find_unique( + where={"config_type": "hashicorp_vault"} + ) + + if db_record is not None and db_record.config_value is not None: + config_data = _parse_config_value(db_record.config_value) + + # Decrypt then mask sensitive fields so plaintext secrets are never sent to the UI + decrypted_data = proxy_config._decrypt_db_variables(config_data) + masked_data = _mask_sensitive_fields( + decrypted_data, HASHICORP_SENSITIVE_FIELDS + ) + + return ConfigOverrideSettingsResponse( + config_type="hashicorp_vault", + values=masked_data, + field_schema=field_schema, + ) + + # Fallback to env vars — also mask sensitive values + env_values = _get_current_env_values(HASHICORP_ENV_VAR_MAPPING) + masked_env_values = _mask_sensitive_fields( + env_values, HASHICORP_SENSITIVE_FIELDS + ) + + return ConfigOverrideSettingsResponse( + config_type="hashicorp_vault", + values=masked_env_values, + field_schema=field_schema, + ) + + +@router.delete( + "/config_overrides/hashicorp_vault", + tags=["Config Overrides"], + dependencies=[Depends(user_api_key_auth)], +) +async def delete_hashicorp_vault_config( + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """Delete Hashicorp Vault configuration. Idempotent.""" + from litellm.proxy.proxy_server import prisma_client, proxy_config + + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + raise HTTPException( + status_code=403, + detail="Only admin users can delete config overrides", + ) + + if prisma_client is None: + raise HTTPException( + status_code=500, + detail=CommonProxyErrors.db_not_connected_error.value, + ) + + # Delete DB record if it exists — ignore if not found + try: + await prisma_client.db.litellm_configoverrides.delete( + where={"config_type": "hashicorp_vault"} + ) + except RecordNotFoundError: + verbose_proxy_logger.debug( + "No existing Hashicorp Vault config record to delete" + ) + + _clear_hashicorp_vault_state(proxy_config) + + return { + "message": "Hashicorp Vault configuration deleted successfully", + "status": "success", + } + + +@router.post( + "/config_overrides/hashicorp_vault/test_connection", + tags=["Config Overrides"], + dependencies=[Depends(user_api_key_auth)], +) +async def test_hashicorp_vault_connection( + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Test the connection to the currently configured Hashicorp Vault. + Uses the already-initialized secret manager client. Does not modify any state. + """ + from litellm.secret_managers.hashicorp_secret_manager import ( + HashicorpSecretManager, + ) + + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: + raise HTTPException( + status_code=403, + detail="Only admin users can test Vault connection", + ) + + client = litellm.secret_manager_client + if not isinstance(client, HashicorpSecretManager): + raise HTTPException( + status_code=400, + detail="Hashicorp Vault is not configured. Save a configuration first.", + ) + + # Step 1: Authenticate (exercises AppRole login, TLS cert login, or direct token) + try: + headers = client._get_request_headers() + except Exception as e: + raise HTTPException( + status_code=502, + detail="Vault authentication failed", + ) + + # Step 2: Verify the token is valid via token/lookup-self + try: + sync_client = _get_httpx_client() + lookup_url = f"{client.vault_addr}/v1/auth/token/lookup-self" + if client.vault_namespace: + headers["X-Vault-Namespace"] = client.vault_namespace + response = sync_client.get(lookup_url, headers=headers) + response.raise_for_status() + except Exception as e: + raise HTTPException( + status_code=502, + detail="Vault token validation failed", + ) + + return { + "status": "success", + "message": f"Successfully connected to Vault at {client.vault_addr}", + } diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index bc2728c2203..f409774c7c1 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -346,6 +346,9 @@ from litellm.proxy.management_endpoints.cache_settings_endpoints import ( from litellm.proxy.management_endpoints.callback_management_endpoints import ( router as callback_management_endpoints_router, ) +from litellm.proxy.management_endpoints.config_override_endpoints import ( + router as config_override_router, +) from litellm.proxy.management_endpoints.common_utils import ( _user_has_admin_privileges, admin_can_invite_user, @@ -2235,6 +2238,7 @@ class ProxyConfig: def __init__(self) -> None: self.config: Dict[str, Any] = {} self._last_semantic_filter_config: Optional[Dict[str, Any]] = None + self._last_hashicorp_vault_config: Optional[Dict[str, Any]] = None def is_yaml(self, config_file_path: str) -> bool: if not os.path.isfile(config_file_path): @@ -4432,6 +4436,11 @@ class ProxyConfig: if self._should_load_db_object(object_type="semantic_filter_settings"): await self._init_semantic_filter_settings_in_db(prisma_client=prisma_client) + if self._should_load_db_object(object_type="config_overrides"): + await self._init_hashicorp_vault_config_override( + prisma_client=prisma_client + ) + async def _init_semantic_filter_settings_in_db(self, prisma_client: PrismaClient): """ Initialize MCP semantic filter settings from database. @@ -4541,6 +4550,65 @@ class ProxyConfig: ) ) + async def _init_hashicorp_vault_config_override( + self, prisma_client: PrismaClient + ): + """ + Load Hashicorp Vault config override from DB. + Decrypts sensitive fields, sets HCP_VAULT_* env vars, and reinitializes the secret manager. + Called periodically via _init_non_llm_objects_in_db to sync config across pods. + """ + from litellm.proxy.management_endpoints.config_override_endpoints import ( + HASHICORP_ENV_VAR_MAPPING, + _clear_hashicorp_vault_state, + _get_current_env_values, + _parse_config_value, + _set_env_vars, + ) + + try: + db_record = await prisma_client.db.litellm_configoverrides.find_unique( + where={"config_type": "hashicorp_vault"} + ) + + if db_record is None or db_record.config_value is None: + if self._last_hashicorp_vault_config is not None: + _clear_hashicorp_vault_state(self) + return + + config_data = _parse_config_value(db_record.config_value) + + # Skip reinit if config hasn't changed since last poll + if self._last_hashicorp_vault_config == config_data: + return + + # Decrypt all fields and set env vars + decrypted_data = self._decrypt_db_variables(config_data) + + # Snapshot current env vars so we can restore on failure + previous_env = _get_current_env_values(HASHICORP_ENV_VAR_MAPPING) + _set_env_vars(decrypted_data) + + # Reinitialize the secret manager + try: + self.initialize_secret_manager( + key_management_system="hashicorp_vault" + ) + except Exception: + # Restore previous working env vars instead of wiping all + _set_env_vars(previous_env) + raise + + self._last_hashicorp_vault_config = config_data.copy() + verbose_proxy_logger.debug( + "Hashicorp Vault config override loaded from DB" + ) + except Exception as e: + verbose_proxy_logger.exception( + "Error loading Hashicorp Vault config override from DB: %s", + str(e), + ) + async def _check_and_reload_model_cost_map(self, prisma_client: PrismaClient): """ Check if model cost map needs to be reloaded based on database configuration. @@ -12971,6 +13039,7 @@ app.include_router(cost_tracking_settings_router) app.include_router(router_settings_router) app.include_router(fallback_management_router) app.include_router(cache_settings_router) +app.include_router(config_override_router) app.include_router(user_agent_analytics_router) app.include_router(enterprise_router) app.include_router(ui_discovery_endpoints_router) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index f18556ac329..fa646808a4a 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -1000,6 +1000,14 @@ model LiteLLM_UISettings { updated_at DateTime @updatedAt } +// Generic config overrides table - one row per config_type +model LiteLLM_ConfigOverrides { + config_type String @id + config_value Json + created_at DateTime @default(now()) + updated_at DateTime @updatedAt +} + // Skills table for storing LiteLLM-managed skills model LiteLLM_SkillsTable { skill_id String @id @default(uuid()) diff --git a/litellm/secret_managers/hashicorp_secret_manager.py b/litellm/secret_managers/hashicorp_secret_manager.py index c59f2ef638a..ccee5018eec 100644 --- a/litellm/secret_managers/hashicorp_secret_manager.py +++ b/litellm/secret_managers/hashicorp_secret_manager.py @@ -44,6 +44,11 @@ class HashicorpSecretManager(BaseSecretManager): self._verify_required_credentials_exist() + if premium_user is not True: + raise ValueError( + f"Hashicorp secret manager is only available for premium users. {CommonProxyErrors.not_premium_user.value}" + ) + litellm.secret_manager_client = self litellm._key_management_system = KeyManagementSystem.HASHICORP_VAULT _refresh_interval = os.environ.get( @@ -58,11 +63,6 @@ class HashicorpSecretManager(BaseSecretManager): default_ttl=_refresh_interval ) # store in memory for 1 day - if premium_user is not True: - raise ValueError( - f"Hashicorp secret manager is only available for premium users. {CommonProxyErrors.not_premium_user.value}" - ) - def _verify_required_credentials_exist(self) -> None: """ Validate that at least one authentication method is configured. @@ -70,13 +70,16 @@ class HashicorpSecretManager(BaseSecretManager): Raises: ValueError: If no valid authentication credentials are provided """ - if not self.vault_token and not ( - self.approle_role_id and self.approle_secret_id - ): + has_token = bool(self.vault_token) + has_approle = bool(self.approle_role_id and self.approle_secret_id) + has_tls_cert = bool(self.tls_cert_path and self.tls_key_path) + + if not has_token and not has_approle and not has_tls_cert: raise ValueError( "Missing Vault authentication credentials. Please set either:\n" " - HCP_VAULT_TOKEN for token-based auth, or\n" - " - HCP_VAULT_APPROLE_ROLE_ID and HCP_VAULT_APPROLE_SECRET_ID for AppRole auth" + " - HCP_VAULT_APPROLE_ROLE_ID and HCP_VAULT_APPROLE_SECRET_ID for AppRole auth, or\n" + " - HCP_VAULT_CLIENT_CERT and HCP_VAULT_CLIENT_KEY for TLS certificate auth" ) def _auth_via_approle(self) -> str: diff --git a/litellm/types/proxy/management_endpoints/config_overrides.py b/litellm/types/proxy/management_endpoints/config_overrides.py new file mode 100644 index 00000000000..6f5d661f57a --- /dev/null +++ b/litellm/types/proxy/management_endpoints/config_overrides.py @@ -0,0 +1,64 @@ +from typing import Any, Dict, Optional + +from pydantic import BaseModel, Field + + +class HashicorpVaultConfig(BaseModel): + """Configuration for Hashicorp Vault secret manager integration.""" + + vault_addr: Optional[str] = Field( + default=None, + description="The address of the Vault server (e.g., https://vault.example.com:8200)", + ) + vault_token: Optional[str] = Field( + default=None, + description="Token for Vault token-based authentication", + ) + approle_role_id: Optional[str] = Field( + default=None, + description="Role ID for Vault AppRole authentication", + ) + approle_secret_id: Optional[str] = Field( + default=None, + description="Secret ID for Vault AppRole authentication", + ) + approle_mount_path: Optional[str] = Field( + default=None, + description="Mount path for the AppRole auth method (default: approle)", + ) + client_cert: Optional[str] = Field( + default=None, + description="Path to the client TLS certificate for Vault", + ) + client_key: Optional[str] = Field( + default=None, + description="Path to the client TLS private key for Vault", + ) + vault_cert_role: Optional[str] = Field( + default=None, + description="Certificate role name for TLS cert authentication", + ) + vault_namespace: Optional[str] = Field( + default=None, + description="Vault namespace (for multi-tenant Vault, sent as X-Vault-Namespace header)", + ) + vault_mount_name: Optional[str] = Field( + default=None, + description="KV engine mount name (default: secret)", + ) + vault_path_prefix: Optional[str] = Field( + default=None, + description="Optional path prefix for secrets (e.g., myapp -> secret/data/myapp/{secret_name})", + ) + + +class ConfigOverrideSettingsResponse(BaseModel): + """Response model for config override settings GET endpoints.""" + + config_type: str = Field(description="The type of config override") + values: Dict[str, Any] = Field( + description="Current configuration values (sensitive fields decrypted)" + ) + field_schema: Dict[str, Any] = Field( + description="Schema information for UI rendering" + ) diff --git a/tests/test_litellm/proxy/management_endpoints/test_config_override_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_config_override_endpoints.py new file mode 100644 index 00000000000..22258dc80c6 --- /dev/null +++ b/tests/test_litellm/proxy/management_endpoints/test_config_override_endpoints.py @@ -0,0 +1,251 @@ +import json +import os +from unittest.mock import AsyncMock, MagicMock + +import pytest +from fastapi.testclient import TestClient +from prisma.errors import RecordNotFoundError + +import litellm +import litellm.proxy.proxy_server as ps +from litellm.proxy._types import KeyManagementSystem, LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.management_endpoints.config_override_endpoints import ( + HASHICORP_ENV_VAR_MAPPING, + _build_field_schema, + _set_env_vars, +) +from litellm.proxy.proxy_server import app +from litellm.types.proxy.management_endpoints.config_overrides import ( + HashicorpVaultConfig, +) + +VAULT_URL = "/config_overrides/hashicorp_vault" + + +@pytest.fixture +def client(): + return TestClient(app) + + +def _make_mock_db(): + mock = MagicMock() + mock.find_unique = AsyncMock(return_value=None) + mock.upsert = AsyncMock(return_value=None) + mock.delete = AsyncMock(return_value=None) + prisma = MagicMock() + prisma.db.litellm_configoverrides = mock + return prisma, mock + + +def _make_mock_proxy_config(): + cfg = MagicMock() + cfg.initialize_secret_manager = MagicMock() + cfg._last_hashicorp_vault_config = None + cfg._encrypt_env_variables = MagicMock( + side_effect=lambda d: {k: f"enc_{v}" for k, v in d.items()} + ) + cfg._decrypt_db_variables = MagicMock( + side_effect=lambda d: { + k: v.replace("enc_", "") if isinstance(v, str) else v + for k, v in d.items() + } + ) + return cfg + + +def _upserted_data(mock_db): + return json.loads(mock_db.upsert.call_args.kwargs["data"]["create"]["config_value"]) + + +def _db_record(data): + rec = MagicMock() + rec.config_value = json.dumps(data) + return rec + + +def _cleanup(): + app.dependency_overrides.pop(ps.user_api_key_auth, None) + for env_var in HASHICORP_ENV_VAR_MAPPING.values(): + os.environ.pop(env_var, None) + + +def _set_admin(): + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin_user" + ) + + +@pytest.mark.asyncio +async def test_hashicorp_vault_crud_lifecycle(client, monkeypatch): + """Create → read (masked) → partial update (merge from DB) → clear field → + only-provided fields → delete → idempotent delete → env fallback → + merge from env → helpers → encrypt/decrypt roundtrip.""" + mock_prisma, mock_db = _make_mock_db() + mock_cfg = _make_mock_proxy_config() + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + monkeypatch.setattr(ps, "proxy_config", mock_cfg) + old_client, old_kms = litellm.secret_manager_client, litellm._key_management_system + _set_admin() + + try: + # 1. POST: create + r = client.post(VAULT_URL, json={ + "vault_addr": "https://vault.example.com", + "vault_token": "my-secret-vault-token", + "vault_namespace": "admin", + "vault_mount_name": "secret", + }) + assert r.status_code == 200 + assert os.environ["HCP_VAULT_ADDR"] == "https://vault.example.com" + data = _upserted_data(mock_db) + assert data["vault_token"] == "enc_my-secret-vault-token" + mock_cfg.initialize_secret_manager.assert_called_with(key_management_system="hashicorp_vault") + assert mock_cfg._last_hashicorp_vault_config is not None + + # 2. GET: sensitive fields masked + mock_db.find_unique = AsyncMock(return_value=_db_record(data)) + r = client.get(VAULT_URL) + assert r.status_code == 200 + vals = r.json()["values"] + assert vals["vault_addr"] == "https://vault.example.com" + assert "*" in vals["vault_token"] + assert "properties" in r.json()["field_schema"] + + # 3. POST partial: omitted fields merge from DB + r = client.post(VAULT_URL, json={"vault_addr": "https://vault.new.com"}) + assert r.status_code == 200 + data = _upserted_data(mock_db) + assert data["vault_addr"] == "enc_https://vault.new.com" + assert data["vault_token"] == "enc_my-secret-vault-token" + assert data["vault_namespace"] == "enc_admin" + + # 4. POST empty string: clears field, preserves others + step3 = {**data, "approle_role_id": "enc_role", "approle_secret_id": "enc_secret"} + mock_db.find_unique = AsyncMock(return_value=_db_record(step3)) + mock_db.upsert = AsyncMock(return_value=None) + r = client.post(VAULT_URL, json={"vault_token": ""}) + assert r.status_code == 200 + data = _upserted_data(mock_db) + assert "vault_token" not in data + assert data["approle_role_id"] == "enc_role" + + # 5. POST only provided fields (clean slate) + for v in HASHICORP_ENV_VAR_MAPPING.values(): + os.environ.pop(v, None) + mock_db.find_unique = AsyncMock(return_value=None) + mock_db.upsert = AsyncMock(return_value=None) + r = client.post(VAULT_URL, json={"vault_addr": "https://v.com", "vault_token": "tok"}) + assert r.status_code == 200 + assert _upserted_data(mock_db) == {"vault_addr": "enc_https://v.com", "vault_token": "enc_tok"} + + # 6. DELETE: clears everything + litellm.secret_manager_client = MagicMock() + litellm._key_management_system = KeyManagementSystem.HASHICORP_VAULT + r = client.delete(VAULT_URL) + assert r.status_code == 200 + assert os.environ.get("HCP_VAULT_ADDR") is None + assert litellm.secret_manager_client is None + + # 7. DELETE idempotent + mock_db.delete = AsyncMock( + side_effect=RecordNotFoundError(data={"clientVersion": "0.0.0"}, message="Not found") + ) + assert client.delete(VAULT_URL).status_code == 200 + + # 8. GET: env var fallback + mock_db.find_unique = AsyncMock(return_value=None) + monkeypatch.setenv("HCP_VAULT_ADDR", "https://vault.env.com") + monkeypatch.setenv("HCP_VAULT_NAMESPACE", "env-ns") + r = client.get(VAULT_URL) + assert r.json()["values"]["vault_addr"] == "https://vault.env.com" + + # 9. POST: merge from env vars + monkeypatch.setenv("HCP_VAULT_TOKEN", "env-token") + monkeypatch.setenv("HCP_VAULT_MOUNT_NAME", "env-mount") + mock_cfg.initialize_secret_manager = MagicMock() + mock_db.upsert = AsyncMock(return_value=None) + r = client.post(VAULT_URL, json={"vault_addr": "https://vault.merged.com"}) + assert r.status_code == 200 + data = _upserted_data(mock_db) + assert data["vault_token"] == "enc_env-token" + assert data["vault_mount_name"] == "enc_env-mount" + + # 10. _set_env_vars: empty string unsets + monkeypatch.setenv("HCP_VAULT_TOKEN", "existing") + _set_env_vars({"vault_token": "", "vault_addr": "https://v.com"}) + assert os.environ.get("HCP_VAULT_TOKEN") is None + assert os.environ["HCP_VAULT_ADDR"] == "https://v.com" + + # 11. _build_field_schema + schema = _build_field_schema(HashicorpVaultConfig) + assert "vault_addr" in schema["properties"] + assert len(schema["properties"]["vault_addr"]["description"]) > 0 + + # 12. encrypt/decrypt roundtrip + from litellm.proxy.proxy_server import ProxyConfig + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-test-salt-key") + pc = ProxyConfig() + orig = {"vault_addr": "https://v.com", "vault_token": "secret"} + encrypted = pc._encrypt_env_variables(orig) + assert all(encrypted[k] != orig[k] for k in orig) + decrypted = pc._decrypt_db_variables(encrypted) + assert all(decrypted[k] == orig[k] for k in orig) + + finally: + litellm.secret_manager_client = old_client + litellm._key_management_system = old_kms + _cleanup() + + +@pytest.mark.asyncio +async def test_hashicorp_vault_validation_errors_and_access_control(client, monkeypatch): + """Validation (missing fields, init failure rollback), DELETE preserves + non-Vault secret managers, non-admin 403 on all endpoints.""" + mock_prisma, mock_db = _make_mock_db() + mock_cfg = MagicMock() + mock_cfg._last_hashicorp_vault_config = {"vault_addr": "old"} + monkeypatch.setattr(ps, "prisma_client", mock_prisma) + monkeypatch.setattr(ps, "proxy_config", mock_cfg) + old_client, old_kms = litellm.secret_manager_client, litellm._key_management_system + _set_admin() + + try: + # 1. Missing vault_addr → 400 + r = client.post(VAULT_URL, json={"vault_token": "tok"}) + assert r.status_code == 400 + assert "Vault Address" in r.json()["detail"] + + # 2. Missing auth → 400 + r = client.post(VAULT_URL, json={"vault_addr": "https://v.com"}) + assert r.status_code == 400 + assert "authentication" in r.json()["detail"].lower() + + # 3. Init failure → 500, env vars restored + mock_cfg.initialize_secret_manager = MagicMock(side_effect=Exception("fail")) + monkeypatch.setenv("HCP_VAULT_ADDR", "https://vault.old.com") + monkeypatch.setenv("HCP_VAULT_TOKEN", "old-token") + r = client.post(VAULT_URL, json={"vault_addr": "https://bad.com", "vault_token": "bad"}) + assert r.status_code == 500 + assert os.environ["HCP_VAULT_ADDR"] == "https://vault.old.com" + mock_db.upsert.assert_not_awaited() + + # 4. DELETE preserves non-Vault secret manager + aws = MagicMock() + litellm.secret_manager_client = aws + litellm._key_management_system = KeyManagementSystem.AWS_SECRET_MANAGER + assert client.delete(VAULT_URL).status_code == 200 + assert litellm.secret_manager_client is aws + assert litellm._key_management_system == KeyManagementSystem.AWS_SECRET_MANAGER + + # 5. Non-admin → 403 + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, user_id="user" + ) + assert client.get(VAULT_URL).status_code == 403 + assert client.post(VAULT_URL, json={"vault_addr": "https://v.com"}).status_code == 403 + assert client.delete(VAULT_URL).status_code == 403 + + finally: + litellm.secret_manager_client = old_client + litellm._key_management_system = old_kms + _cleanup() From 53a1e31729b105cb61decd359031319ae0205c10 Mon Sep 17 00:00:00 2001 From: Krish Dholakia Date: Thu, 5 Mar 2026 16:58:46 -0800 Subject: [PATCH 061/273] feat(spend-logs): add truncation note when error logs are truncated for DB storage (#22936) When the messages or response JSON fields in spend logs are truncated before being written to the database, the truncation marker now includes a note explaining: - This is a DB storage safeguard - Full, untruncated data is still sent to logging callbacks (OTEL, Datadog, etc.) - The MAX_STRING_LENGTH_PROMPT_IN_DB env var can be used to increase the limit Also emits a verbose_proxy_logger.info message when truncation occurs in the request body or response spend log paths. Adds 3 new tests: - test_truncation_includes_db_safeguard_note - test_response_truncation_logs_info_message - test_request_body_truncation_logs_info_message Co-authored-by: Cursor Agent --- litellm/constants.py | 5 ++ .../spend_tracking/spend_tracking_utils.py | 28 +++++- .../test_spend_tracking_utils.py | 86 +++++++++++++++++-- 3 files changed, 110 insertions(+), 9 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index c1bb7da1b73..2ae365300ef 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1242,6 +1242,11 @@ X_LITELLM_DISABLE_CALLBACKS = "x-litellm-disable-callbacks" LITELLM_METADATA_FIELD = "litellm_metadata" OLD_LITELLM_METADATA_FIELD = "metadata" LITELLM_TRUNCATED_PAYLOAD_FIELD = "litellm_truncated" +LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE = ( + "Truncation is a DB storage safeguard. " + "Full, untruncated data is logged to logging callbacks (OTEL, Datadog, etc.). " + "To increase the truncation limit, set `MAX_STRING_LENGTH_PROMPT_IN_DB` in your env." +) ########################### LiteLLM Proxy Specific Constants ########################### ######################################################################################## diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 131841f7b59..f381432a089 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -11,6 +11,10 @@ from pydantic import BaseModel import litellm from litellm._logging import verbose_proxy_logger +from litellm.constants import ( + LITELLM_TRUNCATED_PAYLOAD_FIELD, + LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE, +) from litellm.constants import \ MAX_STRING_LENGTH_PROMPT_IN_DB as DEFAULT_MAX_STRING_LENGTH_PROMPT_IN_DB from litellm.constants import REDACTED_BY_LITELM_STRING @@ -628,7 +632,10 @@ def _sanitize_request_body_for_spend_logs_payload( Recursively sanitize request body to prevent logging large base64 strings or other large values. Truncates strings longer than MAX_STRING_LENGTH_PROMPT_IN_DB characters and handles nested dictionaries. """ - from litellm.constants import LITELLM_TRUNCATED_PAYLOAD_FIELD + from litellm.constants import ( + LITELLM_TRUNCATED_PAYLOAD_FIELD, + LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE, + ) if visited is None: visited = set() @@ -674,7 +681,8 @@ def _sanitize_request_body_for_spend_logs_payload( # Build the truncated string: beginning + truncation marker + end truncated_value = ( f"{value[:start_chars]}" - f"... ({LITELLM_TRUNCATED_PAYLOAD_FIELD} skipped {skipped_chars} chars) ..." + f"... ({LITELLM_TRUNCATED_PAYLOAD_FIELD} skipped {skipped_chars} chars. " + f"{LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE}) ..." f"{value[-end_chars:]}" ) return truncated_value @@ -791,6 +799,11 @@ def _get_proxy_server_request_for_spend_logs_payload( _request_body = _sanitize_request_body_for_spend_logs_payload(_request_body) _request_body_json_str = json.dumps(_request_body, default=str) + if LITELLM_TRUNCATED_PAYLOAD_FIELD in _request_body_json_str: + verbose_proxy_logger.info( + "Spend Log: request body was truncated before storing in DB. %s", + LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE, + ) return _request_body_json_str return "{}" @@ -866,8 +879,15 @@ def _get_response_for_spend_logs_payload( if sanitized_response is None: return "{}" if isinstance(sanitized_response, str): - return sanitized_response - return safe_dumps(sanitized_response) + result_str = sanitized_response + else: + result_str = safe_dumps(sanitized_response) + if LITELLM_TRUNCATED_PAYLOAD_FIELD in result_str: + verbose_proxy_logger.info( + "Spend Log: response was truncated before storing in DB. %s", + LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE, + ) + return result_str return "{}" diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index 24f45cc5c91..9a64e641b5e 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -16,7 +16,11 @@ sys.path.insert( from unittest.mock import AsyncMock, MagicMock, patch import litellm -from litellm.constants import LITELLM_TRUNCATED_PAYLOAD_FIELD, REDACTED_BY_LITELM_STRING +from litellm.constants import ( + LITELLM_TRUNCATED_PAYLOAD_FIELD, + LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE, + REDACTED_BY_LITELM_STRING, +) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy.spend_tracking.spend_tracking_utils import ( _get_messages_for_spend_logs_payload, @@ -60,7 +64,7 @@ def test_sanitize_request_body_for_spend_logs_payload_long_string(): end_chars = MAX_STRING_LENGTH_PROMPT_IN_DB - start_chars skipped_chars = len(long_string) - (start_chars + end_chars) - expected_truncation_message = f"... ({LITELLM_TRUNCATED_PAYLOAD_FIELD} skipped {skipped_chars} chars) ..." + expected_truncation_message = f"... ({LITELLM_TRUNCATED_PAYLOAD_FIELD} skipped {skipped_chars} chars. {LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE}) ..." expected_length = start_chars + len(expected_truncation_message) + end_chars assert len(sanitized["text"]) == expected_length @@ -86,7 +90,7 @@ def test_sanitize_request_body_for_spend_logs_payload_nested_dict(): end_chars = MAX_STRING_LENGTH_PROMPT_IN_DB - start_chars skipped_chars = len(long_string) - total_keep - expected_truncation_message = f"... ({LITELLM_TRUNCATED_PAYLOAD_FIELD} skipped {skipped_chars} chars) ..." + expected_truncation_message = f"... ({LITELLM_TRUNCATED_PAYLOAD_FIELD} skipped {skipped_chars} chars. {LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE}) ..." expected_length = start_chars + len(expected_truncation_message) + end_chars assert len(sanitized["outer"]["inner"]["text"]) == expected_length @@ -111,7 +115,7 @@ def test_sanitize_request_body_for_spend_logs_payload_nested_list(): end_chars = MAX_STRING_LENGTH_PROMPT_IN_DB - start_chars skipped_chars = len(long_string) - total_keep - expected_truncation_message = f"... ({LITELLM_TRUNCATED_PAYLOAD_FIELD} skipped {skipped_chars} chars) ..." + expected_truncation_message = f"... ({LITELLM_TRUNCATED_PAYLOAD_FIELD} skipped {skipped_chars} chars. {LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE}) ..." expected_length = start_chars + len(expected_truncation_message) + end_chars assert len(sanitized["items"][0]["text"]) == expected_length @@ -151,7 +155,7 @@ def test_sanitize_request_body_for_spend_logs_payload_mixed_types(): end_chars = MAX_STRING_LENGTH_PROMPT_IN_DB - start_chars skipped_chars = len(long_string) - total_keep - expected_truncation_message = f"... ({LITELLM_TRUNCATED_PAYLOAD_FIELD} skipped {skipped_chars} chars) ..." + expected_truncation_message = f"... ({LITELLM_TRUNCATED_PAYLOAD_FIELD} skipped {skipped_chars} chars. {LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE}) ..." expected_length = start_chars + len(expected_truncation_message) + end_chars assert len(sanitized["text"]) == expected_length @@ -396,6 +400,78 @@ def test_get_response_for_spend_logs_payload_truncates_large_embedding(mock_shou assert parsed["data"][0]["other_field"] == "value" +def test_truncation_includes_db_safeguard_note(): + """ + Test that truncated content includes the DB safeguard note explaining + that full data is available in OTEL/other logging integrations. + """ + from litellm.constants import MAX_STRING_LENGTH_PROMPT_IN_DB + + large_error = "Error: " + "x" * (MAX_STRING_LENGTH_PROMPT_IN_DB + 1000) + request_body = {"error_trace": large_error} + sanitized = _sanitize_request_body_for_spend_logs_payload(request_body) + + truncated = sanitized["error_trace"] + assert LITELLM_TRUNCATED_PAYLOAD_FIELD in truncated + assert LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE in truncated + assert "DB storage safeguard" in truncated + assert "logging callbacks" in truncated.lower() or "logging integrations" in truncated.lower() or "logging callbacks" in truncated + + +@patch( + "litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs" +) +def test_response_truncation_logs_info_message(mock_should_store): + """ + Test that when response is truncated before DB storage, an info log is emitted + noting that full data is available in OTEL/other integrations. + """ + from litellm.constants import MAX_STRING_LENGTH_PROMPT_IN_DB + + mock_should_store.return_value = True + large_text = "B" * (MAX_STRING_LENGTH_PROMPT_IN_DB + 500) + payload = cast( + StandardLoggingPayload, + {"response": {"data": [{"content": large_text}]}}, + ) + + with patch( + "litellm.proxy.spend_tracking.spend_tracking_utils.verbose_proxy_logger" + ) as mock_logger: + _get_response_for_spend_logs_payload(payload) + mock_logger.info.assert_called_once() + log_msg = mock_logger.info.call_args[0][0] + assert "response was truncated" in log_msg + + +@patch( + "litellm.proxy.spend_tracking.spend_tracking_utils._should_store_prompts_and_responses_in_spend_logs" +) +def test_request_body_truncation_logs_info_message(mock_should_store): + """ + Test that when request body is truncated before DB storage, an info log is emitted. + """ + from litellm.constants import MAX_STRING_LENGTH_PROMPT_IN_DB + + mock_should_store.return_value = True + large_prompt = "C" * (MAX_STRING_LENGTH_PROMPT_IN_DB + 500) + litellm_params = { + "proxy_server_request": { + "body": {"messages": [{"role": "user", "content": large_prompt}]} + } + } + + with patch( + "litellm.proxy.spend_tracking.spend_tracking_utils.verbose_proxy_logger" + ) as mock_logger: + _get_proxy_server_request_for_spend_logs_payload( + metadata={}, litellm_params=litellm_params, kwargs={} + ) + mock_logger.info.assert_called_once() + log_msg = mock_logger.info.call_args[0][0] + assert "request body was truncated" in log_msg + + def test_safe_dumps_handles_circular_references(): """Test that safe_dumps can handle circular references without raising exceptions""" From d0e480414ce23c2c278d1c7f5885afc6d1dd1e4e Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Thu, 5 Mar 2026 17:00:51 -0800 Subject: [PATCH 062/273] Fix team usage spend showing lower than expected values The /team/daily/activity endpoint used Prisma pagination (page_size=1000) but the UI only fetched page 1. Teams with many keys/models easily exceed 1000 rows in LiteLLM_DailyTeamSpend, causing truncated totals. Switches the endpoint to use SQL GROUP BY via get_daily_activity_aggregated with include_entity_breakdown=True, returning all data in a single response while preserving per-team breakdown. Also adds timezone parameter support. Co-Authored-By: Claude Opus 4.6 --- .../common_daily_activity.py | 44 +++++-- .../management_endpoints/team_endpoints.py | 23 ++-- .../test_team_endpoints.py | 110 +++++++++++++----- 3 files changed, 126 insertions(+), 51 deletions(-) diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index 02961748e7c..a4fbeb7e28f 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -474,16 +474,21 @@ def _build_aggregated_sql_query( start_date: str, end_date: str, model: Optional[str], - api_key: Optional[str], + api_key: Optional[Union[str, List[str]]], exclude_entity_ids: Optional[List[str]] = None, timezone_offset_minutes: Optional[int] = None, + include_entity_id: bool = False, ) -> Tuple[str, List[Any]]: """Build a parameterized SQL GROUP BY query for aggregated daily activity. Groups by (date, api_key, model, model_group, custom_llm_provider, mcp_namespaced_tool_name, endpoint) with SUMs on all metric columns. - The entity_id column is intentionally omitted from GROUP BY to collapse - rows across entities — this is where the biggest row reduction comes from. + + When include_entity_id is False (default), the entity_id column is omitted + from GROUP BY to collapse rows across entities. + + When include_entity_id is True, the entity_id column is included in both + SELECT and GROUP BY, preserving per-entity breakdown in the results. Returns: Tuple of (sql_query, params_list) ready for prisma_client.db.query_raw(). @@ -538,14 +543,24 @@ def _build_aggregated_sql_query( # Optional api_key filter if api_key: - sql_conditions.append(f"api_key = ${p}") - sql_params.append(api_key) - p += 1 + if isinstance(api_key, list): + placeholders = ", ".join(f"${p + i}" for i in range(len(api_key))) + sql_conditions.append(f"api_key IN ({placeholders})") + sql_params.extend(api_key) + p += len(api_key) + else: + sql_conditions.append(f"api_key = ${p}") + sql_params.append(api_key) + p += 1 where_clause = " AND ".join(sql_conditions) + entity_select = f'"{entity_id_field}",' if include_entity_id else "" + entity_group_by = f'"{entity_id_field}",' if include_entity_id else "" + sql_query = f""" SELECT + {entity_select} date, api_key, model, @@ -563,7 +578,7 @@ def _build_aggregated_sql_query( SUM(failed_requests)::bigint AS failed_requests FROM "{pg_table}" WHERE {where_clause} - GROUP BY date, api_key, model, model_group, custom_llm_provider, + GROUP BY {entity_group_by} date, api_key, model, model_group, custom_llm_provider, mcp_namespaced_tool_name, endpoint ORDER BY date DESC """ @@ -735,9 +750,10 @@ async def get_daily_activity_aggregated( start_date: Optional[str], end_date: Optional[str], model: Optional[str], - api_key: Optional[str], + api_key: Optional[Union[str, List[str]]], exclude_entity_ids: Optional[List[str]] = None, timezone_offset_minutes: Optional[int] = None, + include_entity_breakdown: bool = False, ) -> SpendAnalyticsPaginatedResponse: """Aggregated variant that returns the full result set (no pagination). @@ -745,6 +761,11 @@ async def get_daily_activity_aggregated( all individual rows into Python. This collapses rows across entities (users/teams/orgs), reducing ~150k rows to ~2-3k grouped rows. + When include_entity_breakdown is True, the entity_id column is included + in the GROUP BY so that per-entity breakdown data is preserved in the + response (e.g. per-team spend). This is needed for entity-specific views + like the team usage dashboard. + Matches the response model of the paginated endpoint so the UI does not need to transform. """ if prisma_client is None: @@ -770,6 +791,7 @@ async def get_daily_activity_aggregated( api_key=api_key, exclude_entity_ids=exclude_entity_ids, timezone_offset_minutes=timezone_offset_minutes, + include_entity_id=include_entity_breakdown, ) # Execute GROUP BY query — returns pre-aggregated dicts @@ -780,13 +802,11 @@ async def get_daily_activity_aggregated( # Convert dicts to objects for compatibility with _aggregate_spend_records records = [SimpleNamespace(**row) for row in rows] - # entity_id_field=None skips entity breakdown (entity dimension was - # collapsed by the GROUP BY, so per-entity data is not available) aggregated = await _aggregate_spend_records( prisma_client=prisma_client, records=records, - entity_id_field=None, - entity_metadata_field=None, + entity_id_field=entity_id_field if include_entity_breakdown else None, + entity_metadata_field=entity_metadata_field if include_entity_breakdown else None, ) return SpendAnalyticsPaginatedResponse( diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 80d50f31a17..5e7a0931b2a 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -77,8 +77,8 @@ from litellm.proxy.management_endpoints.common_utils import ( _upsert_budget_and_membership, _user_has_admin_view, ) -from litellm.proxy.management_endpoints.tag_management_endpoints import ( - get_daily_activity, +from litellm.proxy.management_endpoints.common_daily_activity import ( + get_daily_activity_aggregated, ) from litellm.proxy.management_helpers.object_permission_utils import ( _set_object_permission, @@ -3890,22 +3890,27 @@ async def get_team_daily_activity( page: int = 1, page_size: int = 10, exclude_team_ids: Optional[str] = None, + timezone: Optional[int] = None, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ Get daily activity for specific teams or all teams. + Uses SQL GROUP BY to aggregate all matching rows without pagination, + ensuring accurate total spend regardless of data volume. + Args: team_ids (Optional[str]): Comma-separated list of team IDs to filter by. If not provided, returns data for all teams. start_date (Optional[str]): Start date for the activity period (YYYY-MM-DD). end_date (Optional[str]): End date for the activity period (YYYY-MM-DD). model (Optional[str]): Filter by model name. api_key (Optional[str]): Filter by API key. - page (int): Page number for pagination. - page_size (int): Number of items per page. + page (int): Deprecated, kept for backward compatibility. All results are returned in a single page. + page_size (int): Deprecated, kept for backward compatibility. exclude_team_ids (Optional[str]): Comma-separated list of team IDs to exclude. + timezone (Optional[int]): Timezone offset in minutes from UTC (e.g., 480 for PST). Returns: - SpendAnalyticsPaginatedResponse: Paginated response containing daily activity data. + SpendAnalyticsPaginatedResponse: Response containing daily activity data with per-team breakdown. """ from litellm.proxy.proxy_server import ( prisma_client, @@ -4009,17 +4014,17 @@ async def get_team_daily_activity( if final_api_key_filter is None and user_api_keys is not None: final_api_key_filter = user_api_keys - return await get_daily_activity( + return await get_daily_activity_aggregated( prisma_client=prisma_client, table_name="litellm_dailyteamspend", entity_id_field="team_id", entity_id=team_ids_list, entity_metadata_field=team_alias_metadata, - exclude_entity_ids=exclude_team_ids_list, start_date=start_date, end_date=end_date, model=model, api_key=final_api_key_filter, - page=page, - page_size=page_size, + exclude_entity_ids=exclude_team_ids_list, + timezone_offset_minutes=timezone, + include_entity_breakdown=True, ) diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index b6ac974e2cf..0a2a7e0c432 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -5379,10 +5379,10 @@ async def test_get_team_daily_activity_non_admin_filters_by_user_api_keys( # Mock get_daily_activity to capture the api_key parameter with patch( - "litellm.proxy.management_endpoints.team_endpoints.get_daily_activity", + "litellm.proxy.management_endpoints.team_endpoints.get_daily_activity_aggregated", new_callable=AsyncMock, - ) as mock_get_daily_activity: - mock_get_daily_activity.return_value = MagicMock() + ) as mock_get_daily_activity_agg: + mock_get_daily_activity_agg.return_value = MagicMock() # Call the endpoint await get_team_daily_activity( @@ -5398,8 +5398,8 @@ async def test_get_team_daily_activity_non_admin_filters_by_user_api_keys( ) # Verify get_daily_activity was called with user's API keys as filter - mock_get_daily_activity.assert_called_once() - call_kwargs = mock_get_daily_activity.call_args[1] + mock_get_daily_activity_agg.assert_called_once() + call_kwargs = mock_get_daily_activity_agg.call_args[1] assert call_kwargs["api_key"] == ["user_key_1", "user_key_2"] assert call_kwargs["entity_id"] == [team_id] @@ -5464,10 +5464,10 @@ async def test_get_team_daily_activity_team_admin_sees_all_spend(mock_db_client) # Mock get_daily_activity to capture the api_key parameter with patch( - "litellm.proxy.management_endpoints.team_endpoints.get_daily_activity", + "litellm.proxy.management_endpoints.team_endpoints.get_daily_activity_aggregated", new_callable=AsyncMock, - ) as mock_get_daily_activity: - mock_get_daily_activity.return_value = MagicMock() + ) as mock_get_daily_activity_agg: + mock_get_daily_activity_agg.return_value = MagicMock() # Call the endpoint await get_team_daily_activity( @@ -5483,8 +5483,8 @@ async def test_get_team_daily_activity_team_admin_sees_all_spend(mock_db_client) ) # Verify get_daily_activity was called WITHOUT API key filtering - mock_get_daily_activity.assert_called_once() - call_kwargs = mock_get_daily_activity.call_args[1] + mock_get_daily_activity_agg.assert_called_once() + call_kwargs = mock_get_daily_activity_agg.call_args[1] assert call_kwargs["api_key"] is None assert call_kwargs["entity_id"] == [team_id] @@ -5553,10 +5553,10 @@ async def test_get_team_daily_activity_member_with_permission_sees_all_spend( # Mock get_daily_activity to capture the api_key parameter with patch( - "litellm.proxy.management_endpoints.team_endpoints.get_daily_activity", + "litellm.proxy.management_endpoints.team_endpoints.get_daily_activity_aggregated", new_callable=AsyncMock, - ) as mock_get_daily_activity: - mock_get_daily_activity.return_value = MagicMock() + ) as mock_get_daily_activity_agg: + mock_get_daily_activity_agg.return_value = MagicMock() # Call the endpoint await get_team_daily_activity( @@ -5572,8 +5572,8 @@ async def test_get_team_daily_activity_member_with_permission_sees_all_spend( ) # Verify get_daily_activity was called WITHOUT API key filtering - mock_get_daily_activity.assert_called_once() - call_kwargs = mock_get_daily_activity.call_args[1] + mock_get_daily_activity_agg.assert_called_once() + call_kwargs = mock_get_daily_activity_agg.call_args[1] assert call_kwargs["api_key"] is None assert call_kwargs["entity_id"] == [team_id] @@ -5652,10 +5652,10 @@ async def test_get_team_daily_activity_member_without_permission_filters_by_keys # Mock get_daily_activity to capture the api_key parameter with patch( - "litellm.proxy.management_endpoints.team_endpoints.get_daily_activity", + "litellm.proxy.management_endpoints.team_endpoints.get_daily_activity_aggregated", new_callable=AsyncMock, - ) as mock_get_daily_activity: - mock_get_daily_activity.return_value = MagicMock() + ) as mock_get_daily_activity_agg: + mock_get_daily_activity_agg.return_value = MagicMock() # Call the endpoint await get_team_daily_activity( @@ -5671,8 +5671,8 @@ async def test_get_team_daily_activity_member_without_permission_filters_by_keys ) # Verify get_daily_activity was called WITH API key filtering - mock_get_daily_activity.assert_called_once() - call_kwargs = mock_get_daily_activity.call_args[1] + mock_get_daily_activity_agg.assert_called_once() + call_kwargs = mock_get_daily_activity_agg.call_args[1] assert call_kwargs["api_key"] == ["user_key_abc", "user_key_def"] assert call_kwargs["entity_id"] == [team_id] @@ -5822,10 +5822,10 @@ async def test_get_team_daily_activity_non_admin_filters_by_user_api_keys( # Mock get_daily_activity to capture the api_key parameter with patch( - "litellm.proxy.management_endpoints.team_endpoints.get_daily_activity", + "litellm.proxy.management_endpoints.team_endpoints.get_daily_activity_aggregated", new_callable=AsyncMock, - ) as mock_get_daily_activity: - mock_get_daily_activity.return_value = MagicMock() + ) as mock_get_daily_activity_agg: + mock_get_daily_activity_agg.return_value = MagicMock() # Call the endpoint await get_team_daily_activity( @@ -5841,8 +5841,8 @@ async def test_get_team_daily_activity_non_admin_filters_by_user_api_keys( ) # Verify get_daily_activity was called with user's API keys as filter - mock_get_daily_activity.assert_called_once() - call_kwargs = mock_get_daily_activity.call_args[1] + mock_get_daily_activity_agg.assert_called_once() + call_kwargs = mock_get_daily_activity_agg.call_args[1] assert call_kwargs["api_key"] == ["user_key_1", "user_key_2"] assert call_kwargs["entity_id"] == [team_id] @@ -5907,10 +5907,10 @@ async def test_get_team_daily_activity_team_admin_sees_all_spend(mock_db_client) # Mock get_daily_activity to capture the api_key parameter with patch( - "litellm.proxy.management_endpoints.team_endpoints.get_daily_activity", + "litellm.proxy.management_endpoints.team_endpoints.get_daily_activity_aggregated", new_callable=AsyncMock, - ) as mock_get_daily_activity: - mock_get_daily_activity.return_value = MagicMock() + ) as mock_get_daily_activity_agg: + mock_get_daily_activity_agg.return_value = MagicMock() # Call the endpoint await get_team_daily_activity( @@ -5926,8 +5926,8 @@ async def test_get_team_daily_activity_team_admin_sees_all_spend(mock_db_client) ) # Verify get_daily_activity was called WITHOUT API key filtering - mock_get_daily_activity.assert_called_once() - call_kwargs = mock_get_daily_activity.call_args[1] + mock_get_daily_activity_agg.assert_called_once() + call_kwargs = mock_get_daily_activity_agg.call_args[1] assert call_kwargs["api_key"] is None assert call_kwargs["entity_id"] == [team_id] @@ -5939,6 +5939,56 @@ async def test_get_team_daily_activity_team_admin_sees_all_spend(mock_db_client) assert False, "API keys should not be fetched for team admin users" +@pytest.mark.asyncio +async def test_get_team_daily_activity_uses_aggregated_with_entity_breakdown( + mock_db_client, +): + """ + Test that /team/daily/activity calls get_daily_activity_aggregated + with include_entity_breakdown=True, timezone, and correct parameters. + """ + from litellm.proxy.management_endpoints.team_endpoints import ( + get_team_daily_activity, + ) + + user_api_key_dict = UserAPIKeyAuth( + user_id="admin_user", user_role=LitellmUserRoles.PROXY_ADMIN + ) + + # Mock the team table query for fetching team aliases + mock_db_client.db.litellm_teamtable.find_many = AsyncMock(return_value=[]) + + with patch( + "litellm.proxy.management_endpoints.team_endpoints.get_daily_activity_aggregated", + new_callable=AsyncMock, + ) as mock_get_daily_activity_agg: + mock_get_daily_activity_agg.return_value = MagicMock() + + await get_team_daily_activity( + team_ids="team_1,team_2", + start_date="2024-01-01", + end_date="2024-01-31", + model=None, + api_key=None, + page=1, + page_size=10, + exclude_team_ids="litellm-dashboard", + timezone=480, + user_api_key_dict=user_api_key_dict, + ) + + mock_get_daily_activity_agg.assert_called_once() + call_kwargs = mock_get_daily_activity_agg.call_args[1] + assert call_kwargs["table_name"] == "litellm_dailyteamspend" + assert call_kwargs["entity_id_field"] == "team_id" + assert call_kwargs["entity_id"] == ["team_1", "team_2"] + assert call_kwargs["exclude_entity_ids"] == ["litellm-dashboard"] + assert call_kwargs["start_date"] == "2024-01-01" + assert call_kwargs["end_date"] == "2024-01-31" + assert call_kwargs["timezone_offset_minutes"] == 480 + assert call_kwargs["include_entity_breakdown"] is True + + @pytest.mark.asyncio async def test_validate_and_populate_member_user_info_both_provided_match(): """ From 537be618d4234826fcbb2a654fa2682971f3dc72 Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Thu, 5 Mar 2026 17:17:50 -0800 Subject: [PATCH 063/273] fix(types): add CONFIG_OVERRIDES to SupportedDBObjectType enum Without this, deployments using supported_db_objects filtering would silently skip polling for config_overrides, preventing Hashicorp Vault config from syncing across pods. --- litellm/proxy/_types.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index e408abb3c1b..55d3a61de68 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -77,6 +77,7 @@ class SupportedDBObjectType(str, enum.Enum): PASS_THROUGH_ENDPOINTS = "pass_through_endpoints" PROMPTS = "prompts" MODEL_COST_MAP = "model_cost_map" + CONFIG_OVERRIDES = "config_overrides" def __str__(self): return str(self.value) @@ -2126,7 +2127,7 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): user_header_mappings: Optional[List[UserHeaderMapping]] = None supported_db_objects: Optional[List[SupportedDBObjectType]] = Field( None, - description="Fine-grained control over which object types to load from the database when store_model_in_db is True. Available types: 'models', 'mcp', 'guardrails', 'vector_stores', 'pass_through_endpoints', 'prompts', 'model_cost_map'. If not set, all objects are loaded (default behavior).", + description="Fine-grained control over which object types to load from the database when store_model_in_db is True. Available types: 'models', 'mcp', 'guardrails', 'vector_stores', 'pass_through_endpoints', 'prompts', 'model_cost_map', 'config_overrides'. If not set, all objects are loaded (default behavior).", ) user_mcp_management_mode: Optional[UserMCPManagementMode] = Field( None, From c953388927d44be1877b70f6fb39a16b07e7c11e Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Thu, 5 Mar 2026 17:21:05 -0800 Subject: [PATCH 064/273] fix(vault): remove approle_role_id from sensitive fields, use async HTTP for test_connection - approle_role_id is a non-secret identifier (like a username) per Vault's AppRole model; masking it hinders admin auditing - Use async httpx client for the token lookup-self call to avoid blocking the event loop --- litellm/proxy/_types.py | 3 ++- .../management_endpoints/config_override_endpoints.py | 8 ++++---- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 55d3a61de68..61197738e72 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -77,6 +77,7 @@ class SupportedDBObjectType(str, enum.Enum): PASS_THROUGH_ENDPOINTS = "pass_through_endpoints" PROMPTS = "prompts" MODEL_COST_MAP = "model_cost_map" + TOOLS = "tools" CONFIG_OVERRIDES = "config_overrides" def __str__(self): @@ -2127,7 +2128,7 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): user_header_mappings: Optional[List[UserHeaderMapping]] = None supported_db_objects: Optional[List[SupportedDBObjectType]] = Field( None, - description="Fine-grained control over which object types to load from the database when store_model_in_db is True. Available types: 'models', 'mcp', 'guardrails', 'vector_stores', 'pass_through_endpoints', 'prompts', 'model_cost_map', 'config_overrides'. If not set, all objects are loaded (default behavior).", + description="Fine-grained control over which object types to load from the database when store_model_in_db is True. Available types: 'models', 'mcp', 'guardrails', 'vector_stores', 'pass_through_endpoints', 'prompts', 'model_cost_map', 'tools', 'config_overrides'. If not set, all objects are loaded (default behavior).", ) user_mcp_management_mode: Optional[UserMCPManagementMode] = Field( None, diff --git a/litellm/proxy/management_endpoints/config_override_endpoints.py b/litellm/proxy/management_endpoints/config_override_endpoints.py index 2978a523fb1..78cb91b3483 100644 --- a/litellm/proxy/management_endpoints/config_override_endpoints.py +++ b/litellm/proxy/management_endpoints/config_override_endpoints.py @@ -8,7 +8,8 @@ from pydantic import TypeAdapter import litellm from litellm._logging import verbose_proxy_logger -from litellm.llms.custom_httpx.http_handler import _get_httpx_client +from litellm.llms.custom_httpx.http_handler import get_async_httpx_client +from litellm.llms.custom_httpx.httpx_handler import httpxSpecialProvider from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker from litellm.proxy._types import CommonProxyErrors, KeyManagementSystem, LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth @@ -37,7 +38,6 @@ HASHICORP_ENV_VAR_MAPPING: Dict[str, str] = { HASHICORP_SENSITIVE_FIELDS: Set[str] = { "vault_token", - "approle_role_id", "approle_secret_id", "client_key", } @@ -387,11 +387,11 @@ async def test_hashicorp_vault_connection( # Step 2: Verify the token is valid via token/lookup-self try: - sync_client = _get_httpx_client() + async_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.ProxyServer) lookup_url = f"{client.vault_addr}/v1/auth/token/lookup-self" if client.vault_namespace: headers["X-Vault-Namespace"] = client.vault_namespace - response = sync_client.get(lookup_url, headers=headers) + response = await async_client.get(lookup_url, headers=headers) response.raise_for_status() except Exception as e: raise HTTPException( From 8d539db108dc55cca303e8f2c6757243e7dfaa1e Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Thu, 5 Mar 2026 17:36:46 -0800 Subject: [PATCH 065/273] Fix admin viewer unable to see all organizations The /organization/list endpoint only checked for PROXY_ADMIN role, causing PROXY_ADMIN_VIEW_ONLY users to fall into the else branch which restricts results to orgs the user is a member of. Use the existing _user_has_admin_view() helper to include both roles. --- litellm/proxy/management_endpoints/organization_endpoints.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/management_endpoints/organization_endpoints.py b/litellm/proxy/management_endpoints/organization_endpoints.py index 1c19c4ef313..103b2efcdde 100644 --- a/litellm/proxy/management_endpoints/organization_endpoints.py +++ b/litellm/proxy/management_endpoints/organization_endpoints.py @@ -649,8 +649,8 @@ async def list_organization( "mode": "insensitive", # Case-insensitive search } - # if proxy admin - get all orgs (with optional filters) - if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN: + # if proxy admin or admin viewer - get all orgs (with optional filters) + if _user_has_admin_view(user_api_key_dict): response = await prisma_client.db.litellm_organizationtable.find_many( where=where_conditions if where_conditions else None, include={"litellm_budget_table": True, "members": True, "teams": True}, From 73a8e8cf07535cbd5ab648ed0939219a70543591 Mon Sep 17 00:00:00 2001 From: Ryan Crabbe Date: Thu, 5 Mar 2026 17:40:51 -0800 Subject: [PATCH 066/273] fix(vault): resolve merge conflict, use async auth, include error details - Remove duplicate description kwarg in supported_db_objects Field() that caused SyntaxError preventing proxy startup - Wrap sync _get_request_headers() in asyncio.to_thread to avoid blocking the event loop during AppRole/TLS cert auth - Include exception messages in error responses for admin-only endpoints to aid debugging --- litellm/proxy/_types.py | 1 - .../management_endpoints/config_override_endpoints.py | 9 +++++---- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 12f6cdf600d..da7e1f5a049 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2169,7 +2169,6 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): supported_db_objects: Optional[List[SupportedDBObjectType]] = Field( None, description="Fine-grained control over which object types to load from the database when store_model_in_db is True. Available types: 'models', 'mcp', 'guardrails', 'vector_stores', 'pass_through_endpoints', 'prompts', 'model_cost_map', 'tools', 'config_overrides'. If not set, all objects are loaded (default behavior).", - description="Fine-grained control over which object types to load from the database when store_model_in_db is True. Available types: 'models', 'mcp', 'guardrails', 'vector_stores', 'pass_through_endpoints', 'prompts', 'model_cost_map', 'tools'. If not set, all objects are loaded (default behavior).", ) user_mcp_management_mode: Optional[UserMCPManagementMode] = Field( None, diff --git a/litellm/proxy/management_endpoints/config_override_endpoints.py b/litellm/proxy/management_endpoints/config_override_endpoints.py index 78cb91b3483..f1d6cacf1e7 100644 --- a/litellm/proxy/management_endpoints/config_override_endpoints.py +++ b/litellm/proxy/management_endpoints/config_override_endpoints.py @@ -1,3 +1,4 @@ +import asyncio import json import os from typing import Any, Dict, Set @@ -216,7 +217,7 @@ async def update_hashicorp_vault_config( ) raise HTTPException( status_code=500, - detail="Failed to initialize secret manager", + detail=f"Failed to initialize secret manager: {e}", ) # Only persist to DB after successful init @@ -378,11 +379,11 @@ async def test_hashicorp_vault_connection( # Step 1: Authenticate (exercises AppRole login, TLS cert login, or direct token) try: - headers = client._get_request_headers() + headers = await asyncio.to_thread(client._get_request_headers) except Exception as e: raise HTTPException( status_code=502, - detail="Vault authentication failed", + detail=f"Vault authentication failed: {e}", ) # Step 2: Verify the token is valid via token/lookup-self @@ -396,7 +397,7 @@ async def test_hashicorp_vault_connection( except Exception as e: raise HTTPException( status_code=502, - detail="Vault token validation failed", + detail=f"Vault token validation failed: {e}", ) return { From ec600aa70a06e3c0d92467472f5e75e474b79485 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 5 Mar 2026 18:13:04 -0800 Subject: [PATCH 067/273] =?UTF-8?q?feat(ui):=20add=20Chat=20UI=20=E2=80=94?= =?UTF-8?q?=20ChatGPT-like=20interface=20with=20MCP=20tools=20and=20stream?= =?UTF-8?q?ing=20(#22937)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(ui): add chat message and conversation types * feat(ui): add useChatHistory hook for localStorage-backed conversations * feat(ui): add ConversationList sidebar component * feat(ui): add MCPConnectPicker for attaching MCP servers to chat * feat(ui): add ModelSelector dropdown for chat * feat(ui): add ChatInputBar with MCP tool attachment support * feat(ui): add MCPAppsPanel with list/detail view for MCP servers * feat(ui): add ChatMessages component; remove auto-scrollIntoView that caused scroll-lock bypass * feat(ui): add ChatPage — ChatGPT-like UI with scroll lock, MCP tools, streaming * feat(ui): add /chat route wired to ChatPage * feat(ui): remove chat from leftnav — chat accessible via navbar button * feat(ui): add Chat button to top navbar * feat(ui): add dismissible Chat UI announcement banner to Playground page * feat(proxy): add Chat UI link to Swagger description * feat(ui): add react-markdown and syntax-highlighter deps for chat UI * fix(ui): replace missing BorderOutlined import with inline stop icon div * fix(ui): apply remark-gfm plugin to ReactMarkdown for GFM support * fix(ui): remove unused isEvenRow variable in MCPAppsPanel * fix(ui): add ellipsis when truncating conversation title * fix(ui): wire search button to chats view; remove non-functional keyboard hint * fix(ui): use serverRootPath in navbar chat link for sub-path deployments * fix(ui): remove unused ChatInputBar and ModelSelector files * fix(ui): correct grid bottom-border condition for odd server count * fix(chat): move localStorage writes out of setConversations updater (React purity) * fix(chat): fix stale closure in handleEditAndResend - compute history before async state update * fix(chat): fix 4 issues in ChatMessages - array redaction, clipboard error, inline detection, remove unused ref --- litellm/proxy/proxy_server.py | 9 +- ui/litellm-dashboard/package-lock.json | 295 +++++++ ui/litellm-dashboard/package.json | 4 +- .../src/app/(dashboard)/playground/page.tsx | 64 +- ui/litellm-dashboard/src/app/chat/page.tsx | 19 + .../src/components/chat/ChatMessages.tsx | 577 ++++++++++++ .../src/components/chat/ChatPage.tsx | 823 ++++++++++++++++++ .../src/components/chat/ConversationList.tsx | 483 ++++++++++ .../src/components/chat/MCPAppsPanel.tsx | 274 ++++++ .../src/components/chat/MCPConnectPicker.tsx | 157 ++++ .../src/components/chat/types.ts | 20 + .../src/components/chat/useChatHistory.ts | 230 +++++ .../src/components/leftnav.tsx | 1 + .../src/components/navbar.tsx | 39 +- 14 files changed, 2988 insertions(+), 7 deletions(-) create mode 100644 ui/litellm-dashboard/src/app/chat/page.tsx create mode 100644 ui/litellm-dashboard/src/components/chat/ChatMessages.tsx create mode 100644 ui/litellm-dashboard/src/components/chat/ChatPage.tsx create mode 100644 ui/litellm-dashboard/src/components/chat/ConversationList.tsx create mode 100644 ui/litellm-dashboard/src/components/chat/MCPAppsPanel.tsx create mode 100644 ui/litellm-dashboard/src/components/chat/MCPConnectPicker.tsx create mode 100644 ui/litellm-dashboard/src/components/chat/types.ts create mode 100644 ui/litellm-dashboard/src/components/chat/useChatHistory.ts diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 9683b37dbb4..7fe0ce6d6f5 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -372,6 +372,9 @@ from litellm.proxy.management_endpoints.internal_user_endpoints import ( from litellm.proxy.management_endpoints.internal_user_endpoints import ( user_update, ) +from litellm.proxy.management_endpoints.jwt_key_mapping_endpoints import ( + router as jwt_key_mapping_router, +) from litellm.proxy.management_endpoints.key_management_endpoints import ( delete_verification_tokens, duration_in_seconds, @@ -380,9 +383,6 @@ from litellm.proxy.management_endpoints.key_management_endpoints import ( from litellm.proxy.management_endpoints.key_management_endpoints import ( router as key_management_router, ) -from litellm.proxy.management_endpoints.jwt_key_mapping_endpoints import ( - router as jwt_key_mapping_router, -) from litellm.proxy.management_endpoints.mcp_management_endpoints import ( router as mcp_management_router, ) @@ -661,6 +661,9 @@ ui_message += "\n\n💸 [```LiteLLM Model Cost Map```](https://models.litellm.ai ui_message += f"\n\n🔎 [```LiteLLM Model Hub```]({model_hub_link}). See available models on the proxy. [**Docs**](https://docs.litellm.ai/docs/proxy/ai_hub)" +chat_link = f"{server_root_path}/ui/chat" +ui_message += f"\n\n💬 [```LiteLLM Chat UI```]({chat_link}). ChatGPT-like interface for your users to chat with AI models and MCP tools." + custom_swagger_message = "[**Customize Swagger Docs**](https://docs.litellm.ai/docs/proxy/enterprise#swagger-docs---custom-routes--branding)" ### CUSTOM BRANDING [ENTERPRISE FEATURE] ### diff --git a/ui/litellm-dashboard/package-lock.json b/ui/litellm-dashboard/package-lock.json index 200182cf551..69efbf19c38 100644 --- a/ui/litellm-dashboard/package-lock.json +++ b/ui/litellm-dashboard/package-lock.json @@ -19,6 +19,7 @@ "@types/papaparse": "^5.3.15", "antd": "^5.13.2", "cva": "^1.0.0-beta.3", + "dayjs": "^1.11.19", "jwt-decode": "^4.0.0", "lucide-react": "^0.513.0", "moment": "^2.30.1", @@ -31,6 +32,7 @@ "react-json-view-lite": "^2.5.0", "react-markdown": "^9.0.1", "react-syntax-highlighter": "^15.6.6", + "remark-gfm": "^4.0.1", "tailwind-merge": "^3.2.0", "uuid": "^11.1.0" }, @@ -8281,6 +8283,16 @@ "url": "https://github.com/sponsors/sindresorhus" } }, + "node_modules/markdown-table": { + "version": "3.0.4", + "resolved": "https://registry.npmjs.org/markdown-table/-/markdown-table-3.0.4.tgz", + "integrity": "sha512-wiYz4+JrLyb/DqW2hkFJxP7Vd7JuTDm77fvbM8VfEQdmSMqcImWeeRbHwZjBjIFki/VaMK2BhFi7oUUZeM5bqw==", + "license": "MIT", + "funding": { + "type": "github", + "url": "https://github.com/sponsors/wooorm" + } + }, "node_modules/math-intrinsics": { "version": "1.1.0", "resolved": "https://registry.npmjs.org/math-intrinsics/-/math-intrinsics-1.1.0.tgz", @@ -8290,6 +8302,34 @@ "node": ">= 0.4" } }, + "node_modules/mdast-util-find-and-replace": { + "version": "3.0.2", + "resolved": "https://registry.npmjs.org/mdast-util-find-and-replace/-/mdast-util-find-and-replace-3.0.2.tgz", + "integrity": "sha512-Tmd1Vg/m3Xz43afeNxDIhWRtFZgM2VLyaf4vSTYwudTyeuTneoL3qtWMA5jeLyz/O1vDJmmV4QuScFCA2tBPwg==", + "license": "MIT", + "dependencies": { + "@types/mdast": "^4.0.0", + "escape-string-regexp": "^5.0.0", + "unist-util-is": "^6.0.0", + "unist-util-visit-parents": "^6.0.0" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/unified" + } + }, + "node_modules/mdast-util-find-and-replace/node_modules/escape-string-regexp": { + "version": "5.0.0", + "resolved": "https://registry.npmjs.org/escape-string-regexp/-/escape-string-regexp-5.0.0.tgz", + "integrity": "sha512-/veY75JbMK4j1yjvuUxuVsiS/hr/4iHs9FTT6cgTexxdE0Ly/glccBAkloH/DofkjRbZU3bnoj38mOmhkZ0lHw==", + "license": "MIT", + "engines": { + "node": ">=12" + }, + "funding": { + "url": "https://github.com/sponsors/sindresorhus" + } + }, "node_modules/mdast-util-from-markdown": { "version": "2.0.2", "resolved": "https://registry.npmjs.org/mdast-util-from-markdown/-/mdast-util-from-markdown-2.0.2.tgz", @@ -8314,6 +8354,107 @@ "url": "https://opencollective.com/unified" } }, + "node_modules/mdast-util-gfm": { + "version": "3.1.0", + "resolved": "https://registry.npmjs.org/mdast-util-gfm/-/mdast-util-gfm-3.1.0.tgz", + "integrity": "sha512-0ulfdQOM3ysHhCJ1p06l0b0VKlhU0wuQs3thxZQagjcjPrlFRqY215uZGHHJan9GEAXd9MbfPjFJz+qMkVR6zQ==", + "license": "MIT", + "dependencies": { + "mdast-util-from-markdown": "^2.0.0", + "mdast-util-gfm-autolink-literal": "^2.0.0", + "mdast-util-gfm-footnote": "^2.0.0", + "mdast-util-gfm-strikethrough": "^2.0.0", + "mdast-util-gfm-table": "^2.0.0", + "mdast-util-gfm-task-list-item": "^2.0.0", + "mdast-util-to-markdown": "^2.0.0" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/unified" + } + }, + "node_modules/mdast-util-gfm-autolink-literal": { + "version": "2.0.1", + "resolved": "https://registry.npmjs.org/mdast-util-gfm-autolink-literal/-/mdast-util-gfm-autolink-literal-2.0.1.tgz", + "integrity": "sha512-5HVP2MKaP6L+G6YaxPNjuL0BPrq9orG3TsrZ9YXbA3vDw/ACI4MEsnoDpn6ZNm7GnZgtAcONJyPhOP8tNJQavQ==", + "license": "MIT", + "dependencies": { + "@types/mdast": "^4.0.0", + "ccount": "^2.0.0", + "devlop": "^1.0.0", + "mdast-util-find-and-replace": "^3.0.0", + "micromark-util-character": "^2.0.0" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/unified" + } + }, + "node_modules/mdast-util-gfm-footnote": { + "version": "2.1.0", + "resolved": "https://registry.npmjs.org/mdast-util-gfm-footnote/-/mdast-util-gfm-footnote-2.1.0.tgz", + "integrity": "sha512-sqpDWlsHn7Ac9GNZQMeUzPQSMzR6Wv0WKRNvQRg0KqHh02fpTz69Qc1QSseNX29bhz1ROIyNyxExfawVKTm1GQ==", + "license": "MIT", + "dependencies": { + "@types/mdast": "^4.0.0", + "devlop": "^1.1.0", + "mdast-util-from-markdown": "^2.0.0", + "mdast-util-to-markdown": "^2.0.0", + "micromark-util-normalize-identifier": "^2.0.0" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/unified" + } + }, + "node_modules/mdast-util-gfm-strikethrough": { + "version": "2.0.0", + "resolved": "https://registry.npmjs.org/mdast-util-gfm-strikethrough/-/mdast-util-gfm-strikethrough-2.0.0.tgz", + "integrity": "sha512-mKKb915TF+OC5ptj5bJ7WFRPdYtuHv0yTRxK2tJvi+BDqbkiG7h7u/9SI89nRAYcmap2xHQL9D+QG/6wSrTtXg==", + "license": "MIT", + "dependencies": { + "@types/mdast": "^4.0.0", + "mdast-util-from-markdown": "^2.0.0", + "mdast-util-to-markdown": "^2.0.0" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/unified" + } + }, + "node_modules/mdast-util-gfm-table": { + "version": "2.0.0", + "resolved": "https://registry.npmjs.org/mdast-util-gfm-table/-/mdast-util-gfm-table-2.0.0.tgz", + "integrity": "sha512-78UEvebzz/rJIxLvE7ZtDd/vIQ0RHv+3Mh5DR96p7cS7HsBhYIICDBCu8csTNWNO6tBWfqXPWekRuj2FNOGOZg==", + "license": "MIT", + "dependencies": { + "@types/mdast": "^4.0.0", + "devlop": "^1.0.0", + "markdown-table": "^3.0.0", + "mdast-util-from-markdown": "^2.0.0", + "mdast-util-to-markdown": "^2.0.0" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/unified" + } + }, + "node_modules/mdast-util-gfm-task-list-item": { + "version": "2.0.0", + "resolved": "https://registry.npmjs.org/mdast-util-gfm-task-list-item/-/mdast-util-gfm-task-list-item-2.0.0.tgz", + "integrity": "sha512-IrtvNvjxC1o06taBAVJznEnkiHxLFTzgonUdy8hzFVeDun0uTjxxrRGVaNFqkU1wJR3RBPEfsxmU6jDWPofrTQ==", + "license": "MIT", + "dependencies": { + "@types/mdast": "^4.0.0", + "devlop": "^1.0.0", + "mdast-util-from-markdown": "^2.0.0", + "mdast-util-to-markdown": "^2.0.0" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/unified" + } + }, "node_modules/mdast-util-mdx-expression": { "version": "2.0.1", "resolved": "https://registry.npmjs.org/mdast-util-mdx-expression/-/mdast-util-mdx-expression-2.0.1.tgz", @@ -8528,6 +8669,127 @@ "micromark-util-types": "^2.0.0" } }, + "node_modules/micromark-extension-gfm": { + "version": "3.0.0", + "resolved": "https://registry.npmjs.org/micromark-extension-gfm/-/micromark-extension-gfm-3.0.0.tgz", + "integrity": "sha512-vsKArQsicm7t0z2GugkCKtZehqUm31oeGBV/KVSorWSy8ZlNAv7ytjFhvaryUiCUJYqs+NoE6AFhpQvBTM6Q4w==", + "license": "MIT", + "dependencies": { + "micromark-extension-gfm-autolink-literal": "^2.0.0", + "micromark-extension-gfm-footnote": "^2.0.0", + "micromark-extension-gfm-strikethrough": "^2.0.0", + "micromark-extension-gfm-table": "^2.0.0", + "micromark-extension-gfm-tagfilter": "^2.0.0", + "micromark-extension-gfm-task-list-item": "^2.0.0", + "micromark-util-combine-extensions": "^2.0.0", + "micromark-util-types": "^2.0.0" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/unified" + } + }, + "node_modules/micromark-extension-gfm-autolink-literal": { + "version": "2.1.0", + "resolved": "https://registry.npmjs.org/micromark-extension-gfm-autolink-literal/-/micromark-extension-gfm-autolink-literal-2.1.0.tgz", + "integrity": "sha512-oOg7knzhicgQ3t4QCjCWgTmfNhvQbDDnJeVu9v81r7NltNCVmhPy1fJRX27pISafdjL+SVc4d3l48Gb6pbRypw==", + "license": "MIT", + "dependencies": { + "micromark-util-character": "^2.0.0", + "micromark-util-sanitize-uri": "^2.0.0", + "micromark-util-symbol": "^2.0.0", + "micromark-util-types": "^2.0.0" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/unified" + } + }, + "node_modules/micromark-extension-gfm-footnote": { + "version": "2.1.0", + "resolved": "https://registry.npmjs.org/micromark-extension-gfm-footnote/-/micromark-extension-gfm-footnote-2.1.0.tgz", + "integrity": "sha512-/yPhxI1ntnDNsiHtzLKYnE3vf9JZ6cAisqVDauhp4CEHxlb4uoOTxOCJ+9s51bIB8U1N1FJ1RXOKTIlD5B/gqw==", + "license": "MIT", + "dependencies": { + "devlop": "^1.0.0", + "micromark-core-commonmark": "^2.0.0", + "micromark-factory-space": "^2.0.0", + "micromark-util-character": "^2.0.0", + "micromark-util-normalize-identifier": "^2.0.0", + "micromark-util-sanitize-uri": "^2.0.0", + "micromark-util-symbol": "^2.0.0", + "micromark-util-types": "^2.0.0" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/unified" + } + }, + "node_modules/micromark-extension-gfm-strikethrough": { + "version": "2.1.0", + "resolved": "https://registry.npmjs.org/micromark-extension-gfm-strikethrough/-/micromark-extension-gfm-strikethrough-2.1.0.tgz", + "integrity": "sha512-ADVjpOOkjz1hhkZLlBiYA9cR2Anf8F4HqZUO6e5eDcPQd0Txw5fxLzzxnEkSkfnD0wziSGiv7sYhk/ktvbf1uw==", + "license": "MIT", + "dependencies": { + "devlop": "^1.0.0", + "micromark-util-chunked": "^2.0.0", + "micromark-util-classify-character": "^2.0.0", + "micromark-util-resolve-all": "^2.0.0", + "micromark-util-symbol": "^2.0.0", + "micromark-util-types": "^2.0.0" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/unified" + } + }, + "node_modules/micromark-extension-gfm-table": { + "version": "2.1.1", + "resolved": "https://registry.npmjs.org/micromark-extension-gfm-table/-/micromark-extension-gfm-table-2.1.1.tgz", + "integrity": "sha512-t2OU/dXXioARrC6yWfJ4hqB7rct14e8f7m0cbI5hUmDyyIlwv5vEtooptH8INkbLzOatzKuVbQmAYcbWoyz6Dg==", + "license": "MIT", + "dependencies": { + "devlop": "^1.0.0", + "micromark-factory-space": "^2.0.0", + "micromark-util-character": "^2.0.0", + "micromark-util-symbol": "^2.0.0", + "micromark-util-types": "^2.0.0" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/unified" + } + }, + "node_modules/micromark-extension-gfm-tagfilter": { + "version": "2.0.0", + "resolved": "https://registry.npmjs.org/micromark-extension-gfm-tagfilter/-/micromark-extension-gfm-tagfilter-2.0.0.tgz", + "integrity": "sha512-xHlTOmuCSotIA8TW1mDIM6X2O1SiX5P9IuDtqGonFhEK0qgRI4yeC6vMxEV2dgyr2TiD+2PQ10o+cOhdVAcwfg==", + "license": "MIT", + "dependencies": { + "micromark-util-types": "^2.0.0" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/unified" + } + }, + "node_modules/micromark-extension-gfm-task-list-item": { + "version": "2.1.0", + "resolved": "https://registry.npmjs.org/micromark-extension-gfm-task-list-item/-/micromark-extension-gfm-task-list-item-2.1.0.tgz", + "integrity": "sha512-qIBZhqxqI6fjLDYFTBIa4eivDMnP+OZqsNwmQ3xNLE4Cxwc+zfQEfbs6tzAo2Hjq+bh6q5F+Z8/cksrLFYWQQw==", + "license": "MIT", + "dependencies": { + "devlop": "^1.0.0", + "micromark-factory-space": "^2.0.0", + "micromark-util-character": "^2.0.0", + "micromark-util-symbol": "^2.0.0", + "micromark-util-types": "^2.0.0" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/unified" + } + }, "node_modules/micromark-factory-destination": { "version": "2.0.1", "resolved": "https://registry.npmjs.org/micromark-factory-destination/-/micromark-factory-destination-2.0.1.tgz", @@ -11006,6 +11268,24 @@ "url": "https://github.com/sponsors/ljharb" } }, + "node_modules/remark-gfm": { + "version": "4.0.1", + "resolved": "https://registry.npmjs.org/remark-gfm/-/remark-gfm-4.0.1.tgz", + "integrity": "sha512-1quofZ2RQ9EWdeN34S79+KExV1764+wCUGop5CPL1WGdD0ocPpu91lzPGbwWMECpEpd42kJGQwzRfyov9j4yNg==", + "license": "MIT", + "dependencies": { + "@types/mdast": "^4.0.0", + "mdast-util-gfm": "^3.0.0", + "micromark-extension-gfm": "^3.0.0", + "remark-parse": "^11.0.0", + "remark-stringify": "^11.0.0", + "unified": "^11.0.0" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/unified" + } + }, "node_modules/remark-parse": { "version": "11.0.0", "resolved": "https://registry.npmjs.org/remark-parse/-/remark-parse-11.0.0.tgz", @@ -11039,6 +11319,21 @@ "url": "https://opencollective.com/unified" } }, + "node_modules/remark-stringify": { + "version": "11.0.0", + "resolved": "https://registry.npmjs.org/remark-stringify/-/remark-stringify-11.0.0.tgz", + "integrity": "sha512-1OSmLd3awB/t8qdoEOMazZkNsfVTeY4fTsgzcQFdXNq8ToTN4ZGwrMnlda4K6smTFKD+GRV6O48i6Z4iKgPPpw==", + "license": "MIT", + "dependencies": { + "@types/mdast": "^4.0.0", + "mdast-util-to-markdown": "^2.0.0", + "unified": "^11.0.0" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/unified" + } + }, "node_modules/require-from-string": { "version": "2.0.2", "resolved": "https://registry.npmjs.org/require-from-string/-/require-from-string-2.0.2.tgz", diff --git a/ui/litellm-dashboard/package.json b/ui/litellm-dashboard/package.json index 567673c0989..ea84ea6f401 100644 --- a/ui/litellm-dashboard/package.json +++ b/ui/litellm-dashboard/package.json @@ -31,6 +31,7 @@ "@types/papaparse": "^5.3.15", "antd": "^5.13.2", "cva": "^1.0.0-beta.3", + "dayjs": "^1.11.19", "jwt-decode": "^4.0.0", "lucide-react": "^0.513.0", "moment": "^2.30.1", @@ -43,6 +44,7 @@ "react-json-view-lite": "^2.5.0", "react-markdown": "^9.0.1", "react-syntax-highlighter": "^15.6.6", + "remark-gfm": "^4.0.1", "tailwind-merge": "^3.2.0", "uuid": "^11.1.0" }, @@ -107,4 +109,4 @@ "node": ">=18.17.0", "npm": ">=8.3.0" } -} \ No newline at end of file +} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/playground/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/playground/page.tsx index 555930a576c..6a694d9bee9 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/playground/page.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/playground/page.tsx @@ -8,6 +8,7 @@ import ComplianceUI from "@/components/playground/complianceUI/ComplianceUI"; import { TabGroup, TabList, Tab, TabPanels, TabPanel } from "@tremor/react"; import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; import { fetchProxySettings } from "@/utils/proxyUtils"; +import { MessageOutlined, CloseOutlined } from "@ant-design/icons"; interface ProxySettings { PROXY_BASE_URL?: string; @@ -17,6 +18,7 @@ interface ProxySettings { export default function PlaygroundPage() { const { accessToken, userRole, userId, disabledPersonalKeyCreation, token } = useAuthorized(); const [proxySettings, setProxySettings] = useState(undefined); + const [chatBannerDismissed, setChatBannerDismissed] = useState(false); useEffect(() => { const initializeProxySettings = async () => { @@ -35,7 +37,66 @@ export default function PlaygroundPage() { }, [accessToken]); return ( - +
+ {!chatBannerDismissed && ( +
+ + New + + + Chat UI + {" "}— a ChatGPT-like interface for your users to chat with AI models and MCP tools. Share it with your team. + + + Open Chat UI → + + +
+ )} + Chat Compare @@ -72,5 +133,6 @@ export default function PlaygroundPage() { +
); } diff --git a/ui/litellm-dashboard/src/app/chat/page.tsx b/ui/litellm-dashboard/src/app/chat/page.tsx new file mode 100644 index 00000000000..18fc02e7f73 --- /dev/null +++ b/ui/litellm-dashboard/src/app/chat/page.tsx @@ -0,0 +1,19 @@ +"use client"; + +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import ChatPage from "@/components/chat/ChatPage"; + +const ChatPageRoute = () => { + const { accessToken, userRole, userId, userEmail } = useAuthorized(); + + return ( + + ); +}; + +export default ChatPageRoute; diff --git a/ui/litellm-dashboard/src/components/chat/ChatMessages.tsx b/ui/litellm-dashboard/src/components/chat/ChatMessages.tsx new file mode 100644 index 00000000000..640a8addef0 --- /dev/null +++ b/ui/litellm-dashboard/src/components/chat/ChatMessages.tsx @@ -0,0 +1,577 @@ +"use client"; + +import { ToolOutlined, CopyOutlined, CheckOutlined, EditOutlined } from "@ant-design/icons"; +import { Collapse, Tooltip } from "antd"; +import React, { useEffect, useRef, useState } from "react"; +import ReactMarkdown from "react-markdown"; +import remarkGfm from "remark-gfm"; +import { Prism as SyntaxHighlighter } from "react-syntax-highlighter"; +import { coy } from "react-syntax-highlighter/dist/esm/styles/prism"; +import ReasoningContent from "../playground/chat_ui/ReasoningContent"; +import { ChatMessage } from "./types"; + +const { Panel } = Collapse; + +// Keys whose values must be redacted in tool args display +const REDACTED_KEY_PATTERNS = /token|key|secret|password|auth/i; + +function redactSensitiveValues(obj: Record): Record { + const result: Record = {}; + for (const [k, v] of Object.entries(obj)) { + if (REDACTED_KEY_PATTERNS.test(k)) { + result[k] = "[redacted]"; + } else if (Array.isArray(v)) { + result[k] = v.map((item) => + item !== null && typeof item === "object" && !Array.isArray(item) + ? redactSensitiveValues(item as Record) + : item, + ); + } else if (v !== null && typeof v === "object") { + result[k] = redactSensitiveValues(v as Record); + } else { + result[k] = v; + } + } + return result; +} + +function formatTimestamp(ts: number): string { + const d = new Date(ts); + const hh = String(d.getHours()).padStart(2, "0"); + const mm = String(d.getMinutes()).padStart(2, "0"); + return `${hh}:${mm}`; +} + +// Shared markdown code renderer matching ReasoningContent style. +// react-markdown v9 removed the `inline` prop; detect fenced blocks via language className. +function MarkdownCodeRenderer({ + node, + className, + children, + ...props +}: React.ComponentPropsWithoutRef<"code"> & { node?: unknown }) { + const match = /language-(\w+)/.exec(className || ""); + return match ? ( + } + language={match[1]} + PreTag="div" + className="rounded-md my-2" + {...(props as Record)} + > + {String(children).replace(/\n$/, "")} + + ) : ( + + {children} + + ); +} + +// ------- Sub-components ------- + +interface UserBubbleProps { + message: ChatMessage; + onEdit?: (messageId: string, newContent: string) => void; + isStreaming?: boolean; +} + +function UserBubble({ message, onEdit, isStreaming }: UserBubbleProps) { + const [hovered, setHovered] = useState(false); + const [editing, setEditing] = useState(false); + const [editValue, setEditValue] = useState(message.content); + const textareaRef = useRef(null); + + useEffect(() => { + if (editing && textareaRef.current) { + textareaRef.current.focus(); + textareaRef.current.selectionStart = textareaRef.current.value.length; + } + }, [editing]); + + // Auto-resize textarea + useEffect(() => { + const ta = textareaRef.current; + if (!ta) return; + ta.style.height = "auto"; + ta.style.height = `${ta.scrollHeight}px`; + }, [editValue, editing]); + + const handleSave = () => { + const trimmed = editValue.trim(); + if (trimmed && trimmed !== message.content && onEdit) { + onEdit(message.id, trimmed); + } + setEditing(false); + }; + + const handleKeyDown = (e: React.KeyboardEvent) => { + if (e.key === "Enter" && !e.shiftKey) { + e.preventDefault(); + handleSave(); + } + if (e.key === "Escape") { + setEditValue(message.content); + setEditing(false); + } + }; + + if (editing) { + return ( +
+
+