fix: Add Prefab Stage support for GameObject lookup (#573)

* Enhance Prefab Stage support in GameObject lookup and scene management

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

* fix: unify path matching and restore fast path lookup

---------

Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>

MCPForUnity/Editor/Helpers/GameObjectLookup.cs

@@ -3,6 +3,7 @@ using System.Collections.Generic;
 using System.Linq;
 using Newtonsoft.Json.Linq;
 using UnityEditor;
+using UnityEditor.SceneManagement;
 using UnityEngine;
 using UnityEngine.SceneManagement;
 
@@ -147,21 +148,47 @@ namespace MCPForUnity.Editor.Helpers
 
         private static IEnumerable<int> SearchByPath(string path, bool includeInactive)
         {
-            // NOTE: Unity's GameObject.Find(path) only finds ACTIVE GameObjects.
+            // Check Prefab Stage first - GameObject.Find() doesn't work in Prefab Stage
-            // The includeInactive parameter has no effect here due to Unity API limitations.
+            var prefabStage = PrefabStageUtility.GetCurrentPrefabStage();
-            // Consider using by_name search with includeInactive if you need to find inactive objects.
+            if (prefabStage != null)
-            if (includeInactive)
             {
-                McpLog.Warn("[GameObjectLookup] SearchByPath with includeInactive=true: " +
+                // Use GetAllSceneObjects which already handles Prefab Stage
-                    "GameObject.Find() cannot find inactive objects. Use by_name search instead.");
+                var allObjects = GetAllSceneObjects(includeInactive);
+                foreach (var go in allObjects)
+                {
+                    if (MatchesPath(go, path))
+                    {
+                        yield return go.GetInstanceID();
+                    }
+                }
+                yield break;
             }
 
+            // Normal scene mode
+            // NOTE: Unity's GameObject.Find(path) only finds ACTIVE GameObjects.
+            // If includeInactive=true, we need to search manually to find inactive objects.
+            if (includeInactive)
+            {
+                // Search manually to support inactive objects
+                var allObjects = GetAllSceneObjects(true);
+                foreach (var go in allObjects)
+                {
+                    if (MatchesPath(go, path))
+                    {
+                        yield return go.GetInstanceID();
+                    }
+                }
+            }
+            else
+            {
+                // Use GameObject.Find for active objects only (Unity API limitation)
                 var found = GameObject.Find(path);
                 if (found != null)
                 {
                     yield return found.GetInstanceID();
                 }
             }
+        }
 
         private static IEnumerable<int> SearchByTag(string tag, bool includeInactive, int maxResults)
         {
@@ -249,6 +276,19 @@ namespace MCPForUnity.Editor.Helpers
         /// </summary>
         public static IEnumerable<GameObject> GetAllSceneObjects(bool includeInactive)
         {
+            // Check Prefab Stage first
+            var prefabStage = PrefabStageUtility.GetCurrentPrefabStage();
+            if (prefabStage != null && prefabStage.prefabContentsRoot != null)
+            {
+                // Use Prefab Stage's prefabContentsRoot
+                foreach (var go in GetObjectAndDescendants(prefabStage.prefabContentsRoot, includeInactive))
+                {
+                    yield return go;
+                }
+                yield break;
+            }
+
+            // Normal scene mode
             var scene = SceneManager.GetActiveScene();
             if (!scene.IsValid())
                 yield break;
@@ -290,6 +330,18 @@ namespace MCPForUnity.Editor.Helpers
             return UnityTypeResolver.ResolveComponent(typeName);
         }
 
+        /// <summary>
+        /// Checks whether a GameObject matches a path or trailing path segment.
+        /// </summary>
+        internal static bool MatchesPath(GameObject go, string path)
+        {
+            if (go == null || string.IsNullOrEmpty(path))
+                return false;
+
+            var goPath = GetGameObjectPath(go);
+            return goPath == path || goPath.EndsWith("/" + path);
+        }
+
         /// <summary>
         /// Gets the hierarchical path of a GameObject.
         /// </summary>

MCPForUnity/Editor/Tools/GameObjects/ManageGameObjectCommon.cs

@@ -5,6 +5,7 @@ using System.Linq;
 using MCPForUnity.Editor.Helpers;
 using MCPForUnity.Editor.Tools;
 using Newtonsoft.Json.Linq;
+using UnityEditor.SceneManagement;
 using UnityEngine;
 using UnityEngine.SceneManagement;
 
@@ -83,11 +84,34 @@ namespace MCPForUnity.Editor.Tools.GameObjects
                     break;
 
                 case "by_path":
-                    Transform foundTransform = rootSearchObject
+                    if (rootSearchObject != null)
-                        ? rootSearchObject.transform.Find(searchTerm)
+                    {
-                        : GameObject.Find(searchTerm)?.transform;
+                        Transform foundTransform = rootSearchObject.transform.Find(searchTerm);
                         if (foundTransform != null)
                             results.Add(foundTransform.gameObject);
+                    }
+                    else
+                    {
+                        var prefabStage = PrefabStageUtility.GetCurrentPrefabStage();
+                        if (prefabStage != null || searchInactive)
+                        {
+                            // In Prefab Stage, GameObject.Find() doesn't work, need to search manually
+                            var allObjects = GetAllSceneObjects(searchInactive);
+                            foreach (var go in allObjects)
+                            {
+                                if (GameObjectLookup.MatchesPath(go, searchTerm))
+                                {
+                                    results.Add(go);
+                                }
+                            }
+                        }
+                        else
+                        {
+                            var found = GameObject.Find(searchTerm);
+                            if (found != null)
+                                results.Add(found);
+                        }
+                    }
                     break;
 
                 case "by_tag":
@@ -154,7 +178,12 @@ namespace MCPForUnity.Editor.Tools.GameObjects
                         }
                     }
 
-                    GameObject objByPath = GameObject.Find(searchTerm);
+                    // Try path search - in Prefab Stage, GameObject.Find() doesn't work
+                    var allObjectsForPath = GetAllSceneObjects(true);
+                    GameObject objByPath = allObjectsForPath.FirstOrDefault(go =>
+                    {
+                        return GameObjectLookup.MatchesPath(go, searchTerm);
+                    });
                     if (objByPath != null)
                     {
                         results.Add(objByPath);
@@ -180,16 +209,8 @@ namespace MCPForUnity.Editor.Tools.GameObjects
 
         private static IEnumerable<GameObject> GetAllSceneObjects(bool includeInactive)
         {
-            var rootObjects = SceneManager.GetActiveScene().GetRootGameObjects();
+            // Delegate to GameObjectLookup to avoid code duplication and ensure consistent behavior
-            var allObjects = new List<GameObject>();
+            return GameObjectLookup.GetAllSceneObjects(includeInactive);
-            foreach (var root in rootObjects)
-            {
-                allObjects.AddRange(
-                    root.GetComponentsInChildren<Transform>(includeInactive)
-                        .Select(t => t.gameObject)
-                );
-            }
-            return allObjects;
         }
 
         private static Type FindType(string typeName)

MCPForUnity/Editor/Tools/ManageScene.cs

@@ -489,9 +489,22 @@ namespace MCPForUnity.Editor.Tools
         private static object GetSceneHierarchyPaged(SceneCommand cmd)
         {
             try
+            {
+                // Check Prefab Stage first
+                var prefabStage = PrefabStageUtility.GetCurrentPrefabStage();
+                Scene activeScene;
+                
+                if (prefabStage != null)
+                {
+                    activeScene = prefabStage.scene;
+                    try { McpLog.Info("[ManageScene] get_hierarchy: using Prefab Stage scene", always: false); } catch { }
+                }
+                else
                 {
                     try { McpLog.Info("[ManageScene] get_hierarchy: querying EditorSceneManager.GetActiveScene", always: false); } catch { }
-                Scene activeScene = EditorSceneManager.GetActiveScene();
+                    activeScene = EditorSceneManager.GetActiveScene();
+                }
+                
                 try { McpLog.Info($"[ManageScene] get_hierarchy: got scene valid={activeScene.IsValid()} loaded={activeScene.isLoaded} name='{activeScene.name}'", always: false); } catch { }
                 if (!activeScene.IsValid() || !activeScene.isLoaded)
                 {
@@ -599,7 +612,16 @@ namespace MCPForUnity.Editor.Tools
             // Path-based find (e.g., "Root/Child/GrandChild")
             if (s.Contains("/"))
             {
-                try { return GameObject.Find(s); } catch { }
+                try
+                {
+                    var ids = GameObjectLookup.SearchGameObjects("by_path", s, includeInactive: true, maxResults: 1);
+                    if (ids.Count > 0)
+                    {
+                        var byPath = GameObjectLookup.FindById(ids[0]);
+                        if (byPath != null) return byPath;
+                    }
+                }
+                catch { }
             }
 
             // Name-based find (first match, includes inactive)