diff --git a/MCPForUnity/Editor/Tools/CommandRegistry.cs b/MCPForUnity/Editor/Tools/CommandRegistry.cs index 7588e2cfb..a77b1f082 100644 --- a/MCPForUnity/Editor/Tools/CommandRegistry.cs +++ b/MCPForUnity/Editor/Tools/CommandRegistry.cs @@ -302,7 +302,18 @@ public static object ExecuteCommand(string commandName, JObject @params, TaskCom throw new InvalidOperationException($"Handler for '{commandName}' does not provide a synchronous implementation"); } - return handlerInfo.SyncHandler(@params); + object result = handlerInfo.SyncHandler(@params); + if (result is Task returnedTask) + { + ExecuteAsyncHandler( + new HandlerInfo(commandName, null, _ => returnedTask), + @params, + commandName, + tcs); + return null; + } + + return result; } /// @@ -332,6 +343,11 @@ public static Task InvokeCommandAsync(string commandName, JObject @param } object result = handlerInfo.SyncHandler(payload); + if (result is Task returnedTask) + { + return returnedTask; + } + return Task.FromResult(result); } diff --git a/MCPForUnity/Editor/Tools/ManageScene.cs b/MCPForUnity/Editor/Tools/ManageScene.cs index c3c66b3fa..e50fe24e5 100644 --- a/MCPForUnity/Editor/Tools/ManageScene.cs +++ b/MCPForUnity/Editor/Tools/ManageScene.cs @@ -2,6 +2,7 @@ using System.Collections.Generic; using System.IO; using System.Linq; +using System.Threading.Tasks; using MCPForUnity.Editor.Helpers; // For Response class using MCPForUnity.Runtime.Helpers; // For ScreenshotUtility using Newtonsoft.Json.Linq; @@ -613,18 +614,7 @@ private static object CaptureScreenshot(SceneCommand cmd) if (includeImage && Application.isPlaying) { if (!Application.isBatchMode) EnsureGameView(); - - string folderOverride = ScreenshotPreferences.Resolve(cmd.outputFolder); - ScreenshotCaptureResult result = ScreenshotUtility.CaptureComposited( - fileName, resolvedSuperSize, ensureUniqueFileName: true, - includeImage: true, maxResolution: maxResolution, - folderOverride: folderOverride); - - if (ScreenshotUtility.IsUnderAssets(result.ProjectRelativePath)) - AssetDatabase.ImportAsset(result.ProjectRelativePath, ImportAssetOptions.ForceSynchronousImport); - string cameraName = Camera.main != null ? Camera.main.name : "composited"; - string message = $"Screenshot captured to '{result.ProjectRelativePath}' (camera: {cameraName})."; - return new SuccessResponse(message, BuildScreenshotResponseData(result, cameraName, includeImage: true)); + return CaptureCompositedScreenshotAsync(cmd, fileName, resolvedSuperSize, maxResolution); } if (includeImage) @@ -756,6 +746,38 @@ private static Dictionary BuildScreenshotResponseData( return data; } + private static async Task CaptureCompositedScreenshotAsync( + SceneCommand cmd, + string fileName, + int resolvedSuperSize, + int maxResolution) + { + string folderOverride = ScreenshotPreferences.Resolve(cmd.outputFolder); + ScreenshotCaptureResult result; + try + { + result = await ScreenshotUtility.CaptureCompositedAsync( + fileName, resolvedSuperSize, ensureUniqueFileName: true, + includeImage: true, maxResolution: maxResolution, + folderOverride: folderOverride).ConfigureAwait(true); + } + catch (TimeoutException ex) + { + return new ErrorResponse(ex.Message); + } + catch (InvalidOperationException ex) + { + return new ErrorResponse(ex.Message); + } + + if (ScreenshotUtility.IsUnderAssets(result.ProjectRelativePath)) + AssetDatabase.ImportAsset(result.ProjectRelativePath, ImportAssetOptions.ForceSynchronousImport); + + string cameraName = Camera.main != null ? Camera.main.name : "composited"; + string message = $"Screenshot captured to '{result.ProjectRelativePath}' (camera: {cameraName})."; + return new SuccessResponse(message, BuildScreenshotResponseData(result, cameraName, includeImage: true)); + } + private static object CaptureSceneViewScreenshot( SceneCommand cmd, string fileName, diff --git a/MCPForUnity/Editor/Tools/ManageUI.cs b/MCPForUnity/Editor/Tools/ManageUI.cs index cad174139..f497516c4 100644 --- a/MCPForUnity/Editor/Tools/ManageUI.cs +++ b/MCPForUnity/Editor/Tools/ManageUI.cs @@ -866,6 +866,14 @@ private static object RenderUI(JObject @params) playFullPath = EnsureUniqueFilePath(playFullPath); string playProjectRelPath = ScreenshotUtility.ToProjectRelativePath(playFullPath); + if (s_pendingCaptureDone && s_pendingCaptureTex == null) + { + s_pendingCaptureDone = false; + s_pendingCaptureStarted = false; + return new ErrorResponse( + "Play-mode screenshot timed out or captured nothing. Keep the Game view visible and the editor unpaused."); + } + // ── Case 1: capture is ready ────────────────────────────────────── if (s_pendingCaptureDone && s_pendingCaptureTex != null) { diff --git a/MCPForUnity/Runtime/Helpers/ScreenshotUtility.cs b/MCPForUnity/Runtime/Helpers/ScreenshotUtility.cs index 79b6a6160..b602e0c13 100644 --- a/MCPForUnity/Runtime/Helpers/ScreenshotUtility.cs +++ b/MCPForUnity/Runtime/Helpers/ScreenshotUtility.cs @@ -2,6 +2,8 @@ using System.Collections.Generic; using System.IO; using System.Linq; +using System.Threading; +using System.Threading.Tasks; using UnityEngine; namespace MCPForUnity.Runtime.Helpers @@ -50,6 +52,7 @@ public static class ScreenshotUtility /// or globally via ScreenshotPreferences in the Editor assembly. /// public const string DefaultFolder = "Assets/Screenshots"; + private static readonly SemaphoreSlim CompositedCaptureGate = new SemaphoreSlim(1, 1); private static Camera FindAvailableCamera() { @@ -166,46 +169,6 @@ public static ScreenshotCaptureResult CaptureFromCameraToProjectFolder( return result; } -#if UNITY_EDITOR - // Synchronously drive a WaitForEndOfFrame ScreenshotCapturer by pumping the editor's - // player loop. Play-mode only; EditorApplication.Step is a no-op in edit mode. - private static Texture2D CaptureCompositedAfterFrame(int superSize, int timeoutSteps = 5) - { - Texture2D result = null; - bool done = false; - bool callerReturned = false; - ScreenshotCapturer.Begin(superSize, tex => - { - // Late completion after the spin loop timed out: caller will never consume - // the texture, so destroy it here to avoid leaking a Unity object. - if (callerReturned) - { - if (tex != null) DestroyTexture(tex); - return; - } - result = tex; - done = true; - }); - // Step() pauses play mode as a side effect; restore the prior state so a screenshot - // doesn't leave a running game paused (an already-paused game stays paused). - bool wasPaused = UnityEditor.EditorApplication.isPaused; - try - { - for (int i = 0; i < timeoutSteps && !done; i++) - { - UnityEditor.EditorApplication.Step(); - } - } - finally - { - if (!wasPaused) - UnityEditor.EditorApplication.isPaused = false; - } - callerReturned = true; - return result; - } -#endif - /// /// Captures a screenshot using ScreenCapture.CaptureScreenshotAsTexture, which captures the /// final composited frame including UI Toolkit overlays, post-processing, etc. @@ -222,67 +185,162 @@ public static ScreenshotCaptureResult CaptureComposited( ScreenshotCaptureResult result = PrepareCaptureResult(fileName, superSize, ensureUniqueFileName, folderOverride: folderOverride, isAsync: false); Texture2D tex = null; Texture2D downscaled = null; - string imageBase64 = null; - int imgW = 0, imgH = 0; try { -#if UNITY_EDITOR - // In play mode, inline ScreenCapture reads a backbuffer before UITK has - // composited; route through WaitForEndOfFrame instead. - tex = Application.isPlaying - ? CaptureCompositedAfterFrame(result.SuperSize) - : ScreenCapture.CaptureScreenshotAsTexture(result.SuperSize); -#else + // Direct capture is safe in edit mode. Play-mode MCP callers must use + // CaptureCompositedAsync so WaitForEndOfFrame can run without + // EditorApplication.Step re-entering the PlayerLoop. tex = ScreenCapture.CaptureScreenshotAsTexture(result.SuperSize); -#endif if (tex == null) { - // Fallback to camera-based if ScreenCapture fails - var cam = FindAvailableCamera(); - if (cam != null) - return CaptureFromCameraToProjectFolder(cam, fileName, superSize, ensureUniqueFileName, - includeImage, maxResolution, folderOverride: folderOverride); - throw new InvalidOperationException("ScreenCapture.CaptureScreenshotAsTexture returned null and no fallback camera available."); + return CaptureCompositedOrCameraFallback( + fileName, superSize, ensureUniqueFileName, includeImage, maxResolution, folderOverride); } - int width = tex.width; - int height = tex.height; + return EncodeAndSaveComposited(tex, result, includeImage, maxResolution, ref downscaled); + } + finally + { + DestroyTexture(tex); + DestroyTexture(downscaled); + } + } - byte[] png = tex.EncodeToPNG(); - File.WriteAllBytes(result.FullPath, png); + /// + /// Play-mode composited capture that waits for end-of-frame without pumping + /// EditorApplication.Step. MCP commands run inside + /// UnitySynchronizationContext.ExecuteTasks, so a synchronous Step() + /// re-enters the PlayerLoop and can flood Editor.log until the Editor dies. + /// + public static async Task CaptureCompositedAsync( + string fileName = null, + int superSize = 1, + bool ensureUniqueFileName = true, + bool includeImage = false, + int maxResolution = 0, + string folderOverride = null) + { + if (!await CompositedCaptureGate + .WaitAsync(TimeSpan.FromSeconds(ScreenshotCapturer.DefaultTimeoutSeconds * 4)) + .ConfigureAwait(true)) + { + throw new TimeoutException( + "Another composited screenshot capture is still in progress. Retry shortly."); + } + try + { + return await CaptureCompositedAsyncUngated( + fileName, superSize, ensureUniqueFileName, includeImage, maxResolution, folderOverride) + .ConfigureAwait(true); + } + finally + { + CompositedCaptureGate.Release(); + } + } - if (includeImage) + private static Task CaptureCompositedAsyncUngated( + string fileName, + int superSize, + bool ensureUniqueFileName, + bool includeImage, + int maxResolution, + string folderOverride) + { + var prepared = PrepareCaptureResult(fileName, superSize, ensureUniqueFileName, folderOverride: folderOverride, isAsync: false); + var tcs = new TaskCompletionSource( + TaskCreationOptions.RunContinuationsAsynchronously); + + ScreenshotCapturer.Begin(prepared.SuperSize, (tex, timedOut) => + { + Texture2D downscaled = null; + try { - int targetMax = maxResolution > 0 ? maxResolution : 640; - if (width > targetMax || height > targetMax) + if (timedOut) { - downscaled = DownscaleTexture(tex, targetMax); - byte[] smallPng = downscaled.EncodeToPNG(); - imageBase64 = System.Convert.ToBase64String(smallPng); - imgW = downscaled.width; - imgH = downscaled.height; + tcs.TrySetException(new TimeoutException( + "Play-mode screenshot timed out waiting for end of frame. Keep the Game view visible and the editor unpaused.")); + return; } - else + + if (tex == null) { - imageBase64 = System.Convert.ToBase64String(png); - imgW = width; - imgH = height; + tcs.TrySetResult(CaptureCompositedOrCameraFallback( + fileName, superSize, ensureUniqueFileName, includeImage, maxResolution, folderOverride)); + return; } + + tcs.TrySetResult(EncodeAndSaveComposited(tex, prepared, includeImage, maxResolution, ref downscaled)); } - } - finally + catch (Exception ex) + { + tcs.TrySetException(ex); + } + finally + { + DestroyTexture(tex); + DestroyTexture(downscaled); + } + }); + + return tcs.Task; + } + + private static ScreenshotCaptureResult CaptureCompositedOrCameraFallback( + string fileName, + int superSize, + bool ensureUniqueFileName, + bool includeImage, + int maxResolution, + string folderOverride) + { + var cam = FindAvailableCamera(); + if (cam != null) { - DestroyTexture(tex); - DestroyTexture(downscaled); + return CaptureFromCameraToProjectFolder(cam, fileName, superSize, ensureUniqueFileName, + includeImage, maxResolution, folderOverride: folderOverride); } - if (includeImage && imageBase64 != null) + throw new InvalidOperationException( + "ScreenCapture.CaptureScreenshotAsTexture returned null and no fallback camera available."); + } + + private static ScreenshotCaptureResult EncodeAndSaveComposited( + Texture2D tex, + ScreenshotCaptureResult prepared, + bool includeImage, + int maxResolution, + ref Texture2D downscaled) + { + int width = tex.width; + int height = tex.height; + byte[] png = tex.EncodeToPNG(); + File.WriteAllBytes(prepared.FullPath, png); + + if (!includeImage) + return prepared; + + int targetMax = maxResolution > 0 ? maxResolution : 640; + string imageBase64; + int imgW; + int imgH; + if (width > targetMax || height > targetMax) { - return new ScreenshotCaptureResult( - result.FullPath, result.ProjectRelativePath, result.SuperSize, false, - imageBase64, imgW, imgH); + downscaled = DownscaleTexture(tex, targetMax); + imageBase64 = Convert.ToBase64String(downscaled.EncodeToPNG()); + imgW = downscaled.width; + imgH = downscaled.height; } - return result; + else + { + imageBase64 = Convert.ToBase64String(png); + imgW = width; + imgH = height; + } + + return new ScreenshotCaptureResult( + prepared.FullPath, prepared.ProjectRelativePath, prepared.SuperSize, false, + imageBase64, imgW, imgH); } /// @@ -754,29 +812,120 @@ private static string GetProjectRootPath() /// /// Transient MonoBehaviour that yields WaitForEndOfFrame, calls /// ScreenCapture.CaptureScreenshotAsTexture, invokes the callback, and self-destructs. + /// Times out via the editor update loop so a paused or unfocused PlayerLoop cannot leak + /// hidden capturer objects for the rest of the session. /// public sealed class ScreenshotCapturer : MonoBehaviour { + public const float DefaultTimeoutSeconds = 2f; + private int _superSize = 1; - private Action _onComplete; + private Action _onComplete; + private float _timeoutSeconds = DefaultTimeoutSeconds; + private float _startedAt; + private bool _finished; + private bool _destroying; + + /// Spawns a hidden GameObject, attaches a capturer, returns immediately. + public static ScreenshotCapturer Begin(int superSize, Action onComplete, float timeoutSeconds = DefaultTimeoutSeconds) + { + return Begin(superSize, (tex, _) => onComplete?.Invoke(tex), timeoutSeconds); + } /// Spawns a hidden GameObject, attaches a capturer, returns immediately. - public static void Begin(int superSize, Action onComplete) + public static ScreenshotCapturer Begin(int superSize, Action onComplete, float timeoutSeconds = DefaultTimeoutSeconds) { var go = new GameObject("__MCP_ScreenshotCapturer__") { hideFlags = HideFlags.HideAndDontSave }; var c = go.AddComponent(); c._superSize = Mathf.Max(1, superSize); c._onComplete = onComplete; + c._timeoutSeconds = Mathf.Max(0.05f, timeoutSeconds); + c._startedAt = Time.realtimeSinceStartup; + c.ArmTimeout(); + return c; } + private void ArmTimeout() + { +#if UNITY_EDITOR + UnityEditor.EditorApplication.update += TickTimeout; +#else + StartCoroutine(TimeoutWatch()); +#endif + } + + private void OnDestroy() + { + _destroying = true; + DisarmTimeout(); + if (!_finished) + Complete(null, timedOut: true); + } + + private void DisarmTimeout() + { +#if UNITY_EDITOR + UnityEditor.EditorApplication.update -= TickTimeout; +#endif + } + +#if UNITY_EDITOR + private void TickTimeout() + { + if (_finished) return; + if (Time.realtimeSinceStartup - _startedAt < _timeoutSeconds) return; + Complete(null, timedOut: true); + } +#else + private System.Collections.IEnumerator TimeoutWatch() + { + yield return new WaitForSecondsRealtime(_timeoutSeconds); + if (!_finished) + Complete(null, timedOut: true); + } +#endif + private System.Collections.IEnumerator Start() { yield return new WaitForEndOfFrame(); + if (_finished) yield break; + Texture2D tex = null; try { tex = ScreenCapture.CaptureScreenshotAsTexture(_superSize); } catch (Exception ex) { Debug.LogError($"[MCP for Unity] CaptureScreenshotAsTexture failed: {ex.Message}"); } - _onComplete?.Invoke(tex); - Destroy(gameObject); + Complete(tex, timedOut: false); + } + + private void Complete(Texture2D tex, bool timedOut) + { + if (_finished) + { + if (tex != null) + { + if (Application.isPlaying) + Destroy(tex); + else + DestroyImmediate(tex); + } + return; + } + _finished = true; + DisarmTimeout(); + try + { + _onComplete?.Invoke(tex, timedOut); + } + finally + { + if (!_destroying) + { +#if UNITY_EDITOR + DestroyImmediate(gameObject); +#else + Destroy(gameObject); +#endif + } + } } } } diff --git a/TestProjects/UnityMCPTests/Assets/Tests/EditMode/Helpers/ScreenshotCapturerTests.cs b/TestProjects/UnityMCPTests/Assets/Tests/EditMode/Helpers/ScreenshotCapturerTests.cs new file mode 100644 index 000000000..61654f48d --- /dev/null +++ b/TestProjects/UnityMCPTests/Assets/Tests/EditMode/Helpers/ScreenshotCapturerTests.cs @@ -0,0 +1,61 @@ +using System.Collections; +using NUnit.Framework; +using UnityEngine; +using UnityEngine.TestTools; +using MCPForUnity.Runtime.Helpers; + +namespace MCPForUnityTests.Editor.Helpers +{ + public class ScreenshotCapturerTests + { + [TearDown] + public void TearDown() + { + foreach (var capturer in Resources.FindObjectsOfTypeAll()) + { + if (capturer != null) + Object.DestroyImmediate(capturer.gameObject); + } + } + + [UnityTest] + public IEnumerator Begin_DoesNotLeakCapturerWhenFrameNeverCompletes() + { + LogAssert.ignoreFailingMessages = true; + + bool called = false; + var capturer = ScreenshotCapturer.Begin(1, _ => called = true, timeoutSeconds: 0.15f); + Assert.IsNotNull(capturer, "Begin should return the live capturer."); + + float deadline = Time.realtimeSinceStartup + 2f; + while (!called && Time.realtimeSinceStartup < deadline) + yield return null; + + Assert.IsTrue(called, "Capturer must complete even if WaitForEndOfFrame never resumes."); + yield return null; + + Assert.IsTrue(capturer == null, "Hidden __MCP_ScreenshotCapturer__ must destroy itself after completion."); + Assert.AreEqual(0, Resources.FindObjectsOfTypeAll().Length); + } + + [Test] + public void Destroy_CompletesPendingCallback() + { + bool called = false; + Texture2D received = null; + + var capturer = ScreenshotCapturer.Begin(1, tex => + { + received = tex; + called = true; + }, timeoutSeconds: 5f); + + Assert.IsNotNull(capturer); + Object.DestroyImmediate(capturer.gameObject); + + Assert.IsTrue(called, "Destroying the capturer must complete the waiter so MCP commands cannot hang."); + Assert.IsNull(received); + Assert.AreEqual(0, Resources.FindObjectsOfTypeAll().Length); + } + } +} diff --git a/TestProjects/UnityMCPTests/Assets/Tests/EditMode/Helpers/ScreenshotCapturerTests.cs.meta b/TestProjects/UnityMCPTests/Assets/Tests/EditMode/Helpers/ScreenshotCapturerTests.cs.meta new file mode 100644 index 000000000..33db75ebe --- /dev/null +++ b/TestProjects/UnityMCPTests/Assets/Tests/EditMode/Helpers/ScreenshotCapturerTests.cs.meta @@ -0,0 +1,11 @@ +fileFormatVersion: 2 +guid: b8e4c91a2d7f4a3e9c5b1f0e8d7a6c5b +MonoImporter: + externalObjects: {} + serializedVersion: 2 + defaultReferences: [] + executionOrder: 0 + icon: {instanceID: 0} + userData: + assetBundleName: + assetBundleVariant: