diff --git a/enterprise/storage/lite_llm_manager.py b/enterprise/storage/lite_llm_manager.py index 725b8147a3..b515b7a7d9 100644 --- a/enterprise/storage/lite_llm_manager.py +++ b/enterprise/storage/lite_llm_manager.py @@ -164,9 +164,33 @@ class LiteLlmManager: ) if create_user: - await LiteLlmManager._create_user( + user_created = await LiteLlmManager._create_user( client, keycloak_user_info.get('email'), keycloak_user_id ) + if not user_created: + logger.error( + 'create_entries_failed_user_creation', + extra={ + 'org_id': org_id, + 'user_id': keycloak_user_id, + }, + ) + return None + + # Verify user exists before proceeding with key generation + user_exists = await LiteLlmManager._user_exists( + client, keycloak_user_id + ) + if not user_exists: + logger.error( + 'create_entries_user_not_found_before_key_generation', + extra={ + 'org_id': org_id, + 'user_id': keycloak_user_id, + 'create_user_flag': create_user, + }, + ) + return None await LiteLlmManager._add_user_to_team( client, keycloak_user_id, org_id, team_budget @@ -655,15 +679,48 @@ class LiteLlmManager: ) response.raise_for_status() + @staticmethod + async def _user_exists( + client: httpx.AsyncClient, + user_id: str, + ) -> bool: + """Check if a user exists in LiteLLM. + + Returns True if the user exists, False otherwise. + """ + if LITE_LLM_API_KEY is None or LITE_LLM_API_URL is None: + return False + try: + response = await client.get( + f'{LITE_LLM_API_URL}/user/info?user_id={user_id}', + ) + if response.is_success: + user_data = response.json() + # Check that user_info exists and has the user_id + user_info = user_data.get('user_info', {}) + return user_info.get('user_id') == user_id + return False + except Exception as e: + logger.warning( + 'litellm_user_exists_check_failed', + extra={'user_id': user_id, 'error': str(e)}, + ) + return False + @staticmethod async def _create_user( client: httpx.AsyncClient, email: str | None, keycloak_user_id: str, - ): + ) -> bool: + """Create a user in LiteLLM. + + Returns True if the user was created or already exists and is verified, + False if creation failed and user does not exist. + """ if LITE_LLM_API_KEY is None or LITE_LLM_API_URL is None: logger.warning('LiteLLM API configuration not found') - return + return False response = await client.post( f'{LITE_LLM_API_URL}/user/new', json={ @@ -716,17 +773,33 @@ class LiteLlmManager: 'user_id': keycloak_user_id, }, ) - return + # Verify the user actually exists before returning success + user_exists = await LiteLlmManager._user_exists( + client, keycloak_user_id + ) + if not user_exists: + logger.error( + 'litellm_user_claimed_exists_but_not_found', + extra={ + 'user_id': keycloak_user_id, + 'status_code': response.status_code, + 'text': response.text, + }, + ) + return False + return True logger.error( 'error_creating_litellm_user', extra={ 'status_code': response.status_code, 'text': response.text, - 'user_id': [keycloak_user_id], + 'user_id': keycloak_user_id, 'email': None, }, ) + return False response.raise_for_status() + return True @staticmethod async def _get_user(client: httpx.AsyncClient, user_id: str) -> dict | None: @@ -1450,6 +1523,7 @@ class LiteLlmManager: create_team = staticmethod(with_http_client(_create_team)) get_team = staticmethod(with_http_client(_get_team)) update_team = staticmethod(with_http_client(_update_team)) + user_exists = staticmethod(with_http_client(_user_exists)) create_user = staticmethod(with_http_client(_create_user)) get_user = staticmethod(with_http_client(_get_user)) update_user = staticmethod(with_http_client(_update_user)) diff --git a/enterprise/tests/unit/test_lite_llm_manager.py b/enterprise/tests/unit/test_lite_llm_manager.py index 0cfc9fe58b..3da159421d 100644 --- a/enterprise/tests/unit/test_lite_llm_manager.py +++ b/enterprise/tests/unit/test_lite_llm_manager.py @@ -239,6 +239,16 @@ class TestLiteLlmManager: mock_404_response = MagicMock() mock_404_response.status_code = 404 mock_404_response.is_success = False + mock_404_response.raise_for_status.side_effect = httpx.HTTPStatusError( + message='Not Found', request=MagicMock(), response=mock_404_response + ) + + # Mock user exists check response + mock_user_exists_response = MagicMock() + mock_user_exists_response.is_success = True + mock_user_exists_response.json.return_value = { + 'user_info': {'user_id': 'test-user-id'} + } mock_token_manager = MagicMock() mock_token_manager.return_value.get_user_info_from_user_id = AsyncMock( @@ -246,12 +256,8 @@ class TestLiteLlmManager: ) mock_client = AsyncMock() - mock_client.get.return_value = mock_404_response - mock_client.get.return_value.raise_for_status.side_effect = ( - httpx.HTTPStatusError( - message='Not Found', request=MagicMock(), response=mock_404_response - ) - ) + # First GET is for _get_team (404), second GET is for _user_exists (success) + mock_client.get.side_effect = [mock_404_response, mock_user_exists_response] mock_client.post.return_value = mock_response mock_client_class = MagicMock() @@ -274,8 +280,8 @@ class TestLiteLlmManager: assert result.llm_api_key.get_secret_value() == 'test-api-key' assert result.llm_base_url == 'http://test.com' - # Verify API calls were made (get_team + 4 posts) - assert mock_client.get.call_count == 1 # get_team + # Verify API calls were made (get_team + user_exists + 4 posts) + assert mock_client.get.call_count == 2 # get_team + user_exists assert ( mock_client.post.call_count == 4 ) # create_team, add_user_to_team, delete_key_by_alias, generate_key @@ -294,13 +300,21 @@ class TestLiteLlmManager: } mock_team_response.raise_for_status = MagicMock() + # Mock user exists check response + mock_user_exists_response = MagicMock() + mock_user_exists_response.is_success = True + mock_user_exists_response.json.return_value = { + 'user_info': {'user_id': 'test-user-id'} + } + mock_token_manager = MagicMock() mock_token_manager.return_value.get_user_info_from_user_id = AsyncMock( return_value={'email': 'test@example.com'} ) mock_client = AsyncMock() - mock_client.get.return_value = mock_team_response + # First GET is for _get_team (success), second GET is for _user_exists (success) + mock_client.get.side_effect = [mock_team_response, mock_user_exists_response] mock_client.post.return_value = mock_response mock_client_class = MagicMock() @@ -320,8 +334,8 @@ class TestLiteLlmManager: assert result is not None # Verify _get_team was called first - mock_client.get.assert_called_once() - get_call_url = mock_client.get.call_args[0][0] + assert mock_client.get.call_count == 2 # get_team + user_exists + get_call_url = mock_client.get.call_args_list[0][0][0] assert 'team/info' in get_call_url assert 'test-org-id' in get_call_url @@ -343,19 +357,25 @@ class TestLiteLlmManager: mock_404_response = MagicMock() mock_404_response.status_code = 404 mock_404_response.is_success = False + mock_404_response.raise_for_status.side_effect = httpx.HTTPStatusError( + message='Not Found', request=MagicMock(), response=mock_404_response + ) mock_token_manager = MagicMock() mock_token_manager.return_value.get_user_info_from_user_id = AsyncMock( return_value={'email': 'test@example.com'} ) + # Mock user exists check response + mock_user_exists_response = MagicMock() + mock_user_exists_response.is_success = True + mock_user_exists_response.json.return_value = { + 'user_info': {'user_id': 'test-user-id'} + } + mock_client = AsyncMock() - mock_client.get.return_value = mock_404_response - mock_client.get.return_value.raise_for_status.side_effect = ( - httpx.HTTPStatusError( - message='Not Found', request=MagicMock(), response=mock_404_response - ) - ) + # First GET is for _get_team (404), second GET is for _user_exists (success) + mock_client.get.side_effect = [mock_404_response, mock_user_exists_response] mock_client.post.return_value = mock_response mock_client_class = MagicMock() @@ -393,6 +413,16 @@ class TestLiteLlmManager: mock_404_response = MagicMock() mock_404_response.status_code = 404 mock_404_response.is_success = False + mock_404_response.raise_for_status.side_effect = httpx.HTTPStatusError( + message='Not Found', request=MagicMock(), response=mock_404_response + ) + + # Mock user exists check response + mock_user_exists_response = MagicMock() + mock_user_exists_response.is_success = True + mock_user_exists_response.json.return_value = { + 'user_info': {'user_id': 'test-user-id'} + } mock_token_manager = MagicMock() mock_token_manager.return_value.get_user_info_from_user_id = AsyncMock( @@ -400,12 +430,8 @@ class TestLiteLlmManager: ) mock_client = AsyncMock() - mock_client.get.return_value = mock_404_response - mock_client.get.return_value.raise_for_status.side_effect = ( - httpx.HTTPStatusError( - message='Not Found', request=MagicMock(), response=mock_404_response - ) - ) + # First GET is for _get_team (404), second GET is for _user_exists (success) + mock_client.get.side_effect = [mock_404_response, mock_user_exists_response] mock_client.post.return_value = mock_response mock_client_class = MagicMock() @@ -833,15 +859,16 @@ class TestLiteLlmManager: @pytest.mark.asyncio async def test_create_user_success(self, mock_http_client, mock_response): - """Test successful _create_user operation.""" + """Test successful _create_user operation returns True.""" mock_http_client.post.return_value = mock_response with patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key'): with patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com'): - await LiteLlmManager._create_user( + result = await LiteLlmManager._create_user( mock_http_client, 'test@example.com', 'test-user-id' ) + assert result is True mock_http_client.post.assert_called_once() call_args = mock_http_client.post.call_args assert 'http://test.com/user/new' in call_args[0] @@ -850,7 +877,7 @@ class TestLiteLlmManager: @pytest.mark.asyncio async def test_create_user_duplicate_email(self, mock_http_client, mock_response): - """Test _create_user with duplicate email handling.""" + """Test _create_user with duplicate email handling returns True.""" # First call fails with duplicate email error_response = MagicMock() error_response.is_success = False @@ -862,23 +889,81 @@ class TestLiteLlmManager: with patch('storage.lite_llm_manager.LITE_LLM_API_KEY', 'test-key'): with patch('storage.lite_llm_manager.LITE_LLM_API_URL', 'http://test.com'): - await LiteLlmManager._create_user( + result = await LiteLlmManager._create_user( mock_http_client, 'test@example.com', 'test-user-id' ) + assert result is True assert mock_http_client.post.call_count == 2 # Second call should have None email second_call_args = mock_http_client.post.call_args_list[1] assert second_call_args[1]['json']['user_email'] is None + @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_user_exists_returns_true(self, mock_http_client): + """Test _user_exists returns True when user exists in LiteLLM.""" + # Arrange + user_response = MagicMock() + user_response.is_success = True + user_response.json.return_value = { + 'user_info': {'user_id': 'test-user-id', 'email': 'test@example.com'} + } + mock_http_client.get.return_value = user_response + + # Act + result = await LiteLlmManager._user_exists(mock_http_client, 'test-user-id') + + # Assert + assert result is True + mock_http_client.get.assert_called_once() + + @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_user_exists_returns_false_when_not_found(self, mock_http_client): + """Test _user_exists returns False when user not found.""" + # Arrange + user_response = MagicMock() + user_response.is_success = False + mock_http_client.get.return_value = user_response + + # Act + result = await LiteLlmManager._user_exists(mock_http_client, 'test-user-id') + + # Assert + assert result is False + + @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_user_exists_returns_false_on_mismatched_user_id( + self, mock_http_client + ): + """Test _user_exists returns False when returned user_id doesn't match.""" + # Arrange + user_response = MagicMock() + user_response.is_success = True + user_response.json.return_value = { + 'user_info': {'user_id': 'different-user-id'} + } + mock_http_client.get.return_value = user_response + + # Act + result = await LiteLlmManager._user_exists(mock_http_client, 'test-user-id') + + # Assert + assert result is False + @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_create_user_already_exists_with_409_status_code( + async def test_create_user_already_exists_and_verified( self, mock_logger, mock_http_client ): - """Test _create_user handles 409 Conflict when user already exists.""" + """Test _create_user returns True when user already exists and is verified.""" # Arrange first_response = MagicMock() first_response.is_success = False @@ -890,14 +975,141 @@ class TestLiteLlmManager: second_response.status_code = 409 second_response.text = 'User with id test-user-id already exists' + user_exists_response = MagicMock() + user_exists_response.is_success = True + user_exists_response.json.return_value = { + 'user_info': {'user_id': 'test-user-id'} + } + mock_http_client.post.side_effect = [first_response, second_response] + mock_http_client.get.return_value = user_exists_response # Act - await LiteLlmManager._create_user( + result = await LiteLlmManager._create_user( mock_http_client, 'test@example.com', 'test-user-id' ) # Assert + assert result is True + mock_logger.warning.assert_any_call( + 'litellm_user_already_exists', + extra={'user_id': 'test-user-id'}, + ) + + @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_create_user_already_exists_but_not_found_returns_false( + self, mock_logger, mock_http_client + ): + """Test _create_user returns False when LiteLLM claims user exists but verification fails.""" + # Arrange + first_response = MagicMock() + first_response.is_success = False + first_response.status_code = 400 + first_response.text = 'duplicate email' + + second_response = MagicMock() + second_response.is_success = False + second_response.status_code = 409 + second_response.text = 'User with id test-user-id already exists' + + user_not_exists_response = MagicMock() + user_not_exists_response.is_success = False + + mock_http_client.post.side_effect = [first_response, second_response] + mock_http_client.get.return_value = user_not_exists_response + + # Act + result = await LiteLlmManager._create_user( + mock_http_client, 'test@example.com', 'test-user-id' + ) + + # Assert + assert result is False + mock_logger.error.assert_any_call( + 'litellm_user_claimed_exists_but_not_found', + extra={ + 'user_id': 'test-user-id', + 'status_code': 409, + 'text': 'User with id test-user-id already exists', + }, + ) + + @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_create_user_failure_returns_false( + self, mock_logger, mock_http_client + ): + """Test _create_user returns False when creation fails with non-'already exists' error.""" + # Arrange + first_response = MagicMock() + first_response.is_success = False + first_response.status_code = 400 + first_response.text = 'duplicate email' + + second_response = MagicMock() + second_response.is_success = False + second_response.status_code = 500 + second_response.text = 'Internal server error' + + mock_http_client.post.side_effect = [first_response, second_response] + + # Act + result = await LiteLlmManager._create_user( + mock_http_client, 'test@example.com', 'test-user-id' + ) + + # Assert + assert result is False + mock_logger.error.assert_any_call( + 'error_creating_litellm_user', + extra={ + 'status_code': 500, + 'text': 'Internal server error', + 'user_id': 'test-user-id', + 'email': None, + }, + ) + + @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_create_user_already_exists_with_409_status_code( + self, mock_logger, mock_http_client + ): + """Test _create_user handles 409 Conflict when user already exists and verifies.""" + # Arrange + first_response = MagicMock() + first_response.is_success = False + first_response.status_code = 400 + first_response.text = 'duplicate email' + + second_response = MagicMock() + second_response.is_success = False + second_response.status_code = 409 + second_response.text = 'User with id test-user-id already exists' + + user_exists_response = MagicMock() + user_exists_response.is_success = True + user_exists_response.json.return_value = { + 'user_info': {'user_id': 'test-user-id'} + } + + mock_http_client.post.side_effect = [first_response, second_response] + mock_http_client.get.return_value = user_exists_response + + # Act + result = await LiteLlmManager._create_user( + mock_http_client, 'test@example.com', 'test-user-id' + ) + + # Assert + assert result is True mock_logger.warning.assert_any_call( 'litellm_user_already_exists', extra={'user_id': 'test-user-id'}, @@ -910,7 +1122,7 @@ class TestLiteLlmManager: async def test_create_user_already_exists_with_400_status_code( self, mock_logger, mock_http_client ): - """Test _create_user handles 400 Bad Request when user already exists.""" + """Test _create_user handles 400 Bad Request when user already exists and verifies.""" # Arrange first_response = MagicMock() first_response.is_success = False @@ -922,14 +1134,22 @@ class TestLiteLlmManager: second_response.status_code = 400 second_response.text = 'User already exists' + user_exists_response = MagicMock() + user_exists_response.is_success = True + user_exists_response.json.return_value = { + 'user_info': {'user_id': 'test-user-id'} + } + mock_http_client.post.side_effect = [first_response, second_response] + mock_http_client.get.return_value = user_exists_response # Act - await LiteLlmManager._create_user( + result = await LiteLlmManager._create_user( mock_http_client, 'test@example.com', 'test-user-id' ) # Assert + assert result is True mock_logger.warning.assert_any_call( 'litellm_user_already_exists', extra={'user_id': 'test-user-id'},