fix: ensure LiteLLM user exists before generating API keys (#12667)

Co-authored-by: openhands <openhands@all-hands.dev>

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))

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
+        # First GET is for _get_team (404), second GET is for _user_exists (success)
-        mock_client.get.return_value.raise_for_status.side_effect = (
+        mock_client.get.side_effect = [mock_404_response, mock_user_exists_response]
-            httpx.HTTPStatusError(
-                message='Not Found', request=MagicMock(), response=mock_404_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)
+            # Verify API calls were made (get_team + user_exists + 4 posts)
-            assert mock_client.get.call_count == 1  # get_team
+            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()
+            assert mock_client.get.call_count == 2  # get_team + user_exists
-            get_call_url = mock_client.get.call_args[0][0]
+            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
+        # First GET is for _get_team (404), second GET is for _user_exists (success)
-        mock_client.get.return_value.raise_for_status.side_effect = (
+        mock_client.get.side_effect = [mock_404_response, mock_user_exists_response]
-            httpx.HTTPStatusError(
-                message='Not Found', request=MagicMock(), response=mock_404_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
+        # First GET is for _get_team (404), second GET is for _user_exists (success)
-        mock_client.get.return_value.raise_for_status.side_effect = (
+        mock_client.get.side_effect = [mock_404_response, mock_user_exists_response]
-            httpx.HTTPStatusError(
-                message='Not Found', request=MagicMock(), response=mock_404_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'},