Merge pull request #1040 from KennerMiner/feature/game-view-uitoolkit-screenshot-capture

fix: include UI Toolkit overlays in game_view screenshots with include_image

Author: Scriptwonder

MCPForUnity/Editor/Tools/ManageScene.cs

@@ -591,22 +591,48 @@ private static object CaptureScreenshot(SceneCommand cmd)
                     }
                 }
 
-                // When a specific camera is requested or include_image is true, always use camera-based capture
-                // (synchronous, gives us bytes in memory for base64).
-                if (targetCamera != null || includeImage)
+                // When include_image is requested but no specific camera, use composited capture
+                // (ScreenCapture.CaptureScreenshotAsTexture) which captures UI Toolkit overlays.
+                // When a specific camera IS requested, use camera-based capture.
+                if (targetCamera != null)
                 {
+                    if (!Application.isBatchMode) EnsureGameView();
+
+                    ScreenshotCaptureResult result = ScreenshotUtility.CaptureFromCameraToAssetsFolder(
+                        targetCamera, fileName, resolvedSuperSize, ensureUniqueFileName: true,
+                        includeImage: includeImage, maxResolution: maxResolution);
+
+                    AssetDatabase.ImportAsset(result.AssetsRelativePath, ImportAssetOptions.ForceSynchronousImport);
+                    string message = $"Screenshot captured to '{result.AssetsRelativePath}' (camera: {targetCamera.name}).";
+                    return new SuccessResponse(message, BuildScreenshotResponseData(result, targetCamera.name, includeImage));
+                }
+
+                if (includeImage && Application.isPlaying)
+                {
+                    if (!Application.isBatchMode) EnsureGameView();
+
+                    ScreenshotCaptureResult result = ScreenshotUtility.CaptureComposited(
+                        fileName, resolvedSuperSize, ensureUniqueFileName: true,
+                        includeImage: true, maxResolution: maxResolution);
+
+                    AssetDatabase.ImportAsset(result.AssetsRelativePath, ImportAssetOptions.ForceSynchronousImport);
+                    string cameraName = Camera.main != null ? Camera.main.name : "composited";
+                    string message = $"Screenshot captured to '{result.AssetsRelativePath}' (camera: {cameraName}).";
+                    return new SuccessResponse(message, BuildScreenshotResponseData(result, cameraName, includeImage: true));
+                }
+
+                if (includeImage)
+                {
+                    // Not in play mode — fall back to camera-based capture
+                    targetCamera = Camera.main;
                     if (targetCamera == null)
                     {
-                        targetCamera = Camera.main;
-                        if (targetCamera == null)
-                        {
-                            var allCams = UnityFindObjectsCompat.FindAll<Camera>();
-                            targetCamera = allCams.Length > 0 ? allCams[0] : null;
-                        }
+                        var allCams = UnityFindObjectsCompat.FindAll<Camera>();
+                        targetCamera = allCams.Length > 0 ? allCams[0] : null;
                     }
                     if (targetCamera == null)
                     {
-                        return new ErrorResponse("No camera found in the scene. Add a Camera to use screenshot with camera or include_image.");
+                        return new ErrorResponse("No camera found in the scene. Add a Camera to use screenshot with include_image outside of Play mode.");
                     }
 
                     if (!Application.isBatchMode) EnsureGameView();
@@ -624,7 +650,6 @@ private static object CaptureScreenshot(SceneCommand cmd)
                     {
                         return new ErrorResponse(ex.Message);
                     }
-
                     if (ScreenshotUtility.IsUnderAssets(result.ProjectRelativePath))
                         AssetDatabase.ImportAsset(result.ProjectRelativePath, ImportAssetOptions.ForceSynchronousImport);
                     string message = $"Screenshot captured to '{result.ProjectRelativePath}' (camera: {targetCamera.name}).";
@@ -644,7 +669,7 @@ private static object CaptureScreenshot(SceneCommand cmd)
                         data["imageWidth"] = result.ImageWidth;
                         data["imageHeight"] = result.ImageHeight;
                     }
-                    return new SuccessResponse(message, data);
+                    return new SuccessResponse(message, BuildScreenshotResponseData(result, targetCamera.name, includeImage));
                 }
 
                 // Default path: use ScreenCapture API if available, camera fallback otherwise
@@ -715,6 +740,31 @@ private static object CaptureScreenshot(SceneCommand cmd)
             }
         }
 
+        private static Dictionary<string, object> BuildScreenshotResponseData(
+            ScreenshotCaptureResult result,
+            string cameraName,
+            bool includeImage)
+        {
+            var data = new Dictionary<string, object>
+            {
+                { "path", result.AssetsRelativePath },
+                { "fullPath", result.FullPath },
+                { "superSize", result.SuperSize },
+                { "isAsync", false },
+                { "camera", cameraName },
+                { "captureSource", "game_view" },
+            };
+
+            if (includeImage && result.ImageBase64 != null)
+            {
+                data["imageBase64"] = result.ImageBase64;
+                data["imageWidth"] = result.ImageWidth;
+                data["imageHeight"] = result.ImageHeight;
+            }
+
+            return data;
+        }
+
         private static object CaptureSceneViewScreenshot(
             SceneCommand cmd,
             string fileName,

MCPForUnity/Runtime/Helpers/ScreenshotUtility.cs

@@ -236,6 +236,84 @@ public static ScreenshotCaptureResult CaptureFromCameraToProjectFolder(
             return result;
         }
 
+        /// <summary>
+        /// Captures a screenshot using ScreenCapture.CaptureScreenshotAsTexture, which captures the
+        /// final composited frame including UI Toolkit overlays, post-processing, etc.
+        /// Falls back to camera-based capture if ScreenCapture is unavailable.
+        /// </summary>
+        public static ScreenshotCaptureResult CaptureComposited(
+            string fileName = null,
+            int superSize = 1,
+            bool ensureUniqueFileName = true,
+            bool includeImage = false,
+            int maxResolution = 0)
+        {
+            if (!IsScreenCaptureModuleAvailable)
+            {
+                var fallbackCamera = FindAvailableCamera();
+                if (fallbackCamera != null)
+                    return CaptureFromCameraToAssetsFolder(fallbackCamera, fileName, superSize, ensureUniqueFileName, includeImage, maxResolution);
+
+                throw new InvalidOperationException("ScreenCapture module is unavailable and no fallback camera found.");
+            }
+
+            ScreenshotCaptureResult result = PrepareCaptureResult(fileName, superSize, ensureUniqueFileName, isAsync: false);
+            Texture2D tex = null;
+            Texture2D downscaled = null;
+            string imageBase64 = null;
+            int imgW = 0, imgH = 0;
+            try
+            {
+                tex = ScreenCapture.CaptureScreenshotAsTexture(result.SuperSize);
+                if (tex == null)
+                {
+                    // Fallback to camera-based if ScreenCapture fails
+                    var cam = FindAvailableCamera();
+                    if (cam != null)
+                        return CaptureFromCameraToAssetsFolder(cam, fileName, superSize, ensureUniqueFileName, includeImage, maxResolution);
+                    throw new InvalidOperationException("ScreenCapture.CaptureScreenshotAsTexture returned null and no fallback camera available.");
+                }
+
+                int width = tex.width;
+                int height = tex.height;
+
+                byte[] png = tex.EncodeToPNG();
+                File.WriteAllBytes(result.FullPath, png);
+
+                if (includeImage)
+                {
+                    int targetMax = maxResolution > 0 ? maxResolution : 640;
+                    if (width > targetMax || height > targetMax)
+                    {
+                        downscaled = DownscaleTexture(tex, targetMax);
+                        byte[] smallPng = downscaled.EncodeToPNG();
+                        imageBase64 = System.Convert.ToBase64String(smallPng);
+                        imgW = downscaled.width;
+                        imgH = downscaled.height;
+                    }
+                    else
+                    {
+                        imageBase64 = System.Convert.ToBase64String(png);
+                        imgW = width;
+                        imgH = height;
+                    }
+                }
+            }
+            finally
+            {
+                DestroyTexture(tex);
+                DestroyTexture(downscaled);
+            }
+
+            if (includeImage && imageBase64 != null)
+            {
+                return new ScreenshotCaptureResult(
+                    result.FullPath, result.AssetsRelativePath, result.SuperSize, false,
+                    imageBase64, imgW, imgH);
+            }
+            return result;
+        }
+
         /// <summary>
         /// Renders a camera to a Texture2D without saving to disk. Used for multi-angle captures.
         /// Returns the base64-encoded PNG, downscaled to fit within <paramref name="maxResolution"/>.