Merge pull request #1056 from jacklaplante/codex/fix-prefab-stage-save

fix: save prefab stage edits using the correct prefab-stage workflow

Author: Scriptwonder

MCPForUnity/Editor/Tools/Prefabs/ManagePrefabs.cs

@@ -1348,15 +1348,12 @@ private static object SavePrefabStage()
                     return new ErrorResponse("Not currently in prefab editing mode. Open a prefab stage first with open_prefab_stage.");
                 }
 
-                string prefabPath = prefabStage.assetPath;
-                EditorSceneManager.MarkSceneDirty(prefabStage.scene);
-                bool saved = EditorSceneManager.SaveScene(prefabStage.scene);
-                if (!saved)
+                if (!TrySavePrefabStage(prefabStage, out string prefabPath, out string errorMessage))
                 {
-                    return new ErrorResponse($"Failed to save prefab stage for '{prefabPath}'. The file may be read-only or the disk may be full.");
+                    return new ErrorResponse(errorMessage);
                 }
 
-                return new SuccessResponse($"Saved prefab stage changes for '{prefabPath}'.", new { prefabPath, saved });
+                return new SuccessResponse($"Saved prefab stage changes for '{prefabPath}'.", new { prefabPath, saved = true });
             }
             catch (Exception e)
             {
@@ -1376,10 +1373,9 @@ private static object ClosePrefabStage(bool saveBeforeClose = false)
 
                 if (saveBeforeClose)
                 {
-                    var saveResult = SavePrefabStage();
-                    if (saveResult is ErrorResponse)
+                    if (!TrySavePrefabStage(prefabStage, out _, out string errorMessage))
                     {
-                        return saveResult;
+                        return new ErrorResponse(errorMessage);
                     }
                 }
 
@@ -1393,6 +1389,31 @@ private static object ClosePrefabStage(bool saveBeforeClose = false)
             }
         }
 
+        private static bool TrySavePrefabStage(PrefabStage prefabStage, out string prefabPath, out string errorMessage)
+        {
+            prefabPath = prefabStage.assetPath;
+            errorMessage = null;
+
+            if (prefabStage.prefabContentsRoot == null)
+            {
+                errorMessage = $"Failed to save prefab stage for '{prefabPath}'. The prefab contents root is missing.";
+                return false;
+            }
+
+            bool saved;
+            PrefabUtility.SaveAsPrefabAsset(prefabStage.prefabContentsRoot, prefabPath, out saved);
+            if (!saved)
+            {
+                errorMessage = $"Failed to save prefab stage for '{prefabPath}'. The file may be read-only or the disk may be full.";
+                return false;
+            }
+
+            prefabStage.ClearDirtiness();
+            AssetDatabase.SaveAssets();
+            AssetDatabase.Refresh();
+            return true;
+        }
+
         #endregion
     }
 }

TestProjects/UnityMCPTests/Assets/Tests/EditMode/Tools/ManagePrefabsStageTests.cs

@@ -151,6 +151,86 @@ public void OpenPrefabStage_PrefabPathTakesPrecedenceOverPath()
             }
         }
 
+        [Test]
+        public void SavePrefabStage_PersistsPrefabContentChanges()
+        {
+            string prefabPath = CreateTestPrefab("SaveStageRoot");
+
+            try
+            {
+                AssertOpenPrefabStage(prefabPath);
+
+                var stage = PrefabStageUtility.GetCurrentPrefabStage();
+                Assert.IsNotNull(stage, "Expected prefab stage to be open.");
+
+                var child = new GameObject("SavedChild");
+                child.transform.SetParent(stage.prefabContentsRoot.transform, false);
+
+                var saveResult = ToJObject(ManagePrefabs.HandleCommand(new JObject
+                {
+                    ["action"] = "save_prefab_stage"
+                }));
+
+                Assert.IsTrue(saveResult.Value<bool>("success"), $"Expected save to succeed but got: {saveResult}");
+                Assert.AreEqual(prefabPath, saveResult["data"].Value<string>("prefabPath"));
+
+                StageUtility.GoToMainStage();
+
+                GameObject reloaded = AssetDatabase.LoadAssetAtPath<GameObject>(prefabPath);
+                Assert.IsNotNull(reloaded.transform.Find("SavedChild"), "Saved prefab should contain the new child after save_prefab_stage.");
+            }
+            finally
+            {
+                StageUtility.GoToMainStage();
+                SafeDeleteAsset(prefabPath);
+            }
+        }
+
+        [Test]
+        public void ClosePrefabStage_SaveBeforeClose_PersistsChanges()
+        {
+            string prefabPath = CreateTestPrefab("CloseSaveRoot");
+
+            try
+            {
+                AssertOpenPrefabStage(prefabPath);
+
+                var stage = PrefabStageUtility.GetCurrentPrefabStage();
+                Assert.IsNotNull(stage, "Expected prefab stage to be open.");
+
+                var child = new GameObject("CloseSavedChild");
+                child.transform.SetParent(stage.prefabContentsRoot.transform, false);
+
+                var closeResult = ToJObject(ManagePrefabs.HandleCommand(new JObject
+                {
+                    ["action"] = "close_prefab_stage",
+                    ["saveBeforeClose"] = true
+                }));
+
+                Assert.IsTrue(closeResult.Value<bool>("success"), $"Expected close with save to succeed but got: {closeResult}");
+                Assert.IsNull(PrefabStageUtility.GetCurrentPrefabStage(), "Prefab stage should be closed after close_prefab_stage.");
+
+                GameObject reloaded = AssetDatabase.LoadAssetAtPath<GameObject>(prefabPath);
+                Assert.IsNotNull(reloaded.transform.Find("CloseSavedChild"), "Saved prefab should contain the new child after close_prefab_stage(saveBeforeClose: true).");
+            }
+            finally
+            {
+                StageUtility.GoToMainStage();
+                SafeDeleteAsset(prefabPath);
+            }
+        }
+
+        private static void AssertOpenPrefabStage(string prefabPath)
+        {
+            var openResult = ToJObject(ManagePrefabs.HandleCommand(new JObject
+            {
+                ["action"] = "open_prefab_stage",
+                ["prefabPath"] = prefabPath
+            }));
+
+            Assert.IsTrue(openResult.Value<bool>("success"), $"Expected open to succeed but got: {openResult}");
+        }
+
         private static string CreateTestPrefab(string rootName)
         {
             string prefabPath = Path.Combine(TempDirectory, $"{rootName}.prefab").Replace('\\', '/');