diff --git a/Source/Editor/Modules/SceneModule.cs b/Source/Editor/Modules/SceneModule.cs
index e41612a01..3a7af939b 100644
--- a/Source/Editor/Modules/SceneModule.cs
+++ b/Source/Editor/Modules/SceneModule.cs
@@ -359,9 +359,7 @@ namespace FlaxEditor.Modules
         {
             Editor.SceneEditing.Deselect();
 
-            // TODO: this works in most cases fine but we still need to handle the case when assembly gets reloaded and type references are invalid
-            // TODO: To solve this: serialize type reference as 'namespace.name' in undo data storage
-            //if (fullCleanup)
+            if (fullCleanup)
             {
                 Undo.Clear();
             }
diff --git a/Source/Editor/Progress/Handlers/CompileScriptsProgress.cs b/Source/Editor/Progress/Handlers/CompileScriptsProgress.cs
index 3661e89df..e48b9c1c3 100644
--- a/Source/Editor/Progress/Handlers/CompileScriptsProgress.cs
+++ b/Source/Editor/Progress/Handlers/CompileScriptsProgress.cs
@@ -36,7 +36,7 @@ namespace FlaxEditor.Progress.Handlers
             _selectionCache.Cache();
 
             // Clear references to the user scripts (we gonna reload an assembly)
-            Editor.Instance.Scene.ClearRefsToSceneObjects();
+            Editor.Instance.Scene.ClearRefsToSceneObjects(true);
         }
 
         private void OnScriptsReload()
diff --git a/Source/Editor/States/PlayingState.cs b/Source/Editor/States/PlayingState.cs
index 4c9a566fd..c5d1eee0d 100644
--- a/Source/Editor/States/PlayingState.cs
+++ b/Source/Editor/States/PlayingState.cs
@@ -119,7 +119,7 @@ namespace FlaxEditor.States
             CacheSelection();
 
             // Remove references to the scene objects
-            Editor.Scene.ClearRefsToSceneObjects(true);
+            Editor.Scene.ClearRefsToSceneObjects();
 
             // Apply game settings (user may modify them before the gameplay)
             GameSettingsApplying?.Invoke();
@@ -158,7 +158,7 @@ namespace FlaxEditor.States
             IsPaused = true;
 
             // Remove references to the scene objects
-            Editor.Scene.ClearRefsToSceneObjects(true);
+            Editor.Scene.ClearRefsToSceneObjects();
 
             // Restore editor scene
             SceneRestoring?.Invoke();