Merge pull request #1013 from Scriptwonder/prefab-fix

Prefab stages integration

Author: Scriptwonder

MCPForUnity/Editor/Tools/GameObjects/ManageGameObject.cs

@@ -80,7 +80,7 @@ public static object HandleCommand(JObject @params)
                 return new ErrorResponse(
                     $"Target '{targetPath}' is a prefab asset. " +
                     $"Use 'manage_asset' with action='modify' for prefab asset modifications, " +
-                    $"or 'manage_prefabs' with action='modify_contents' to edit the prefab headlessly, or 'manage_editor' with action='close_prefab_stage' to exit prefab editing mode."
+                    $"or 'manage_prefabs' with action='modify_contents' to edit the prefab headlessly, or 'manage_prefabs' with action='close_prefab_stage' to exit prefab editing mode."
                 );
             }
             // --- End Prefab Asset Check ---

MCPForUnity/Editor/Tools/ManageEditor.cs

@@ -3,7 +3,6 @@
 using MCPForUnity.Editor.Services;
 using Newtonsoft.Json.Linq;
 using UnityEditor;
-using UnityEditor.SceneManagement;
 using UnityEditorInternal; // Required for tag management
 using UnityEngine;
 
@@ -47,8 +46,6 @@ public static object HandleCommand(JObject @params)
             // Parameters for specific actions
             string tagName = p.Get("tagName");
             string layerName = p.Get("layerName");
-            string prefabPath = p.Get("prefabPath") ?? p.Get("path");
-
             // Route action
             switch (action)
             {
@@ -137,14 +134,6 @@ public static object HandleCommand(JObject @params)
                 //     // Handle string name or int index
                 //     return SetQualityLevel(@params["qualityLevel"]);
 
-                // Prefab Stage
-                case "open_prefab_stage":
-                    return OpenPrefabStage(prefabPath);
-                case "save_prefab_stage":
-                    return SavePrefabStage();
-                case "close_prefab_stage":
-                    return ClosePrefabStage();
-
                 // Package Deployment
                 case "deploy_package":
                     return DeployPackage();
@@ -182,7 +171,7 @@ public static object HandleCommand(JObject @params)
 
                 default:
                     return new ErrorResponse(
-                        $"Unknown action: '{action}'. Supported actions: play, pause, stop, set_active_tool, add_tag, remove_tag, add_layer, remove_layer, open_prefab_stage, save_prefab_stage, close_prefab_stage, deploy_package, restore_package, undo, redo. Use MCP resources for reading editor state, project info, tags, layers, selection, windows, prefab stage, and active tool."
+                        $"Unknown action: '{action}'. Supported actions: play, pause, stop, set_active_tool, add_tag, remove_tag, add_layer, remove_layer, deploy_package, restore_package, undo, redo. For prefab editing (open/save/close prefab stage), use manage_prefabs. Use MCP resources for reading editor state, project info, tags, layers, selection, windows, prefab stage, and active tool."
                     );
             }
         }
@@ -402,112 +391,6 @@ private static object RemoveLayer(string layerName)
             }
         }
 
-        // --- Prefab Stage Methods ---
-
-        private static object OpenPrefabStage(string requestedPath)
-        {
-            if (string.IsNullOrWhiteSpace(requestedPath))
-            {
-                return new ErrorResponse("'prefabPath' parameter is required for open_prefab_stage.");
-            }
-
-            string sanitizedPath = AssetPathUtility.SanitizeAssetPath(requestedPath);
-            if (sanitizedPath == null)
-            {
-                return new ErrorResponse($"Invalid prefab path (path traversal detected): '{requestedPath}'.");
-            }
-
-            if (!sanitizedPath.StartsWith("Assets/", StringComparison.OrdinalIgnoreCase))
-            {
-                return new ErrorResponse($"Prefab path must be within the Assets folder. Got: '{sanitizedPath}'.");
-            }
-
-            if (!sanitizedPath.EndsWith(".prefab", StringComparison.OrdinalIgnoreCase))
-            {
-                return new ErrorResponse($"Prefab path must end with '.prefab'. Got: '{sanitizedPath}'.");
-            }
-
-            try
-            {
-                GameObject prefabAsset = AssetDatabase.LoadAssetAtPath<GameObject>(sanitizedPath);
-                if (prefabAsset == null)
-                {
-                    return new ErrorResponse($"Prefab asset not found at '{sanitizedPath}'.");
-                }
-
-                var prefabStage = PrefabStageUtility.OpenPrefab(sanitizedPath);
-                bool enteredStage = prefabStage != null
-                    && string.Equals(prefabStage.assetPath, sanitizedPath, StringComparison.OrdinalIgnoreCase)
-                    && prefabStage.prefabContentsRoot != null;
-
-                if (!enteredStage)
-                {
-                    return new ErrorResponse($"Failed to open prefab stage for '{sanitizedPath}'. PrefabStageUtility.OpenPrefab did not enter the requested prefab stage.");
-                }
-
-                return new SuccessResponse(
-                    $"Opened prefab stage for '{sanitizedPath}'.",
-                    new
-                    {
-                        prefabPath = sanitizedPath,
-                        openedPrefabPath = prefabStage.assetPath,
-                        rootName = prefabStage.prefabContentsRoot.name,
-                        enteredPrefabStage = enteredStage
-                    }
-                );
-            }
-            catch (Exception e)
-            {
-                return new ErrorResponse($"Error opening prefab stage: {e.Message}");
-            }
-        }
-
-        private static object SavePrefabStage()
-        {
-            try
-            {
-                var prefabStage = PrefabStageUtility.GetCurrentPrefabStage();
-                if (prefabStage == null)
-                {
-                    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)
-                {
-                    return new ErrorResponse($"Failed to save prefab stage for '{prefabPath}'. The file may be read-only or the disk may be full.");
-                }
-
-                return new SuccessResponse($"Saved prefab stage changes for '{prefabPath}'.", new { prefabPath, saved });
-            }
-            catch (Exception e)
-            {
-                return new ErrorResponse($"Error saving prefab stage: {e.Message}");
-            }
-        }
-
-        private static object ClosePrefabStage()
-        {
-            try
-            {
-                var prefabStage = PrefabStageUtility.GetCurrentPrefabStage();
-                if (prefabStage == null)
-                {
-                    return new SuccessResponse("Not currently in prefab editing mode.");
-                }
-
-                string prefabPath = prefabStage.assetPath;
-                StageUtility.GoToMainStage();
-                return new SuccessResponse($"Exited prefab stage for '{prefabPath}'.", new { prefabPath });
-            }
-            catch (Exception e)
-            {
-                return new ErrorResponse($"Error closing prefab stage: {e.Message}");
-            }
-        }
-
         // --- Package Deployment Methods ---
 
         private static object DeployPackage()

MCPForUnity/Editor/Tools/Prefabs/ManagePrefabs.cs

@@ -13,8 +13,8 @@ namespace MCPForUnity.Editor.Tools.Prefabs
 {
     [McpForUnityTool("manage_prefabs", AutoRegister = false)]
     /// <summary>
-    /// Tool to manage Unity Prefabs: create, inspect, and modify prefab assets.
-    /// Uses headless editing (no UI, no dialogs) for reliable automated workflows.
+    /// Tool to manage Unity Prefabs: create, inspect, modify, and open/save/close prefab stage.
+    /// Supports both headless editing (modify_contents) and interactive prefab stage workflows.
     /// </summary>
     public static class ManagePrefabs
     {
@@ -23,7 +23,10 @@ public static class ManagePrefabs
         private const string ACTION_GET_INFO = "get_info";
         private const string ACTION_GET_HIERARCHY = "get_hierarchy";
         private const string ACTION_MODIFY_CONTENTS = "modify_contents";
-        private const string SupportedActions = ACTION_CREATE_FROM_GAMEOBJECT + ", " + ACTION_GET_INFO + ", " + ACTION_GET_HIERARCHY + ", " + ACTION_MODIFY_CONTENTS;
+        private const string ACTION_OPEN_PREFAB_STAGE = "open_prefab_stage";
+        private const string ACTION_SAVE_PREFAB_STAGE = "save_prefab_stage";
+        private const string ACTION_CLOSE_PREFAB_STAGE = "close_prefab_stage";
+        private const string SupportedActions = ACTION_CREATE_FROM_GAMEOBJECT + ", " + ACTION_GET_INFO + ", " + ACTION_GET_HIERARCHY + ", " + ACTION_MODIFY_CONTENTS + ", " + ACTION_OPEN_PREFAB_STAGE + ", " + ACTION_SAVE_PREFAB_STAGE + ", " + ACTION_CLOSE_PREFAB_STAGE;
 
         public static object HandleCommand(JObject @params)
         {
@@ -50,6 +53,18 @@ public static object HandleCommand(JObject @params)
                         return GetHierarchy(@params);
                     case ACTION_MODIFY_CONTENTS:
                         return ModifyContents(@params);
+                    case ACTION_OPEN_PREFAB_STAGE:
+                    {
+                        string prefabPath = @params["prefabPath"]?.ToString() ?? @params["path"]?.ToString();
+                        return OpenPrefabStage(prefabPath);
+                    }
+                    case ACTION_SAVE_PREFAB_STAGE:
+                        return SavePrefabStage();
+                    case ACTION_CLOSE_PREFAB_STAGE:
+                    {
+                        bool saveBeforeClose = @params["saveBeforeClose"]?.ToObject<bool>() ?? false;
+                        return ClosePrefabStage(saveBeforeClose);
+                    }
                     default:
                         return new ErrorResponse($"Unknown action: '{action}'. Valid actions are: {SupportedActions}.");
                 }
@@ -1261,5 +1276,122 @@ private static void BuildHierarchyItemsRecursive(Transform transform, Transform
         }
 
         #endregion
+
+        #region Prefab Stage
+
+        private static object OpenPrefabStage(string requestedPath)
+        {
+            if (string.IsNullOrWhiteSpace(requestedPath))
+            {
+                return new ErrorResponse("Either 'prefabPath' or 'path' parameter is required for open_prefab_stage.");
+            }
+
+            string sanitizedPath = AssetPathUtility.SanitizeAssetPath(requestedPath);
+            if (sanitizedPath == null)
+            {
+                return new ErrorResponse($"Invalid prefab path (path traversal detected): '{requestedPath}'.");
+            }
+
+            if (!sanitizedPath.StartsWith("Assets/", StringComparison.OrdinalIgnoreCase))
+            {
+                return new ErrorResponse($"Prefab path must be within the Assets folder. Got: '{sanitizedPath}'.");
+            }
+
+            if (!sanitizedPath.EndsWith(".prefab", StringComparison.OrdinalIgnoreCase))
+            {
+                return new ErrorResponse($"Prefab path must end with '.prefab'. Got: '{sanitizedPath}'.");
+            }
+
+            try
+            {
+                GameObject prefabAsset = AssetDatabase.LoadAssetAtPath<GameObject>(sanitizedPath);
+                if (prefabAsset == null)
+                {
+                    return new ErrorResponse($"Prefab asset not found at '{sanitizedPath}'.");
+                }
+
+                var prefabStage = PrefabStageUtility.OpenPrefab(sanitizedPath);
+                bool enteredStage = prefabStage != null
+                    && string.Equals(prefabStage.assetPath, sanitizedPath, StringComparison.OrdinalIgnoreCase)
+                    && prefabStage.prefabContentsRoot != null;
+
+                if (!enteredStage)
+                {
+                    return new ErrorResponse($"Failed to open prefab stage for '{sanitizedPath}'. PrefabStageUtility.OpenPrefab did not enter the requested prefab stage.");
+                }
+
+                return new SuccessResponse(
+                    $"Opened prefab stage for '{sanitizedPath}'.",
+                    new
+                    {
+                        prefabPath = sanitizedPath,
+                        openedPrefabPath = prefabStage.assetPath,
+                        rootName = prefabStage.prefabContentsRoot.name,
+                        enteredPrefabStage = enteredStage
+                    }
+                );
+            }
+            catch (Exception e)
+            {
+                return new ErrorResponse($"Error opening prefab stage: {e.Message}");
+            }
+        }
+
+        private static object SavePrefabStage()
+        {
+            try
+            {
+                var prefabStage = PrefabStageUtility.GetCurrentPrefabStage();
+                if (prefabStage == null)
+                {
+                    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)
+                {
+                    return new ErrorResponse($"Failed to save prefab stage for '{prefabPath}'. The file may be read-only or the disk may be full.");
+                }
+
+                return new SuccessResponse($"Saved prefab stage changes for '{prefabPath}'.", new { prefabPath, saved });
+            }
+            catch (Exception e)
+            {
+                return new ErrorResponse($"Error saving prefab stage: {e.Message}");
+            }
+        }
+
+        private static object ClosePrefabStage(bool saveBeforeClose = false)
+        {
+            try
+            {
+                var prefabStage = PrefabStageUtility.GetCurrentPrefabStage();
+                if (prefabStage == null)
+                {
+                    return new SuccessResponse("Not currently in prefab editing mode.");
+                }
+
+                if (saveBeforeClose)
+                {
+                    var saveResult = SavePrefabStage();
+                    if (saveResult is ErrorResponse)
+                    {
+                        return saveResult;
+                    }
+                }
+
+                string prefabPath = prefabStage.assetPath;
+                StageUtility.GoToMainStage();
+                return new SuccessResponse($"Exited prefab stage for '{prefabPath}'.", new { prefabPath });
+            }
+            catch (Exception e)
+            {
+                return new ErrorResponse($"Error closing prefab stage: {e.Message}");
+            }
+        }
+
+        #endregion
     }
 }

Server/src/cli/commands/prefab.py

@@ -29,7 +29,7 @@ def open_stage(path: str):
     config = get_config()
 
     params: dict[str, Any] = {
-        "action": "open_stage",
+        "action": "open_prefab_stage",
         "prefabPath": path,
     }
 
@@ -57,7 +57,7 @@ def close_stage(save: bool):
     config = get_config()
 
     params: dict[str, Any] = {
-        "action": "close_stage",
+        "action": "close_prefab_stage",
     }
     if save:
         params["saveBeforeClose"] = True
@@ -83,7 +83,7 @@ def save_stage():
         "action": "save_prefab_stage",
     }
 
-    result = run_command("manage_editor", params, config)
+    result = run_command("manage_prefabs", params, config)
     click.echo(format_output(result, config.format))
     if result.get("success"):
         print_success("Saved prefab stage")

Server/src/services/resources/prefab.py

@@ -57,7 +57,7 @@ async def get_prefab_api_docs(_ctx: Context) -> MCPResponse:
         "workflow": [
             "1. Use manage_asset action=search filterType=Prefab to find prefabs",
             "2. Use the asset path to access detailed data via resources below",
-            "3. Use manage_editor action=open_prefab_stage / save_prefab_stage / close_prefab_stage for prefab editing UI transitions"
+            "3. Use manage_prefabs action=open_prefab_stage / save_prefab_stage / close_prefab_stage for prefab editing UI transitions"
         ],
         "path_encoding": {
             "note": "Prefab paths must be URL-encoded when used in resource URIs",
@@ -80,8 +80,8 @@ async def get_prefab_api_docs(_ctx: Context) -> MCPResponse:
             }
         },
         "related_tools": {
-            "manage_editor": "Open/save/close prefab stages in the Unity Editor UI",
-            "manage_prefabs": "Headless prefab inspection and modification without opening prefab stages",
+            "manage_editor": "Editor controls (play/pause/stop, active tool, tags/layers, package deploy/restore)",
+            "manage_prefabs": "Prefab stage lifecycle (open/save/close) and headless prefab inspection/modification",
             "manage_asset": "Search for prefab assets, get asset info",
             "manage_gameobject": "Modify GameObjects in open prefab stage",
             "manage_components": "Add/remove/modify components on prefab GameObjects"