using System; using System.Collections.Generic; using System.IO; using MCPForUnity.Editor.Helpers; using Newtonsoft.Json.Linq; using UnityEditor; using UnityEditor.SceneManagement; using UnityEngine; using UnityEngine.SceneManagement; namespace MCPForUnity.Editor.Tools.Prefabs { [McpForUnityTool("manage_prefabs", AutoRegister = false)] /// /// Tool to manage Unity Prefab stages and create prefabs from GameObjects. /// public static class ManagePrefabs { // Action constants private const string ACTION_OPEN_STAGE = "open_stage"; private const string ACTION_CLOSE_STAGE = "close_stage"; private const string ACTION_SAVE_OPEN_STAGE = "save_open_stage"; private const string ACTION_CREATE_FROM_GAMEOBJECT = "create_from_gameobject"; private const string ACTION_GET_INFO = "get_info"; private const string ACTION_GET_HIERARCHY = "get_hierarchy"; private const string SupportedActions = ACTION_OPEN_STAGE + ", " + ACTION_CLOSE_STAGE + ", " + ACTION_SAVE_OPEN_STAGE + ", " + ACTION_CREATE_FROM_GAMEOBJECT + ", " + ACTION_GET_INFO + ", " + ACTION_GET_HIERARCHY; public static object HandleCommand(JObject @params) { if (@params == null) { return new ErrorResponse("Parameters cannot be null."); } string action = @params["action"]?.ToString()?.ToLowerInvariant(); if (string.IsNullOrEmpty(action)) { return new ErrorResponse($"Action parameter is required. Valid actions are: {SupportedActions}."); } try { switch (action) { case ACTION_OPEN_STAGE: return OpenStage(@params); case ACTION_CLOSE_STAGE: return CloseStage(@params); case ACTION_SAVE_OPEN_STAGE: return SaveOpenStage(@params); case ACTION_CREATE_FROM_GAMEOBJECT: return CreatePrefabFromGameObject(@params); case ACTION_GET_INFO: return GetInfo(@params); case ACTION_GET_HIERARCHY: return GetHierarchy(@params); default: return new ErrorResponse($"Unknown action: '{action}'. Valid actions are: {SupportedActions}."); } } catch (Exception e) { McpLog.Error($"[ManagePrefabs] Action '{action}' failed: {e}"); return new ErrorResponse($"Internal error: {e.Message}"); } } /// /// Opens a prefab in prefab mode for editing. /// private static object OpenStage(JObject @params) { string prefabPath = @params["prefabPath"]?.ToString(); if (string.IsNullOrEmpty(prefabPath)) { return new ErrorResponse("'prefabPath' parameter is required for open_stage."); } string sanitizedPath = AssetPathUtility.SanitizeAssetPath(prefabPath); if (string.IsNullOrEmpty(sanitizedPath)) { return new ErrorResponse($"Invalid prefab path: '{prefabPath}'."); } GameObject prefabAsset = AssetDatabase.LoadAssetAtPath(sanitizedPath); if (prefabAsset == null) { return new ErrorResponse($"No prefab asset found at path '{sanitizedPath}'."); } PrefabStage stage = PrefabStageUtility.OpenPrefab(sanitizedPath); if (stage == null) { return new ErrorResponse($"Failed to open prefab stage for '{sanitizedPath}'."); } return new SuccessResponse($"Opened prefab stage for '{sanitizedPath}'.", SerializeStage(stage)); } /// /// Closes the currently open prefab stage, optionally saving first. /// private static object CloseStage(JObject @params) { PrefabStage stage = PrefabStageUtility.GetCurrentPrefabStage(); if (stage == null) { return new SuccessResponse("No prefab stage was open."); } string assetPath = stage.assetPath; bool saveBeforeClose = @params["saveBeforeClose"]?.ToObject() ?? false; if (saveBeforeClose && stage.scene.isDirty) { try { SaveAndRefreshStage(stage); } catch (Exception e) { return new ErrorResponse($"Failed to save prefab before closing: {e.Message}"); } } StageUtility.GoToMainStage(); return new SuccessResponse($"Closed prefab stage for '{assetPath}'."); } /// /// Saves changes to the currently open prefab stage. /// Supports a 'force' parameter for automated workflows where isDirty may not be set. /// private static object SaveOpenStage(JObject @params) { PrefabStage stage = PrefabStageUtility.GetCurrentPrefabStage(); if (stage == null) { return new ErrorResponse("No prefab stage is currently open."); } if (!ValidatePrefabStageForSave(stage)) { return new ErrorResponse("Prefab stage validation failed. Cannot save."); } // Check for force parameter (useful for automated workflows) bool force = @params?["force"]?.ToObject() ?? false; // Check if there are actual changes to save bool wasDirty = stage.scene.isDirty; if (!wasDirty && !force) { return new SuccessResponse($"Prefab stage for '{stage.assetPath}' has no unsaved changes.", SerializeStage(stage)); } try { SaveAndRefreshStage(stage, force); return new SuccessResponse($"Saved prefab stage for '{stage.assetPath}'.", SerializeStage(stage)); } catch (Exception e) { return new ErrorResponse($"Failed to save prefab: {e.Message}"); } } #region Prefab Save Operations /// /// Saves the prefab stage and refreshes the asset database. /// Uses PrefabUtility.SaveAsPrefabAsset for reliable prefab saving without dialogs. /// /// The prefab stage to save. /// If true, marks the prefab dirty before saving to ensure changes are captured. private static void SaveAndRefreshStage(PrefabStage stage, bool force = false) { if (stage == null) { throw new ArgumentNullException(nameof(stage), "Prefab stage cannot be null."); } if (stage.prefabContentsRoot == null) { throw new InvalidOperationException("Cannot save prefab stage without a prefab root."); } if (string.IsNullOrEmpty(stage.assetPath)) { throw new InvalidOperationException("Prefab stage has invalid asset path."); } // When force=true, mark the prefab root dirty to ensure changes are saved // This is useful for automated workflows where isDirty may not be set correctly if (force) { EditorUtility.SetDirty(stage.prefabContentsRoot); EditorSceneManager.MarkSceneDirty(stage.scene); } // Mark all children as dirty to ensure their changes are captured foreach (Transform child in stage.prefabContentsRoot.GetComponentsInChildren(true)) { if (child != stage.prefabContentsRoot.transform) { EditorUtility.SetDirty(child.gameObject); } } // Use PrefabUtility.SaveAsPrefabAsset which saves without dialogs // This is more reliable for automated workflows than EditorSceneManager.SaveScene bool success; PrefabUtility.SaveAsPrefabAsset(stage.prefabContentsRoot, stage.assetPath, out success); if (!success) { throw new InvalidOperationException($"Failed to save prefab asset for '{stage.assetPath}'."); } // Ensure changes are persisted to disk AssetDatabase.SaveAssets(); AssetDatabase.Refresh(); McpLog.Info($"[ManagePrefabs] Successfully saved prefab '{stage.assetPath}'."); } /// /// Validates prefab stage before saving. /// private static bool ValidatePrefabStageForSave(PrefabStage stage) { if (stage == null) { McpLog.Warn("[ManagePrefabs] No prefab stage is open."); return false; } if (stage.prefabContentsRoot == null) { McpLog.Error($"[ManagePrefabs] Prefab stage '{stage.assetPath}' has no root object."); return false; } if (string.IsNullOrEmpty(stage.assetPath)) { McpLog.Error("[ManagePrefabs] Prefab stage has invalid asset path."); return false; } return true; } #endregion #region Create Prefab from GameObject /// /// Creates a prefab asset from a GameObject in the scene. /// private static object CreatePrefabFromGameObject(JObject @params) { // 1. Validate and parse parameters var validation = ValidateCreatePrefabParams(@params); if (!validation.isValid) { return new ErrorResponse(validation.errorMessage); } string targetName = validation.targetName; string finalPath = validation.finalPath; bool includeInactive = validation.includeInactive; bool replaceExisting = validation.replaceExisting; bool unlinkIfInstance = validation.unlinkIfInstance; // 2. Find the source object GameObject sourceObject = FindSceneObjectByName(targetName, includeInactive); if (sourceObject == null) { return new ErrorResponse($"GameObject '{targetName}' not found in the active scene or prefab stage{(includeInactive ? " (including inactive objects)" : "")}."); } // 3. Validate source object state var objectValidation = ValidateSourceObjectForPrefab(sourceObject, unlinkIfInstance); if (!objectValidation.isValid) { return new ErrorResponse(objectValidation.errorMessage); } // 4. Check for path conflicts and track if file will be replaced bool fileExistedAtPath = AssetDatabase.LoadAssetAtPath(finalPath) != null; if (!replaceExisting && fileExistedAtPath) { finalPath = AssetDatabase.GenerateUniqueAssetPath(finalPath); McpLog.Info($"[ManagePrefabs] Generated unique path: {finalPath}"); } // 5. Ensure directory exists EnsureAssetDirectoryExists(finalPath); // 6. Unlink from existing prefab if needed if (unlinkIfInstance && objectValidation.shouldUnlink) { try { // UnpackPrefabInstance requires the prefab instance root, not a child object GameObject rootToUnlink = PrefabUtility.GetOutermostPrefabInstanceRoot(sourceObject); if (rootToUnlink != null) { PrefabUtility.UnpackPrefabInstance(rootToUnlink, PrefabUnpackMode.Completely, InteractionMode.AutomatedAction); McpLog.Info($"[ManagePrefabs] Unpacked prefab instance '{rootToUnlink.name}' before creating new prefab."); } } catch (Exception e) { return new ErrorResponse($"Failed to unlink prefab instance: {e.Message}"); } } // 7. Create the prefab try { GameObject result = CreatePrefabAsset(sourceObject, finalPath, replaceExisting); if (result == null) { return new ErrorResponse($"Failed to create prefab asset at '{finalPath}'."); } // 8. Select the newly created instance Selection.activeGameObject = result; return new SuccessResponse( $"Prefab created at '{finalPath}' and instance linked.", new { prefabPath = finalPath, instanceId = result.GetInstanceID(), instanceName = result.name, wasUnlinked = unlinkIfInstance && objectValidation.shouldUnlink, wasReplaced = replaceExisting && fileExistedAtPath, componentCount = result.GetComponents().Length, childCount = result.transform.childCount } ); } catch (Exception e) { McpLog.Error($"[ManagePrefabs] Error creating prefab at '{finalPath}': {e}"); return new ErrorResponse($"Error saving prefab asset: {e.Message}"); } } /// /// Validates parameters for creating a prefab from GameObject. /// private static (bool isValid, string errorMessage, string targetName, string finalPath, bool includeInactive, bool replaceExisting, bool unlinkIfInstance) ValidateCreatePrefabParams(JObject @params) { string targetName = @params["target"]?.ToString() ?? @params["name"]?.ToString(); if (string.IsNullOrEmpty(targetName)) { return (false, "'target' parameter is required for create_from_gameobject.", null, null, false, false, false); } string requestedPath = @params["prefabPath"]?.ToString(); if (string.IsNullOrWhiteSpace(requestedPath)) { return (false, "'prefabPath' parameter is required for create_from_gameobject.", targetName, null, false, false, false); } string sanitizedPath = AssetPathUtility.SanitizeAssetPath(requestedPath); if (sanitizedPath == null) { return (false, $"Invalid prefab path (path traversal detected): '{requestedPath}'", targetName, null, false, false, false); } if (string.IsNullOrEmpty(sanitizedPath)) { return (false, $"Invalid prefab path '{requestedPath}'. Path cannot be empty.", targetName, null, false, false, false); } if (!sanitizedPath.EndsWith(".prefab", StringComparison.OrdinalIgnoreCase)) { sanitizedPath += ".prefab"; } // Validate path is within Assets folder if (!sanitizedPath.StartsWith("Assets/", StringComparison.OrdinalIgnoreCase)) { return (false, $"Prefab path must be within the Assets folder. Got: '{sanitizedPath}'", targetName, null, false, false, false); } bool includeInactive = @params["searchInactive"]?.ToObject() ?? false; bool replaceExisting = @params["allowOverwrite"]?.ToObject() ?? false; bool unlinkIfInstance = @params["unlinkIfInstance"]?.ToObject() ?? false; return (true, null, targetName, sanitizedPath, includeInactive, replaceExisting, unlinkIfInstance); } /// /// Validates source object can be converted to prefab. /// private static (bool isValid, string errorMessage, bool shouldUnlink, string existingPrefabPath) ValidateSourceObjectForPrefab(GameObject sourceObject, bool unlinkIfInstance) { // Check if this is a Prefab Asset (the .prefab file itself in the editor) if (PrefabUtility.IsPartOfPrefabAsset(sourceObject)) { return (false, $"GameObject '{sourceObject.name}' is part of a prefab asset. " + "Open the prefab stage to save changes instead.", false, null); } // Check if this is already a Prefab Instance PrefabInstanceStatus status = PrefabUtility.GetPrefabInstanceStatus(sourceObject); if (status != PrefabInstanceStatus.NotAPrefab) { string existingPath = PrefabUtility.GetPrefabAssetPathOfNearestInstanceRoot(sourceObject); if (!unlinkIfInstance) { return (false, $"GameObject '{sourceObject.name}' is already linked to prefab '{existingPath}'. " + "Set 'unlinkIfInstance' to true to unlink it first, or modify the existing prefab instead.", false, existingPath); } // Needs to be unlinked return (true, null, true, existingPath); } return (true, null, false, null); } /// /// Creates a prefab asset from a GameObject. /// private static GameObject CreatePrefabAsset(GameObject sourceObject, string path, bool replaceExisting) { GameObject result = PrefabUtility.SaveAsPrefabAssetAndConnect( sourceObject, path, InteractionMode.AutomatedAction ); string action = replaceExisting ? "Replaced existing" : "Created new"; McpLog.Info($"[ManagePrefabs] {action} prefab at '{path}'."); if (result != null) { AssetDatabase.SaveAssets(); AssetDatabase.Refresh(); } return result; } #endregion /// /// Ensures the directory for an asset path exists, creating it if necessary. /// private static void EnsureAssetDirectoryExists(string assetPath) { string directory = Path.GetDirectoryName(assetPath); if (string.IsNullOrEmpty(directory)) { return; } // Use Application.dataPath for more reliable path resolution // Application.dataPath points to the Assets folder (e.g., ".../ProjectName/Assets") string assetsPath = Application.dataPath; string projectRoot = Path.GetDirectoryName(assetsPath); string fullDirectory = Path.Combine(projectRoot, directory); if (!Directory.Exists(fullDirectory)) { Directory.CreateDirectory(fullDirectory); AssetDatabase.Refresh(ImportAssetOptions.ForceSynchronousImport); McpLog.Info($"[ManagePrefabs] Created directory: {directory}"); } } /// /// Finds a GameObject by name in the active scene or current prefab stage. /// private static GameObject FindSceneObjectByName(string name, bool includeInactive) { // First check if we're in Prefab Stage PrefabStage stage = PrefabStageUtility.GetCurrentPrefabStage(); if (stage?.prefabContentsRoot != null) { foreach (Transform transform in stage.prefabContentsRoot.GetComponentsInChildren(includeInactive)) { if (transform.name == name && (includeInactive || transform.gameObject.activeSelf)) { return transform.gameObject; } } } // Search in the active scene Scene activeScene = SceneManager.GetActiveScene(); foreach (GameObject root in activeScene.GetRootGameObjects()) { // Check the root object itself if (root.name == name && (includeInactive || root.activeSelf)) { return root; } // Check children foreach (Transform transform in root.GetComponentsInChildren(includeInactive)) { if (transform.name == name && (includeInactive || transform.gameObject.activeSelf)) { return transform.gameObject; } } } return null; } #region Read Operations /// /// Gets basic metadata information about a prefab asset. /// private static object GetInfo(JObject @params) { string prefabPath = @params["prefabPath"]?.ToString() ?? @params["path"]?.ToString(); if (string.IsNullOrEmpty(prefabPath)) { return new ErrorResponse("'prefabPath' parameter is required for get_info."); } string sanitizedPath = AssetPathUtility.SanitizeAssetPath(prefabPath); if (string.IsNullOrEmpty(sanitizedPath)) { return new ErrorResponse($"Invalid prefab path: '{prefabPath}'."); } GameObject prefabAsset = AssetDatabase.LoadAssetAtPath(sanitizedPath); if (prefabAsset == null) { return new ErrorResponse($"No prefab asset found at path '{sanitizedPath}'."); } string guid = PrefabUtilityHelper.GetPrefabGUID(sanitizedPath); PrefabAssetType assetType = PrefabUtility.GetPrefabAssetType(prefabAsset); string prefabTypeString = assetType.ToString(); var componentTypes = PrefabUtilityHelper.GetComponentTypeNames(prefabAsset); int childCount = PrefabUtilityHelper.CountChildrenRecursive(prefabAsset.transform); var (isVariant, parentPrefab, _) = PrefabUtilityHelper.GetVariantInfo(prefabAsset); return new SuccessResponse( $"Successfully retrieved prefab info.", new { assetPath = sanitizedPath, guid = guid, prefabType = prefabTypeString, rootObjectName = prefabAsset.name, rootComponentTypes = componentTypes, childCount = childCount, isVariant = isVariant, parentPrefab = parentPrefab } ); } /// /// Gets the hierarchical structure of a prefab asset. /// Returns all objects in the prefab for full client-side filtering and search. /// private static object GetHierarchy(JObject @params) { string prefabPath = @params["prefabPath"]?.ToString() ?? @params["path"]?.ToString(); if (string.IsNullOrEmpty(prefabPath)) { return new ErrorResponse("'prefabPath' parameter is required for get_hierarchy."); } string sanitizedPath = AssetPathUtility.SanitizeAssetPath(prefabPath); if (string.IsNullOrEmpty(sanitizedPath)) { return new ErrorResponse($"Invalid prefab path '{prefabPath}'. Path traversal sequences are not allowed."); } // Load prefab contents in background (without opening stage UI) GameObject prefabContents = PrefabUtility.LoadPrefabContents(sanitizedPath); if (prefabContents == null) { return new ErrorResponse($"Failed to load prefab contents from '{sanitizedPath}'."); } try { // Build complete hierarchy items (no pagination) var allItems = BuildHierarchyItems(prefabContents.transform, sanitizedPath); return new SuccessResponse( $"Successfully retrieved prefab hierarchy. Found {allItems.Count} objects.", new { prefabPath = sanitizedPath, total = allItems.Count, items = allItems } ); } finally { // Always unload prefab contents to free memory PrefabUtility.UnloadPrefabContents(prefabContents); } } #endregion #region Hierarchy Builder /// /// Builds a flat list of hierarchy items from a transform root. /// /// The root transform of the prefab. /// Asset path of the main prefab. /// List of hierarchy items with prefab information. private static List BuildHierarchyItems(Transform root, string mainPrefabPath) { var items = new List(); BuildHierarchyItemsRecursive(root, root, mainPrefabPath, "", items); return items; } /// /// Recursively builds hierarchy items. /// /// Current transform being processed. /// Root transform of the main prefab asset. /// Asset path of the main prefab. /// Parent path for building full hierarchy path. /// List to accumulate hierarchy items. private static void BuildHierarchyItemsRecursive(Transform transform, Transform mainPrefabRoot, string mainPrefabPath, string parentPath, List items) { if (transform == null) return; string name = transform.gameObject.name; string path = string.IsNullOrEmpty(parentPath) ? name : $"{parentPath}/{name}"; int instanceId = transform.gameObject.GetInstanceID(); bool activeSelf = transform.gameObject.activeSelf; int childCount = transform.childCount; var componentTypes = PrefabUtilityHelper.GetComponentTypeNames(transform.gameObject); // Prefab information bool isNestedPrefab = PrefabUtility.IsAnyPrefabInstanceRoot(transform.gameObject); bool isPrefabRoot = transform == mainPrefabRoot; int nestingDepth = isPrefabRoot ? 0 : PrefabUtilityHelper.GetPrefabNestingDepth(transform.gameObject, mainPrefabRoot); string parentPrefabPath = isNestedPrefab && !isPrefabRoot ? PrefabUtilityHelper.GetParentPrefabPath(transform.gameObject, mainPrefabRoot) : null; string nestedPrefabPath = isNestedPrefab ? PrefabUtilityHelper.GetNestedPrefabPath(transform.gameObject) : null; var item = new { name = name, instanceId = instanceId, path = path, activeSelf = activeSelf, childCount = childCount, componentTypes = componentTypes, prefab = new { isRoot = isPrefabRoot, isNestedRoot = isNestedPrefab, nestingDepth = nestingDepth, assetPath = isNestedPrefab ? nestedPrefabPath : mainPrefabPath, parentPath = parentPrefabPath } }; items.Add(item); // Recursively process children foreach (Transform child in transform) { BuildHierarchyItemsRecursive(child, mainPrefabRoot, mainPrefabPath, path, items); } } #endregion /// /// Serializes the prefab stage information for response. /// private static object SerializeStage(PrefabStage stage) { if (stage == null) { return new { isOpen = false }; } return new { isOpen = true, assetPath = stage.assetPath, prefabRootName = stage.prefabContentsRoot != null ? stage.prefabContentsRoot.name : null, mode = stage.mode.ToString(), isDirty = stage.scene.isDirty }; } } }