From fa20abfe48849ed6794252937de602f60db30ae3 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 22 Sep 2025 10:07:16 -0700 Subject: [PATCH] [Feat] Proxy CLI Auth - Allow re-using cli auth token (#14780) * fix: cli auth with SSO okta * fix: add LITTELM_CLI_SERVICE_ACCOUNT_NAME * fix: get_litellm_cli_user_api_key_auth * use existing_key CLI * fix: use existing key * test auth commands * test_cli_sso_callback_regenerate_vs_create_flow --- litellm/constants.py | 1 + litellm/proxy/_types.py | 16 ++ litellm/proxy/client/cli/commands/auth.py | 7 + litellm/proxy/management_endpoints/ui_sso.py | 153 +++++++++++++----- .../proxy/client/cli/test_auth_commands.py | 102 ++++++++++++ .../proxy/management_endpoints/test_ui_sso.py | 150 +++++++++++++++++ 6 files changed, 392 insertions(+), 37 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 9b44613b85..dd61ba3fa8 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -947,6 +947,7 @@ HEALTH_CHECK_TIMEOUT_SECONDS = int( os.getenv("HEALTH_CHECK_TIMEOUT_SECONDS", 60) ) # 60 seconds LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME = "litellm-internal-health-check" +LITTELM_CLI_SERVICE_ACCOUNT_NAME = "litellm-cli" UI_SESSION_TOKEN_TEAM_ID = "litellm-dashboard" LITELLM_PROXY_ADMIN_NAME = "default_user_id" diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index b0a4e71e23..9d5298c30f 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1915,6 +1915,22 @@ class UserAPIKeyAuth( key_alias=LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME, team_alias=LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME, ) + + @classmethod + def get_litellm_cli_user_api_key_auth(cls) -> "UserAPIKeyAuth": + """ + Returns a `UserAPIKeyAuth` object for the litellm internal health check service account. + + This is used to track number of requests/spend for health check calls. + """ + from litellm.constants import LITTELM_CLI_SERVICE_ACCOUNT_NAME + + return cls( + api_key=LITTELM_CLI_SERVICE_ACCOUNT_NAME, + team_id=LITTELM_CLI_SERVICE_ACCOUNT_NAME, + key_alias=LITTELM_CLI_SERVICE_ACCOUNT_NAME, + team_alias=LITTELM_CLI_SERVICE_ACCOUNT_NAME, + ) class UserInfoResponse(LiteLLMPydanticObjectBase): diff --git a/litellm/proxy/client/cli/commands/auth.py b/litellm/proxy/client/cli/commands/auth.py index 7d89d39ed7..ec57f64946 100644 --- a/litellm/proxy/client/cli/commands/auth.py +++ b/litellm/proxy/client/cli/commands/auth.py @@ -64,6 +64,9 @@ def login(ctx: click.Context): base_url = ctx.obj["base_url"] + # Check if we have an existing key to regenerate + existing_key = get_stored_api_key() + # Generate unique key ID for this login session key_id = f"sk-{str(uuid.uuid4())}" @@ -71,6 +74,10 @@ def login(ctx: click.Context): # Construct SSO login URL with CLI source and pre-generated key sso_url = f"{base_url}/sso/key/generate?source={LITELLM_CLI_SOURCE_IDENTIFIER}&key={key_id}" + # If we have an existing key, include it so the server can regenerate it + if existing_key: + sso_url += f"&existing_key={existing_key}" + click.echo(f"Opening browser to: {sso_url}") click.echo("Please complete the SSO authentication in your browser...") click.echo(f"Session ID: {key_id}") diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index e6fdcef280..3c9fec3dee 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -114,7 +114,7 @@ def process_sso_jwt_access_token( @router.get("/sso/key/generate", tags=["experimental"], include_in_schema=False) async def google_login( - request: Request, source: Optional[str] = None, key: Optional[str] = None + request: Request, source: Optional[str] = None, key: Optional[str] = None, existing_key: Optional[str] = None ): # noqa: PLR0915 """ Create Proxy API Keys using Google Workspace SSO. Requires setting PROXY_BASE_URL in .env @@ -173,12 +173,14 @@ async def google_login( redirect_url = SSOAuthenticationHandler.get_redirect_url_for_sso( request=request, sso_callback_route="sso/callback", + existing_key=existing_key, ) # Store CLI key in state for OAuth flow cli_state: Optional[str] = SSOAuthenticationHandler._get_cli_state( source=source, key=key, + existing_key=existing_key, ) # check if user defined a custom auth sso sign in handler, if yes, use it @@ -590,8 +592,12 @@ async def auth_callback(request: Request, state: Optional[str] = None): # noqa: if state and state.startswith(f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:"): # Extract the key ID from the state key_id = state.split(":", 1)[1] - verbose_proxy_logger.info(f"CLI SSO callback detected for key: {key_id}") - return await cli_sso_callback(request, key=key_id) + + # Get existing_key from query parameters if provided + existing_key = request.query_params.get("existing_key") + + verbose_proxy_logger.info(f"CLI SSO callback detected for key: {key_id}, existing_key: {existing_key}") + return await cli_sso_callback(request, key=key_id, existing_key=existing_key) from litellm.proxy._types import LiteLLM_JWTAuth from litellm.proxy.auth.handle_jwt import JWTHandler @@ -678,13 +684,59 @@ async def auth_callback(request: Request, state: Optional[str] = None): # noqa: ) -async def cli_sso_callback(request: Request, key: Optional[str] = None): - """CLI SSO callback - generates the key with pre-specified ID""" - verbose_proxy_logger.info(f"CLI SSO callback for key: {key}") +async def _regenerate_cli_key(existing_key: str, new_key: str) -> None: + """Regenerate an existing CLI key with a new token""" + from litellm.proxy._types import RegenerateKeyRequest, UserAPIKeyAuth + from litellm.proxy.management_endpoints.key_management_endpoints import ( + regenerate_key_fn, + ) + + verbose_proxy_logger.info(f"Regenerating existing CLI key: {existing_key}") + + admin_user_dict = UserAPIKeyAuth.get_litellm_cli_user_api_key_auth() + + regenerate_request = RegenerateKeyRequest( + key=existing_key, + new_key=new_key, + duration="24hr" + ) + + await regenerate_key_fn( + key=existing_key, + data=regenerate_request, + user_api_key_dict=admin_user_dict + ) + + verbose_proxy_logger.info(f"Regenerated CLI key: {new_key}") + +async def _create_new_cli_key(key: str) -> None: + """Create a new CLI key""" from litellm.proxy.management_endpoints.key_management_endpoints import ( generate_key_helper_fn, ) + + verbose_proxy_logger.info("Creating new CLI key") + + await generate_key_helper_fn( + request_type="key", + duration="24hr", + key_max_budget=litellm.max_ui_session_budget, + aliases={}, + config={}, + spend=0, + team_id="litellm-cli", + table_name="key", + token=key, + ) + + verbose_proxy_logger.info(f"Created new CLI key: {key}") + + +async def cli_sso_callback(request: Request, key: Optional[str] = None, existing_key: Optional[str] = None): + """CLI SSO callback - regenerates existing CLI key or creates new one""" + verbose_proxy_logger.info(f"CLI SSO callback for key: {key}, existing_key: {existing_key}") + from litellm.proxy.proxy_server import prisma_client if not key or not key.startswith("sk-"): @@ -698,21 +750,11 @@ async def cli_sso_callback(request: Request, key: Optional[str] = None): status_code=500, detail=CommonProxyErrors.db_not_connected_error.value ) - # Generate a simple key for CLI usage with the pre-specified key ID try: - await generate_key_helper_fn( - request_type="key", - duration="24hr", - key_max_budget=litellm.max_ui_session_budget, - aliases={}, - config={}, - spend=0, - team_id="litellm-cli", - table_name="key", - token=key, # Use the pre-specified key ID - ) - - verbose_proxy_logger.info(f"Generated CLI key: {key}") + if existing_key: + await _regenerate_cli_key(existing_key, key) + else: + await _create_new_cli_key(key) # Return success page from fastapi.responses import HTMLResponse @@ -725,8 +767,8 @@ async def cli_sso_callback(request: Request, key: Optional[str] = None): return HTMLResponse(content=html_content, status_code=200) except Exception as e: - verbose_proxy_logger.error(f"Error generating CLI key: {e}") - raise HTTPException(status_code=500, detail=f"Failed to generate key: {str(e)}") + verbose_proxy_logger.error(f"Error with CLI key: {e}") + raise HTTPException(status_code=500, detail=f"Failed to process CLI key: {str(e)}") @router.get("/sso/cli/poll/{key_id}", tags=["experimental"], include_in_schema=False) @@ -993,20 +1035,51 @@ class SSOAuthenticationHandler: # or a cryptographicly signed state that we can verify stateless # For simplification we are using a static state, this is not perfect but some # SSO providers do not allow stateless verification - redirect_params = {} - state = os.getenv("GENERIC_CLIENT_STATE", None) - - if state: - redirect_params["state"] = state - elif "okta" in generic_authorization_endpoint: - redirect_params["state"] = ( - uuid.uuid4().hex - ) # set state param for okta - required + redirect_params = SSOAuthenticationHandler._get_generic_sso_redirect_params( + state=state, + generic_authorization_endpoint=generic_authorization_endpoint + ) + return await generic_sso.get_login_redirect(**redirect_params) # type: ignore raise ValueError( "Unknown SSO provider. Please setup SSO with client IDs https://docs.litellm.ai/docs/proxy/admin_ui_sso" ) + @staticmethod + def _get_generic_sso_redirect_params( + state: Optional[str] = None, + generic_authorization_endpoint: Optional[str] = None + ) -> dict: + """ + Get redirect parameters for Generic SSO with proper state priority handling. + + Priority order: + 1. CLI state (if provided) + 2. GENERIC_CLIENT_STATE environment variable + 3. Generated UUID for Okta (if Okta endpoint detected) + + Args: + state: Optional state parameter (e.g., CLI state) + generic_authorization_endpoint: Authorization endpoint URL + + Returns: + dict: Redirect parameters for SSO login + """ + redirect_params = {} + + if state: + # CLI state takes priority + # the litellm proxy cli sends the "state" parameter to the proxy server for auth. We should maintain the state parameter for the cli if it is provided + redirect_params["state"] = state + else: + generic_client_state = os.getenv("GENERIC_CLIENT_STATE", None) + if generic_client_state: + redirect_params["state"] = generic_client_state + elif generic_authorization_endpoint and "okta" in generic_authorization_endpoint: + redirect_params["state"] = uuid.uuid4().hex # set state param for okta - required + + return redirect_params + @staticmethod def should_use_sso_handler( google_client_id: Optional[str] = None, @@ -1025,6 +1098,7 @@ class SSOAuthenticationHandler: def get_redirect_url_for_sso( request: Request, sso_callback_route: str, + existing_key: Optional[str] = None, ) -> str: """ Get the redirect URL for SSO @@ -1036,6 +1110,11 @@ class SSOAuthenticationHandler: redirect_url += sso_callback_route else: redirect_url += "/" + sso_callback_route + + # Append existing_key as query parameter if provided + if existing_key: + redirect_url += f"?existing_key={existing_key}" + return redirect_url @staticmethod @@ -1218,7 +1297,7 @@ class SSOAuthenticationHandler: return team_request @staticmethod - def _get_cli_state(source: Optional[str], key: Optional[str]) -> Optional[str]: + def _get_cli_state(source: Optional[str], key: Optional[str], existing_key: Optional[str] = None) -> Optional[str]: """ Checks the request 'source' if a cli state token was passed in @@ -1229,11 +1308,11 @@ class SSOAuthenticationHandler: LITELLM_CLI_SOURCE_IDENTIFIER, ) - return ( - f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:{key}" - if source == LITELLM_CLI_SOURCE_IDENTIFIER and key - else None - ) + if source == LITELLM_CLI_SOURCE_IDENTIFIER and key: + # Just use the key - existing_key will be passed separately via query params + return f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:{key}" + else: + return None @staticmethod async def get_redirect_response_from_openid( # noqa: PLR0915 diff --git a/tests/test_litellm/proxy/client/cli/test_auth_commands.py b/tests/test_litellm/proxy/client/cli/test_auth_commands.py index 611fde7767..5ef96e4f9a 100644 --- a/tests/test_litellm/proxy/client/cli/test_auth_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_auth_commands.py @@ -435,3 +435,105 @@ class TestWhoamiCommand: assert "✅ Authenticated" in result.output # Should calculate age based on timestamp=0 assert "Token age:" in result.output + + +class TestCLIKeyRegenerationFlow: + """Test the end-to-end CLI key regeneration flow from CLI perspective""" + + def setup_method(self): + """Setup for each test""" + self.runner = CliRunner() + + def test_login_with_existing_key_regeneration_flow(self): + """Test complete login flow when user has existing key - should regenerate it""" + mock_context = Mock() + mock_context.obj = {"base_url": "https://test.example.com"} + + # Mock existing stored key + existing_key = "sk-existing-key-123" + + # Mock successful regeneration response + mock_response = Mock() + mock_response.status_code = 200 + mock_response.json.return_value = { + "status": "ready", + "key": "sk-regenerated-key-456" # New regenerated key + } + + with patch('webbrowser.open') as mock_browser, \ + patch('requests.get', return_value=mock_response) as mock_get, \ + patch('litellm.proxy.client.cli.commands.auth.get_stored_api_key', return_value=existing_key) as mock_get_stored, \ + patch('litellm.proxy.client.cli.commands.auth.save_token') as mock_save, \ + patch('litellm.proxy.client.cli.interface.show_commands') as mock_show_commands, \ + patch('uuid.uuid4', return_value='new-session-uuid-789'): + + result = self.runner.invoke(login, obj=mock_context.obj) + + assert result.exit_code == 0 + assert "✅ Login successful!" in result.output + assert "API Key: sk-regenerated-key-456" in result.output + + # Verify existing key was retrieved + mock_get_stored.assert_called_once() + + # Verify browser was opened with correct URL including existing key + mock_browser.assert_called_once() + call_args = mock_browser.call_args[0][0] + assert "https://test.example.com/sso/key/generate" in call_args + assert "source=litellm-cli" in call_args + assert "key=sk-new-session-uuid-789" in call_args + assert f"existing_key={existing_key}" in call_args + + # Verify polling was done with correct session key + mock_get.assert_called() + poll_url = mock_get.call_args[0][0] + assert "sk-new-session-uuid-789" in poll_url + + # Verify regenerated key was saved + mock_save.assert_called_once() + saved_data = mock_save.call_args[0][0] + assert saved_data['key'] == 'sk-regenerated-key-456' + assert saved_data['user_id'] == 'cli-user' + + mock_show_commands.assert_called_once() + + def test_login_without_existing_key_creation_flow(self): + """Test complete login flow when user has no existing key - should create new one""" + mock_context = Mock() + mock_context.obj = {"base_url": "https://test.example.com"} + + # Mock no existing key + mock_response = Mock() + mock_response.status_code = 200 + mock_response.json.return_value = { + "status": "ready", + "key": "sk-new-created-key-789" + } + + with patch('webbrowser.open') as mock_browser, \ + patch('requests.get', return_value=mock_response), \ + patch('litellm.proxy.client.cli.commands.auth.get_stored_api_key', return_value=None) as mock_get_stored, \ + patch('litellm.proxy.client.cli.commands.auth.save_token') as mock_save, \ + patch('litellm.proxy.client.cli.interface.show_commands'), \ + patch('uuid.uuid4', return_value='new-session-uuid-999'): + + result = self.runner.invoke(login, obj=mock_context.obj) + + assert result.exit_code == 0 + assert "✅ Login successful!" in result.output + + # Verify existing key check was done + mock_get_stored.assert_called_once() + + # Verify browser was opened with correct URL WITHOUT existing key + mock_browser.assert_called_once() + call_args = mock_browser.call_args[0][0] + assert "https://test.example.com/sso/key/generate" in call_args + assert "source=litellm-cli" in call_args + assert "key=sk-new-session-uuid-999" in call_args + assert "existing_key=" not in call_args # Should not include existing_key param + + # Verify new key was saved + mock_save.assert_called_once() + saved_data = mock_save.call_args[0][0] + assert saved_data['key'] == 'sk-new-created-key-789' diff --git a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py index d72e9f7caa..80aebc9849 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/test_litellm/proxy/management_endpoints/test_ui_sso.py @@ -1247,6 +1247,156 @@ class TestCustomUISSO: assert result.status_code == 303 +class TestCLIKeyRegenerationFlow: + """Test the end-to-end CLI key regeneration flow""" + + @pytest.mark.asyncio + async def test_cli_sso_callback_regenerate_existing_key(self): + """Test CLI SSO callback regenerating an existing key""" + from litellm.proxy.management_endpoints.ui_sso import cli_sso_callback + + # Mock request + mock_request = MagicMock(spec=Request) + + # Test data + existing_key = "sk-existing-key-123" + new_key = "sk-new-key-456" + + # Mock the regenerate helper function + with patch("litellm.proxy.management_endpoints.ui_sso._regenerate_cli_key") as mock_regenerate, \ + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), \ + patch("litellm.proxy.common_utils.html_forms.cli_sso_success.render_cli_sso_success_page", return_value="Success"): + + # Act + result = await cli_sso_callback( + request=mock_request, + key=new_key, + existing_key=existing_key + ) + + # Assert + mock_regenerate.assert_called_once_with(existing_key, new_key) + assert result.status_code == 200 + assert "Success" in result.body.decode() + + @pytest.mark.asyncio + async def test_cli_sso_callback_create_new_key(self): + """Test CLI SSO callback creating a new key when no existing key provided""" + from litellm.proxy.management_endpoints.ui_sso import cli_sso_callback + + # Mock request + mock_request = MagicMock(spec=Request) + + # Test data + new_key = "sk-new-key-789" + + # Mock the create helper function + with patch("litellm.proxy.management_endpoints.ui_sso._create_new_cli_key") as mock_create, \ + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), \ + patch("litellm.proxy.common_utils.html_forms.cli_sso_success.render_cli_sso_success_page", return_value="Success"): + + # Act + result = await cli_sso_callback( + request=mock_request, + key=new_key, + existing_key=None + ) + + # Assert + mock_create.assert_called_once_with(new_key) + assert result.status_code == 200 + assert "Success" in result.body.decode() + + @pytest.mark.asyncio + async def test_auth_callback_routes_to_cli_with_existing_key(self): + """Test that auth_callback properly routes CLI requests and preserves existing_key parameter""" + from litellm.constants import LITELLM_CLI_SESSION_TOKEN_PREFIX + from litellm.proxy.management_endpoints.ui_sso import auth_callback + + # Mock request with existing_key query parameter + mock_request = MagicMock(spec=Request) + mock_request.query_params.get.return_value = "sk-existing-cli-key-123" + + # CLI state + cli_state = f"{LITELLM_CLI_SESSION_TOKEN_PREFIX}:sk-new-session-key-456" + + # Mock the CLI callback + with patch("litellm.proxy.management_endpoints.ui_sso.cli_sso_callback") as mock_cli_callback: + mock_cli_callback.return_value = MagicMock() + + # Act + await auth_callback(request=mock_request, state=cli_state) + + # Assert + mock_cli_callback.assert_called_once_with( + mock_request, + key="sk-new-session-key-456", + existing_key="sk-existing-cli-key-123" + ) + + def test_get_redirect_url_preserves_existing_key(self): + """Test that redirect URL generation preserves existing_key parameter""" + from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler + + # Mock request + mock_request = MagicMock() + mock_request.base_url = "https://test.litellm.ai/" + + with patch("litellm.proxy.utils.get_custom_url", return_value="https://test.litellm.ai"): + # Test with existing_key + redirect_url = SSOAuthenticationHandler.get_redirect_url_for_sso( + request=mock_request, + sso_callback_route="sso/callback", + existing_key="sk-existing-123" + ) + + assert "https://test.litellm.ai/sso/callback?existing_key=sk-existing-123" == redirect_url + + def test_get_redirect_url_without_existing_key(self): + """Test that redirect URL generation works without existing_key parameter""" + from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler + + # Mock request + mock_request = MagicMock() + mock_request.base_url = "https://test.litellm.ai/" + + with patch("litellm.proxy.utils.get_custom_url", return_value="https://test.litellm.ai"): + # Test without existing_key + redirect_url = SSOAuthenticationHandler.get_redirect_url_for_sso( + request=mock_request, + sso_callback_route="sso/callback" + ) + + assert "https://test.litellm.ai/sso/callback" == redirect_url + + @pytest.mark.asyncio + async def test_cli_sso_callback_regenerate_vs_create_flow(self): + """Test CLI SSO callback calls regenerate_key_fn when existing_key provided, generate_key_helper_fn when not""" + from litellm.proxy.management_endpoints.ui_sso import cli_sso_callback + + mock_request = MagicMock(spec=Request) + + with patch("litellm.proxy.management_endpoints.key_management_endpoints.regenerate_key_fn") as mock_regenerate, \ + patch("litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn") as mock_generate, \ + patch("litellm.proxy._types.UserAPIKeyAuth.get_litellm_cli_user_api_key_auth"), \ + patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), \ + patch("litellm.proxy.common_utils.html_forms.cli_sso_success.render_cli_sso_success_page", return_value="Success"): + + # Test regeneration path + await cli_sso_callback(mock_request, key="sk-new-123", existing_key="sk-existing-456") + mock_regenerate.assert_called_once() + mock_generate.assert_not_called() + + # Reset mocks + mock_regenerate.reset_mock() + mock_generate.reset_mock() + + # Test creation path + await cli_sso_callback(mock_request, key="sk-new-789", existing_key=None) + mock_regenerate.assert_not_called() + mock_generate.assert_called_once() + + class TestProcessSSOJWTAccessToken: """Test the process_sso_jwt_access_token helper function"""