diff --git a/enterprise/storage/lite_llm_manager.py b/enterprise/storage/lite_llm_manager.py index f1179ebcbc..0ac1ea7a66 100644 --- a/enterprise/storage/lite_llm_manager.py +++ b/enterprise/storage/lite_llm_manager.py @@ -346,8 +346,16 @@ class LiteLlmManager: # User failed to create in litellm - this is an unforseen error state... if not response.is_success: - if response.status_code == 400 and 'already exists' in response.text: - # user already exists, just return + if ( + response.status_code in (400, 409) + and 'already exists' in response.text + ): + logger.warning( + 'litellm_user_already_exists', + extra={ + 'user_id': keycloak_user_id, + }, + ) return logger.error( 'error_creating_litellm_user', @@ -478,6 +486,18 @@ class LiteLlmManager: ) # Failed to add user to team - this is an unforseen error state... if not response.is_success: + if ( + response.status_code == 400 + and 'already in team' in response.text.lower() + ): + logger.warning( + 'user_already_in_team', + extra={ + 'user_id': keycloak_user_id, + 'team_id': team_id, + }, + ) + return logger.error( 'error_adding_litellm_user_to_team', extra={ diff --git a/enterprise/tests/unit/test_lite_llm_manager.py b/enterprise/tests/unit/test_lite_llm_manager.py index 0ec8fe6b65..d148882d8a 100644 --- a/enterprise/tests/unit/test_lite_llm_manager.py +++ b/enterprise/tests/unit/test_lite_llm_manager.py @@ -465,6 +465,149 @@ class TestLiteLlmManager: 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.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.""" + # 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' + + mock_http_client.post.side_effect = [first_response, second_response] + + # Act + await LiteLlmManager._create_user( + mock_http_client, 'test@example.com', 'test-user-id' + ) + + # Assert + 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_with_400_status_code( + self, mock_logger, mock_http_client + ): + """Test _create_user handles 400 Bad Request when user already exists.""" + # 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 = 400 + second_response.text = 'User already exists' + + mock_http_client.post.side_effect = [first_response, second_response] + + # Act + await LiteLlmManager._create_user( + mock_http_client, 'test@example.com', 'test-user-id' + ) + + # Assert + mock_logger.warning.assert_any_call( + 'litellm_user_already_exists', + extra={'user_id': 'test-user-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_add_user_to_team_success(self, mock_http_client, mock_response): + """Test successful _add_user_to_team operation.""" + # Arrange + mock_http_client.post.return_value = mock_response + + # Act + await LiteLlmManager._add_user_to_team( + mock_http_client, 'test-user-id', 'test-team-id', 100.0 + ) + + # Assert + mock_http_client.post.assert_called_once() + call_args = mock_http_client.post.call_args + assert 'http://test.com/team/member_add' in call_args[0] + assert call_args[1]['json']['team_id'] == 'test-team-id' + assert call_args[1]['json']['member'] == { + 'user_id': 'test-user-id', + 'role': 'user', + } + assert call_args[1]['json']['max_budget_in_team'] == 100.0 + + @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_add_user_to_team_already_in_team( + self, mock_logger, mock_http_client + ): + """Test _add_user_to_team handles 'already in team' error gracefully.""" + # Arrange + error_response = MagicMock() + error_response.is_success = False + error_response.status_code = 400 + error_response.text = ( + '{"error":{"message":"User already in team. Member: ' + 'user_id=test-user-id","type":"team_member_already_in_team"}}' + ) + mock_http_client.post.return_value = error_response + + # Act + await LiteLlmManager._add_user_to_team( + mock_http_client, 'test-user-id', 'test-team-id', 100.0 + ) + + # Assert + mock_logger.warning.assert_called_once_with( + 'user_already_in_team', + extra={ + 'user_id': 'test-user-id', + 'team_id': 'test-team-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_add_user_to_team_other_error_raises_exception( + self, mock_http_client + ): + """Test _add_user_to_team raises exception for non-'already in team' errors.""" + # Arrange + error_response = MagicMock() + error_response.is_success = False + error_response.status_code = 500 + error_response.text = 'Internal server error' + error_response.raise_for_status.side_effect = httpx.HTTPStatusError( + 'Server error', request=MagicMock(), response=error_response + ) + mock_http_client.post.return_value = error_response + + # Act & Assert + with pytest.raises(httpx.HTTPStatusError): + await LiteLlmManager._add_user_to_team( + mock_http_client, 'test-user-id', 'test-team-id', 100.0 + ) + @pytest.mark.asyncio async def test_generate_key_success(self, mock_http_client, mock_response): """Test successful _generate_key operation."""