From 70f7c73defbf95af55afcf5997fbe18eeac90539 Mon Sep 17 00:00:00 2001 From: Krish Dholakia Date: Fri, 25 Apr 2025 18:20:41 -0700 Subject: [PATCH] Move UI to encrypted token usage (#10302) * test(test_auth_checks.py): add unit tests for ExperimentalUIJWTToken * test: add appropriate flag * fix: fix ruff check * test: add autouse fixture to test salt key * fix(user_api_key_auth.py): fix auth flow logic * test: skip flaky test - anthropic does not reliably return 'redacted_thinking' --- litellm/proxy/_new_secret_config.yaml | 6 +- litellm/proxy/auth/auth_checks.py | 56 +++++++++ litellm/proxy/auth/user_api_key_auth.py | 8 ++ .../common_utils/encrypt_decrypt_utils.py | 16 ++- .../management_endpoints/team_endpoints.py | 1 + litellm/proxy/management_endpoints/ui_sso.py | 4 +- litellm/proxy/proxy_server.py | 40 ++++++- test-results/.last-run.json | 4 - tests/litellm/proxy/auth/test_auth_checks.py | 111 ++++++++++++++++++ .../test_anthropic_completion.py | 27 ----- 10 files changed, 230 insertions(+), 43 deletions(-) delete mode 100644 test-results/.last-run.json create mode 100644 tests/litellm/proxy/auth/test_auth_checks.py diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index ad9ae99f7a..67fab9732f 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -4,6 +4,11 @@ model_list: model: azure/gpt-4o api_key: os.environ/AZURE_API_KEY api_base: os.environ/AZURE_API_BASE + - model_name: fake-openai-endpoint + litellm_params: + model: openai/fake + api_key: fake-key + api_base: https://exampleopenaiendpoint-production.up.railway.app/ - model_name: "gpt-4o-mini-openai" litellm_params: model: gpt-4o-mini @@ -33,7 +38,6 @@ model_list: litellm_settings: num_retries: 0 - callbacks: ["datadog_llm_observability"] check_provider_endpoint: true files_settings: diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index e46cc858ac..dac817e811 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -953,6 +953,62 @@ async def get_team_object( ) +class ExperimentalUIJWTToken: + @staticmethod + def get_experimental_ui_login_jwt_auth_token(user_info: LiteLLM_UserTable) -> str: + from datetime import UTC, datetime, timedelta + + from litellm.proxy.common_utils.encrypt_decrypt_utils import ( + encrypt_value_helper, + ) + + if user_info.user_role is None: + raise Exception("User role is required for experimental UI login") + + # Calculate expiration time (10 minutes from now) + expiration_time = datetime.now(UTC) + timedelta(minutes=10) + + # Format the expiration time as ISO 8601 string + expires = expiration_time.strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3] + "+00:00" + + valid_token = UserAPIKeyAuth( + token="ui-token", + key_name="ui-token", + key_alias="ui-token", + max_budget=litellm.max_ui_session_budget, + rpm_limit=100, # allow user to have a conversation on test key pane of UI + expires=expires, + user_id=user_info.user_id, + team_id="litellm-dashboard", + models=user_info.models, + max_parallel_requests=None, + user_role=LitellmUserRoles(user_info.user_role), + ) + + return encrypt_value_helper(valid_token.model_dump_json(exclude_none=True)) + + @staticmethod + def get_key_object_from_ui_hash_key( + hashed_token: str, + ) -> Optional[UserAPIKeyAuth]: + import json + + from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth + from litellm.proxy.common_utils.encrypt_decrypt_utils import ( + decrypt_value_helper, + ) + + decrypted_token = decrypt_value_helper(hashed_token, exception_type="debug") + if decrypted_token is None: + return None + try: + return UserAPIKeyAuth(**json.loads(decrypted_token)) + except Exception as e: + raise Exception( + f"Invalid hash key. Hash key={hashed_token}. Decrypted token={decrypted_token}. Error: {e}" + ) + + @log_db_metrics async def get_key_object( hashed_token: str, diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 97e9fb8c73..2b47d527ab 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -24,6 +24,7 @@ from litellm.caching import DualCache from litellm.litellm_core_utils.dd_tracing import tracer from litellm.proxy._types import * from litellm.proxy.auth.auth_checks import ( + ExperimentalUIJWTToken, _cache_key_object, _get_user_role, _is_user_proxy_admin, @@ -51,6 +52,7 @@ from litellm.proxy.auth.oauth2_check import check_oauth2_token from litellm.proxy.auth.oauth2_proxy_hook import handle_oauth2_proxy_request from litellm.proxy.common_utils.http_parsing_utils import _read_request_body from litellm.proxy.utils import PrismaClient, ProxyLogging +from litellm.secret_managers.main import get_secret_bool from litellm.types.services import ServiceTypes user_api_key_service_logger_obj = ServiceLogging() # used for tracking latency on OTEL @@ -553,6 +555,12 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 verbose_logger.debug("api key not found in cache.") valid_token = None + ## Check UI Hash Key + if valid_token is None and get_secret_bool("EXPERIMENTAL_UI_LOGIN"): + valid_token = ExperimentalUIJWTToken.get_key_object_from_ui_hash_key( + api_key + ) + if ( valid_token is not None and isinstance(valid_token, UserAPIKeyAuth) diff --git a/litellm/proxy/common_utils/encrypt_decrypt_utils.py b/litellm/proxy/common_utils/encrypt_decrypt_utils.py index 348c81101f..bee098cd32 100644 --- a/litellm/proxy/common_utils/encrypt_decrypt_utils.py +++ b/litellm/proxy/common_utils/encrypt_decrypt_utils.py @@ -1,6 +1,6 @@ import base64 import os -from typing import Optional +from typing import Literal, Optional from litellm._logging import verbose_proxy_logger @@ -39,7 +39,9 @@ def encrypt_value_helper(value: str, new_encryption_key: Optional[str] = None): raise e -def decrypt_value_helper(value: str): +def decrypt_value_helper( + value: str, exception_type: Literal["debug", "error"] = "error" +): signing_key = _get_salt_key() try: @@ -51,11 +53,13 @@ def decrypt_value_helper(value: str): # if it's not str - do not decrypt it, return the value return value except Exception as e: - verbose_proxy_logger.error( - f"Error decrypting value, Did your master_key/salt key change recently? \nError: {str(e)}\nSet permanent salt key - https://docs.litellm.ai/docs/proxy/prod#5-set-litellm-salt-key" - ) + error_message = f"Error decrypting value, Did your master_key/salt key change recently? \nError: {str(e)}\nSet permanent salt key - https://docs.litellm.ai/docs/proxy/prod#5-set-litellm-salt-key" + if exception_type == "debug": + verbose_proxy_logger.debug(error_message) + return None + verbose_proxy_logger.error(error_message) # [Non-Blocking Exception. - this should not block decrypting other values] - pass + return None def encrypt_value(value: str, signing_key: str): diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 81b2072eb4..fbd1eb5f7d 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -538,6 +538,7 @@ async def update_team( detail={"error": "Team doesn't exist. Got={}".format(team_row)}, ) + verbose_proxy_logger.info("Successfully updated team - %s, info", team_row.team_id) await _cache_team_object( team_id=team_row.team_id, team_table=LiteLLM_TeamTableCachedObj(**team_row.model_dump()), diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index cd91bda264..4f53c03702 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -642,7 +642,9 @@ async def auth_callback(request: Request): # noqa: PLR0915 litellm_dashboard_ui += "?login=success" verbose_proxy_logger.info(f"Redirecting to {litellm_dashboard_ui}") redirect_response = RedirectResponse(url=litellm_dashboard_ui, status_code=303) - redirect_response.set_cookie(key="token", value=jwt_token, secure=True) + redirect_response.set_cookie( + key="token", value=jwt_token, secure=True, httponly=True + ) return redirect_response diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 8c5b88be4d..569887f0ac 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -150,7 +150,11 @@ from litellm.proxy.analytics_endpoints.analytics_endpoints import ( router as analytics_router, ) from litellm.proxy.anthropic_endpoints.endpoints import router as anthropic_router -from litellm.proxy.auth.auth_checks import get_team_object, log_db_metrics +from litellm.proxy.auth.auth_checks import ( + ExperimentalUIJWTToken, + get_team_object, + log_db_metrics, +) from litellm.proxy.auth.auth_utils import check_response_size_is_safe from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.proxy.auth.litellm_license import LicenseCheck @@ -6724,7 +6728,7 @@ async def login(request: Request): # noqa: PLR0915 ) # check if we can find the `username` in the db. on the ui, users can enter username=their email - _user_row = None + _user_row: Optional[LiteLLM_UserTable] = None user_role: Optional[ Literal[ LitellmUserRoles.PROXY_ADMIN, @@ -6734,8 +6738,11 @@ async def login(request: Request): # noqa: PLR0915 ] ] = None if prisma_client is not None: - _user_row = await prisma_client.db.litellm_usertable.find_first( - where={"user_email": {"equals": username}} + _user_row = cast( + Optional[LiteLLM_UserTable], + await prisma_client.db.litellm_usertable.find_first( + where={"user_email": {"equals": username}} + ), ) disabled_non_admin_personal_key_creation = ( get_disabled_non_admin_personal_key_creation() @@ -6800,6 +6807,31 @@ async def login(request: Request): # noqa: PLR0915 litellm_dashboard_ui += "/ui/" import jwt + if get_secret_bool("EXPERIMENTAL_UI_LOGIN"): + user_info: Optional[LiteLLM_UserTable] = None + if _user_row is not None: + user_info = _user_row + elif ( + user_id is not None + ): # if user_id is not None, we are using the UI_USERNAME and UI_PASSWORD + user_info = LiteLLM_UserTable( + user_id=user_id, + user_role=user_role, + models=[], + max_budget=litellm.max_ui_session_budget, + ) + if user_info is None: + raise HTTPException( + status_code=401, + detail={ + "error": "User Information is required for experimental UI login" + }, + ) + + key = ExperimentalUIJWTToken.get_experimental_ui_login_jwt_auth_token( + user_info + ) + jwt_token = jwt.encode( # type: ignore { "user_id": user_id, diff --git a/test-results/.last-run.json b/test-results/.last-run.json deleted file mode 100644 index 5fca3f84bc..0000000000 --- a/test-results/.last-run.json +++ /dev/null @@ -1,4 +0,0 @@ -{ - "status": "failed", - "failedTests": [] -} \ No newline at end of file diff --git a/tests/litellm/proxy/auth/test_auth_checks.py b/tests/litellm/proxy/auth/test_auth_checks.py new file mode 100644 index 0000000000..7e8d99d4ee --- /dev/null +++ b/tests/litellm/proxy/auth/test_auth_checks.py @@ -0,0 +1,111 @@ +import asyncio +import json +import os +import sys +from unittest.mock import AsyncMock, MagicMock, patch + +sys.path.insert( + 0, os.path.abspath("../../..") +) # Adds the parent directory to the system path + +from datetime import UTC, datetime, timedelta + +import pytest + +import litellm +from litellm.proxy._types import ( + LiteLLM_UserTable, + LitellmUserRoles, + SSOUserDefinedValues, +) +from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken +from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper + + +@pytest.fixture(autouse=True) +def set_salt_key(monkeypatch): + """Automatically set LITELLM_SALT_KEY for all tests""" + monkeypatch.setenv("LITELLM_SALT_KEY", "sk-1234") + + +@pytest.fixture +def valid_sso_user_defined_values(): + return LiteLLM_UserTable( + user_id="test_user", + user_email="test@example.com", + user_role=LitellmUserRoles.PROXY_ADMIN.value, + models=["gpt-3.5-turbo"], + max_budget=100.0, + ) + + +@pytest.fixture +def invalid_sso_user_defined_values(): + return LiteLLM_UserTable( + user_id="test_user", + user_email="test@example.com", + user_role=None, # Missing user role + models=["gpt-3.5-turbo"], + max_budget=100.0, + ) + + +def test_get_experimental_ui_login_jwt_auth_token_valid(valid_sso_user_defined_values): + """Test generating JWT token with valid user role""" + token = ExperimentalUIJWTToken.get_experimental_ui_login_jwt_auth_token( + valid_sso_user_defined_values + ) + + # Decrypt and verify token contents + decrypted_token = decrypt_value_helper(token, exception_type="debug") + token_data = json.loads(decrypted_token) + + assert token_data["user_id"] == "test_user" + assert token_data["user_role"] == LitellmUserRoles.PROXY_ADMIN.value + assert token_data["models"] == ["gpt-3.5-turbo"] + assert token_data["max_budget"] == litellm.max_ui_session_budget + + # Verify expiration time is set and valid + assert "expires" in token_data + expires = datetime.fromisoformat(token_data["expires"].replace("Z", "+00:00")) + assert expires > datetime.now(UTC) + assert expires <= datetime.now(UTC) + timedelta(minutes=10) + + +def test_get_experimental_ui_login_jwt_auth_token_invalid( + invalid_sso_user_defined_values, +): + """Test generating JWT token with missing user role""" + with pytest.raises(Exception) as exc_info: + ExperimentalUIJWTToken.get_experimental_ui_login_jwt_auth_token( + invalid_sso_user_defined_values + ) + + assert str(exc_info.value) == "User role is required for experimental UI login" + + +def test_get_key_object_from_ui_hash_key_valid( + valid_sso_user_defined_values, monkeypatch +): + """Test getting key object from valid UI hash key""" + monkeypatch.setenv("EXPERIMENTAL_UI_LOGIN", "True") + # Generate a valid token + token = ExperimentalUIJWTToken.get_experimental_ui_login_jwt_auth_token( + valid_sso_user_defined_values + ) + + # Get key object + key_object = ExperimentalUIJWTToken.get_key_object_from_ui_hash_key(token) + + assert key_object is not None + assert key_object.user_id == "test_user" + assert key_object.user_role == LitellmUserRoles.PROXY_ADMIN + assert key_object.models == ["gpt-3.5-turbo"] + assert key_object.max_budget == litellm.max_ui_session_budget + + +def test_get_key_object_from_ui_hash_key_invalid(): + """Test getting key object from invalid UI hash key""" + # Test with invalid token + key_object = ExperimentalUIJWTToken.get_key_object_from_ui_hash_key("invalid_token") + assert key_object is None diff --git a/tests/llm_translation/test_anthropic_completion.py b/tests/llm_translation/test_anthropic_completion.py index 73d0bec5bf..40b2a8a27d 100644 --- a/tests/llm_translation/test_anthropic_completion.py +++ b/tests/llm_translation/test_anthropic_completion.py @@ -997,33 +997,6 @@ def test_anthropic_thinking_output(model): assert resp.choices[0].message.thinking_blocks[0]["signature"] is not None -@pytest.mark.parametrize( - "model", - [ - "anthropic/claude-3-7-sonnet-20250219", - # "bedrock/us.anthropic.claude-3-7-sonnet-20250219-v1:0", - ], -) -def test_anthropic_redacted_thinking_output(model): - from litellm import completion - - litellm._turn_on_debug() - - resp = completion( - model=model, - messages=[{"role": "user", "content": "ANTHROPIC_MAGIC_STRING_TRIGGER_REDACTED_THINKING_46C9A13E193C177646C7398A98432ECCCE4C1253D5E2D82641AC0E52CC2876CB"}], - thinking={"type": "enabled", "budget_tokens": 1024}, - ) - - print(resp) - assert resp.choices[0].message.thinking_blocks is not None - assert isinstance(resp.choices[0].message.thinking_blocks, list) - assert len(resp.choices[0].message.thinking_blocks) > 0 - assert resp.choices[0].message.thinking_blocks[0]["type"] == "redacted_thinking" - assert resp.choices[0].message.thinking_blocks[0]["data"] is not None - - - @pytest.mark.parametrize( "model",