diff --git a/enterprise/server/routes/api_keys.py b/enterprise/server/routes/api_keys.py index 3586cbb094..c07ae3fc51 100644 --- a/enterprise/server/routes/api_keys.py +++ b/enterprise/server/routes/api_keys.py @@ -99,10 +99,19 @@ async def generate_byor_key(user_id: str) -> str | None: async def delete_byor_key_from_litellm(user_id: str, byor_key: str) -> bool: - """Delete the BYOR key from LiteLLM using the key directly.""" + """Delete the BYOR key from LiteLLM using the key directly. + Also attempts to delete by key alias if the key is not found, + to clean up orphaned aliases that could block key regeneration. + """ try: - await LiteLlmManager.delete_key(byor_key) + # Get user to construct the key alias + user = await UserStore.get_user_by_id_async(user_id) + key_alias = None + if user and user.current_org_id: + key_alias = f'BYOR Key - user {user_id}, org {user.current_org_id}' + + await LiteLlmManager.delete_key(byor_key, key_alias=key_alias) logger.info( 'Successfully deleted BYOR key from LiteLLM', extra={'user_id': user_id}, diff --git a/enterprise/storage/lite_llm_manager.py b/enterprise/storage/lite_llm_manager.py index 439ff45f94..f1179ebcbc 100644 --- a/enterprise/storage/lite_llm_manager.py +++ b/enterprise/storage/lite_llm_manager.py @@ -699,10 +699,45 @@ class LiteLlmManager: 'key_spend': key_info.get('spend'), } + @staticmethod + async def _delete_key_by_alias( + client: httpx.AsyncClient, + key_alias: str, + ): + """Delete a key from LiteLLM by its alias. + + This is a best-effort operation that logs but does not raise on failure. + """ + if LITE_LLM_API_KEY is None or LITE_LLM_API_URL is None: + logger.warning('LiteLLM API configuration not found') + return + response = await client.post( + f'{LITE_LLM_API_URL}/key/delete', + json={ + 'key_aliases': [key_alias], + }, + ) + if response.is_success: + logger.info( + 'LiteLlmManager:_delete_key_by_alias:key_deleted', + extra={'key_alias': key_alias}, + ) + elif response.status_code != 404: + # Log non-404 errors but don't fail + logger.warning( + 'error_deleting_key_by_alias', + extra={ + 'key_alias': key_alias, + 'status_code': response.status_code, + 'text': response.text, + }, + ) + @staticmethod async def _delete_key( client: httpx.AsyncClient, key_id: str, + key_alias: str | None = None, ): if LITE_LLM_API_KEY is None or LITE_LLM_API_URL is None: logger.warning('LiteLLM API configuration not found') @@ -713,10 +748,13 @@ class LiteLlmManager: 'keys': [key_id], }, ) - # Failed to key... + # Failed to delete key... if not response.is_success: if response.status_code == 404: - # key doesn't exist, just return + # Key doesn't exist by key_id. If we have a key_alias, + # try deleting by alias to clean up any orphaned alias. + if key_alias: + await LiteLlmManager._delete_key_by_alias(client, key_alias) return logger.error( 'error_deleting_key', diff --git a/enterprise/tests/unit/server/routes/test_api_keys.py b/enterprise/tests/unit/server/routes/test_api_keys.py index 08e92bd56b..7f08425758 100644 --- a/enterprise/tests/unit/server/routes/test_api_keys.py +++ b/enterprise/tests/unit/server/routes/test_api_keys.py @@ -6,6 +6,7 @@ import httpx import pytest from fastapi import HTTPException from server.routes.api_keys import ( + delete_byor_key_from_litellm, get_llm_api_key_for_byor, ) from storage.lite_llm_manager import LiteLlmManager @@ -328,3 +329,99 @@ class TestGetLlmApiKeyForByor: assert exc_info.value.status_code == 500 assert 'Failed to retrieve BYOR LLM API key' in exc_info.value.detail + + +class TestDeleteByorKeyFromLitellm: + """Test the delete_byor_key_from_litellm function with alias cleanup.""" + + @pytest.mark.asyncio + @patch('storage.lite_llm_manager.LiteLlmManager.delete_key') + @patch('storage.user_store.UserStore.get_user_by_id_async') + async def test_delete_constructs_alias_from_user( + self, mock_get_user, mock_delete_key + ): + """Test that delete_byor_key_from_litellm constructs key alias from user.""" + # Arrange + user_id = 'user-123' + org_id = 'org-456' + byor_key = 'sk-byor-key-to-delete' + expected_alias = f'BYOR Key - user {user_id}, org {org_id}' + + mock_user = MagicMock() + mock_user.current_org_id = org_id + mock_get_user.return_value = mock_user + mock_delete_key.return_value = None + + # Act + result = await delete_byor_key_from_litellm(user_id, byor_key) + + # Assert + assert result is True + mock_get_user.assert_called_once_with(user_id) + mock_delete_key.assert_called_once_with(byor_key, key_alias=expected_alias) + + @pytest.mark.asyncio + @patch('storage.lite_llm_manager.LiteLlmManager.delete_key') + @patch('storage.user_store.UserStore.get_user_by_id_async') + async def test_delete_without_user_passes_no_alias( + self, mock_get_user, mock_delete_key + ): + """Test that when user is not found, no alias is passed.""" + # Arrange + user_id = 'user-123' + byor_key = 'sk-byor-key-to-delete' + + mock_get_user.return_value = None + mock_delete_key.return_value = None + + # Act + result = await delete_byor_key_from_litellm(user_id, byor_key) + + # Assert + assert result is True + mock_delete_key.assert_called_once_with(byor_key, key_alias=None) + + @pytest.mark.asyncio + @patch('storage.lite_llm_manager.LiteLlmManager.delete_key') + @patch('storage.user_store.UserStore.get_user_by_id_async') + async def test_delete_without_org_id_passes_no_alias( + self, mock_get_user, mock_delete_key + ): + """Test that when user has no current_org_id, no alias is passed.""" + # Arrange + user_id = 'user-123' + byor_key = 'sk-byor-key-to-delete' + + mock_user = MagicMock() + mock_user.current_org_id = None + mock_get_user.return_value = mock_user + mock_delete_key.return_value = None + + # Act + result = await delete_byor_key_from_litellm(user_id, byor_key) + + # Assert + assert result is True + mock_delete_key.assert_called_once_with(byor_key, key_alias=None) + + @pytest.mark.asyncio + @patch('storage.lite_llm_manager.LiteLlmManager.delete_key') + @patch('storage.user_store.UserStore.get_user_by_id_async') + async def test_delete_returns_false_on_exception( + self, mock_get_user, mock_delete_key + ): + """Test that exceptions during deletion return False.""" + # Arrange + user_id = 'user-123' + byor_key = 'sk-byor-key-to-delete' + + mock_user = MagicMock() + mock_user.current_org_id = 'org-456' + mock_get_user.return_value = mock_user + mock_delete_key.side_effect = Exception('LiteLLM API error') + + # Act + result = await delete_byor_key_from_litellm(user_id, byor_key) + + # Assert + assert result is False diff --git a/enterprise/tests/unit/test_lite_llm_manager.py b/enterprise/tests/unit/test_lite_llm_manager.py index 5ed21d2929..0ec8fe6b65 100644 --- a/enterprise/tests/unit/test_lite_llm_manager.py +++ b/enterprise/tests/unit/test_lite_llm_manager.py @@ -557,6 +557,131 @@ class TestLiteLlmManager: # Should not raise an exception for 404 await LiteLlmManager._delete_key(mock_http_client, 'test-key-id') + @pytest.mark.asyncio + @patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com') + @patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key') + async def test_delete_key_not_found_with_alias_triggers_alias_deletion( + self, mock_http_client + ): + """Test _delete_key falls back to alias deletion when key_id returns 404.""" + # Arrange + not_found_response = MagicMock() + not_found_response.is_success = False + not_found_response.status_code = 404 + not_found_response.text = 'Key not found' + + alias_success_response = MagicMock() + alias_success_response.is_success = True + alias_success_response.status_code = 200 + + mock_http_client.post.side_effect = [not_found_response, alias_success_response] + + # Act + await LiteLlmManager._delete_key( + mock_http_client, 'test-key-id', key_alias='BYOR Key - user 123, org 456' + ) + + # Assert + assert mock_http_client.post.call_count == 2 + first_call = mock_http_client.post.call_args_list[0] + assert first_call[1]['json']['keys'] == ['test-key-id'] + second_call = mock_http_client.post.call_args_list[1] + assert second_call[1]['json']['key_aliases'] == ['BYOR Key - user 123, org 456'] + + @pytest.mark.asyncio + @patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com') + @patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key') + async def test_delete_key_not_found_without_alias_no_fallback( + self, mock_http_client + ): + """Test _delete_key without alias does not attempt alias deletion on 404.""" + # Arrange + not_found_response = MagicMock() + not_found_response.is_success = False + not_found_response.status_code = 404 + not_found_response.text = 'Key not found' + mock_http_client.post.return_value = not_found_response + + # Act + await LiteLlmManager._delete_key(mock_http_client, 'test-key-id') + + # Assert + assert mock_http_client.post.call_count == 1 + + @pytest.mark.asyncio + @patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com') + @patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key') + async def test_delete_key_by_alias_success(self, mock_http_client, mock_response): + """Test successful _delete_key_by_alias operation.""" + # Arrange + mock_http_client.post.return_value = mock_response + + # Act + await LiteLlmManager._delete_key_by_alias( + mock_http_client, 'BYOR Key - user 123, org 456' + ) + + # Assert + mock_http_client.post.assert_called_once() + call_args = mock_http_client.post.call_args + assert 'http://test.com/key/delete' in call_args[0] + assert call_args[1]['json']['key_aliases'] == ['BYOR Key - user 123, org 456'] + + @pytest.mark.asyncio + @patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com') + @patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key') + async def test_delete_key_by_alias_not_found(self, mock_http_client): + """Test _delete_key_by_alias when alias is not found (404).""" + # Arrange + not_found_response = MagicMock() + not_found_response.is_success = False + not_found_response.status_code = 404 + not_found_response.text = 'Key alias not found' + mock_http_client.post.return_value = not_found_response + + # Act & Assert - should not raise exception for 404 + await LiteLlmManager._delete_key_by_alias( + mock_http_client, 'BYOR Key - user 123, org 456' + ) + + @pytest.mark.asyncio + @patch('storage.lite_llm_manager.logger') + @patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com') + @patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key') + async def test_delete_key_by_alias_server_error_logs_warning( + self, mock_logger, mock_http_client + ): + """Test _delete_key_by_alias logs warning for non-404 errors.""" + # Arrange + error_response = MagicMock() + error_response.is_success = False + error_response.status_code = 500 + error_response.text = 'Internal server error' + mock_http_client.post.return_value = error_response + + # Act + await LiteLlmManager._delete_key_by_alias( + mock_http_client, 'BYOR Key - user 123, org 456' + ) + + # Assert + mock_logger.warning.assert_called_once() + call_args = mock_logger.warning.call_args + assert call_args[0][0] == 'error_deleting_key_by_alias' + + @pytest.mark.asyncio + @patch('storage.lite_llm_manager.LITE_LLM_API_URL', None) + @patch('storage.lite_llm_manager.LITE_LLM_API_KEY', None) + async def test_delete_key_by_alias_missing_config(self, mock_http_client): + """Test _delete_key_by_alias returns early when config is missing.""" + # Act + await LiteLlmManager._delete_key_by_alias( + mock_http_client, 'BYOR Key - user 123, org 456' + ) + + # Assert + mock_http_client.post.assert_not_called() + @pytest.mark.asyncio async def test_with_http_client_decorator(self): """Test the with_http_client decorator functionality."""