Skip to content
Open
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
215 changes: 201 additions & 14 deletions unity-connector/Editor/Tools/EditorScreenshot.cs
Original file line number Diff line number Diff line change
@@ -1,5 +1,7 @@
using System;
using System.IO;
using System.Linq;
using System.Threading.Tasks;
using Newtonsoft.Json.Linq;
using UnityEditor;
using UnityEngine;
Expand All @@ -12,22 +14,36 @@ public static class EditorScreenshot
private const int DefaultWidth = 1920;
private const int DefaultHeight = 1080;

// Upper bounds on the wait for ScreenCapture to flush the PNG before we fall back
// to the multi-camera composite. Bounded by BOTH frames and wall-clock so a stalled
// Game view can't hold CommandRouter's global lock for an unbounded time when the
// editor is background-throttled.
private const int MaxCaptureFrames = 120;
private const double MaxCaptureSeconds = 3.0;

// Trailing bytes of every PNG: the IEND chunk type + its (constant) CRC. Their
// presence at end-of-file is our signal that ScreenCapture finished writing.
private static readonly byte[] PngIendSignature = { 0x49, 0x45, 0x4E, 0x44, 0xAE, 0x42, 0x60, 0x82 };

public class Parameters
{
[ToolParameter("View to capture: scene (default), game", Required = false)]
public string View { get; set; }

[ToolParameter("Override width (default 1920)", Required = false)]
[ToolParameter("Override width (default 1920). Game view: applies to the multi-camera fallback only.", Required = false)]
public int Width { get; set; }

[ToolParameter("Override height (default 1080)", Required = false)]
[ToolParameter("Override height (default 1080). Game view: applies to the multi-camera fallback only.", Required = false)]
public int Height { get; set; }

[ToolParameter("Game view supersize multiplier 1-4 (default 1) for the primary ScreenCapture path: captures at native game-view resolution x this. Ignored by the multi-camera fallback.", Required = false)]
public int SuperSize { get; set; }

[ToolParameter("Output file path, absolute or relative to project root (default: Screenshots/screenshot.png)", Required = false)]
public string OutputPath { get; set; }
}

public static object HandleCommand(JObject @params)
public static async Task<object> HandleCommand(JObject @params)
{
if (@params == null)
@params = new JObject();
Expand All @@ -36,6 +52,7 @@ public static object HandleCommand(JObject @params)
var view = p.Get("view", "scene").ToLowerInvariant();
var width = p.GetInt("width", DefaultWidth).Value;
var height = p.GetInt("height", DefaultHeight).Value;
var superSize = Mathf.Clamp(p.GetInt("super_size", 1).Value, 1, 4);
var outputPath = ResolveOutputPath(p.Get("output_path"));

try
Expand All @@ -49,7 +66,7 @@ public static object HandleCommand(JObject @params)
case "scene":
return CaptureSceneView(width, height, outputPath);
case "game":
return CaptureGameView(width, height, outputPath);
return await CaptureGameView(width, height, superSize, outputPath);
default:
return new ErrorResponse($"Unknown view '{view}'. Valid: scene, game.");
}
Expand Down Expand Up @@ -82,29 +99,199 @@ private static object CaptureSceneView(int width, int height, string outputPath)
if (!camera)
return new ErrorResponse("SceneView camera is null.");

return CaptureCamera(camera, width, height, outputPath);
return CaptureSingleCamera(camera, width, height, outputPath);
}

private static object CaptureGameView(int width, int height, string outputPath)
// Capture exactly what the Game view shows: every stacked camera plus ALL UI,
// including Screen Space - Overlay canvases that no single camera.Render() can
// reproduce (overlay UI is composited onto the back buffer after camera render).
// ScreenCapture.CaptureScreenshot grabs that composited back buffer at end of
// frame, so we queue it and wait for the PNG to land. If no frame is produced
// (no Game view rendering, batch mode, etc.) we fall back to compositing the
// cameras ourselves - which the fallback's response explicitly flags as lossy.
private static async Task<object> CaptureGameView(int width, int height, int superSize, string outputPath)
{
var camera = Camera.main;
if (!camera)
try
{
if (File.Exists(outputPath))
File.Delete(outputPath);
}
catch { /* a stale file just means the completion check can't assume freshness */ }

ScreenCapture.CaptureScreenshot(outputPath, superSize);

if (await WaitForScreenshotFile(outputPath, MaxCaptureFrames, MaxCaptureSeconds))
{
return new SuccessResponse($"Screenshot saved to {outputPath}",
new { path = outputPath, mode = "screencapture", superSize });
}

// No frame rendered in time - composite the cameras directly instead.
return CaptureGameViewComposite(width, height, outputPath);
}

// Awaits editor frames until ScreenCapture has fully flushed the PNG (IEND written),
// or until we hit the frame/time budget, nudging the player loop each tick so a frame
// actually renders even in edit mode. Resolves false (-> fallback) on a domain reload
// so an in-flight capture can never leak the update subscription or hang the request.
private static Task<bool> WaitForScreenshotFile(string path, int maxFrames, double maxSeconds)
{
var tcs = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
var sw = System.Diagnostics.Stopwatch.StartNew();
var frames = 0;

void Cleanup()
{
EditorApplication.update -= OnUpdate;
AssemblyReloadEvents.beforeAssemblyReload -= OnBeforeReload;
}

void Finish(bool result)
{
Cleanup();
tcs.TrySetResult(result);
}

void OnBeforeReload() => Finish(false);

void OnUpdate()
{
frames++;

if (IsPngComplete(path))
{
Finish(true);
return;
}

if (frames >= maxFrames || sw.Elapsed.TotalSeconds >= maxSeconds)
{
Finish(false);
return;
}

EditorApplication.QueuePlayerLoopUpdate();
}

EditorApplication.update += OnUpdate;
AssemblyReloadEvents.beforeAssemblyReload += OnBeforeReload;
EditorApplication.QueuePlayerLoopUpdate();
return tcs.Task;
}

// True once the file exists and ends with a PNG IEND chunk, i.e. ScreenCapture has
// finished writing it. Guards against returning a truncated image: ScreenCapture's
// write is not atomic, so a plain length-greater-than-zero check can observe a
// partially written file (especially at large super_size / on slow disks).
private static bool IsPngComplete(string path)
{
try
{
using (var fs = new FileStream(path, FileMode.Open, FileAccess.Read, FileShare.ReadWrite))
{
if (fs.Length < PngIendSignature.Length)
return false;

fs.Seek(-PngIendSignature.Length, SeekOrigin.End);

var tail = new byte[PngIendSignature.Length];
var read = 0;
while (read < tail.Length)
{
var n = fs.Read(tail, read, tail.Length - read);
if (n <= 0)
break;
read += n;
}
if (read < tail.Length)
return false;

for (var i = 0; i < PngIendSignature.Length; i++)
if (tail[i] != PngIendSignature[i])
return false;

return true;
}
}
catch
{
// File missing or still locked exclusively by the writer - not ready yet.
return false;
}
}

// Fallback: composite every enabled screen-targeting camera in depth order into a
// single render texture, honoring each camera's clear flags - this mirrors how the
// built-in render pipeline stacks cameras onto the screen. Captures multi-camera
// rigs and Screen Space - Camera / World Space UI, but NOT Screen Space - Overlay
// UI (only the primary ScreenCapture path can reproduce that) - so the response
// says so out loud. Note: camera.Render() is a built-in-pipeline API; under URP/HDRP
// this per-camera composite is unreliable and the primary path should be relied on.
private static object CaptureGameViewComposite(int width, int height, string outputPath)
{
var cameras = Camera.allCameras
.Where(c => c && c.enabled && !c.targetTexture)
.OrderBy(c => c.depth)
.ToList();

if (cameras.Count == 0)
{
var single = Camera.main;
#if UNITY_2023_1_OR_NEWER
camera = UnityEngine.Object.FindFirstObjectByType<Camera>();
if (!single) single = UnityEngine.Object.FindFirstObjectByType<Camera>();
#else
camera = UnityEngine.Object.FindObjectOfType<Camera>();
if (!single) single = UnityEngine.Object.FindObjectOfType<Camera>();
#endif
if (!camera)
if (!single)
return new ErrorResponse("No camera found in scene.");
cameras.Add(single);
}

return CaptureCamera(camera, width, height, outputPath);
RenderTexture rt = null;
Texture2D tex = null;
var previousActive = RenderTexture.active;
var previousTargets = new RenderTexture[cameras.Count];

try
{
rt = new RenderTexture(width, height, 24);

for (var i = 0; i < cameras.Count; i++)
{
previousTargets[i] = cameras[i].targetTexture;
cameras[i].targetTexture = rt;
cameras[i].Render();
}

RenderTexture.active = rt;
tex = new Texture2D(width, height, TextureFormat.RGB24, false);
tex.ReadPixels(new Rect(0, 0, width, height), 0, 0);
tex.Apply();

File.WriteAllBytes(outputPath, tex.EncodeToPNG());

return new SuccessResponse(
$"Screenshot saved to {outputPath} - WARNING: Game view produced no frame in time, " +
"so this is a multi-camera composite that OMITS Screen Space - Overlay UI. Run in play " +
"mode with a rendering Game view for a full capture.",
new { path = outputPath, width, height, mode = "composite", cameras = cameras.Count });
}
finally
{
for (var i = 0; i < cameras.Count; i++)
if (cameras[i]) cameras[i].targetTexture = previousTargets[i];
RenderTexture.active = previousActive;
if (rt) UnityEngine.Object.DestroyImmediate(rt);
if (tex) UnityEngine.Object.DestroyImmediate(tex);
}
}

private static object CaptureCamera(Camera camera, int width, int height, string outputPath)
// Render a single camera to an offscreen render texture. Used for the Scene view
// (which is genuinely one editor camera) and as a primitive for the composite.
private static object CaptureSingleCamera(Camera camera, int width, int height, string outputPath)
{
var previousRT = camera.targetTexture;
var previousActive = RenderTexture.active;
RenderTexture rt = null;
Texture2D tex = null;

Expand All @@ -127,7 +314,7 @@ private static object CaptureCamera(Camera camera, int width, int height, string
finally
{
camera.targetTexture = previousRT;
RenderTexture.active = null;
RenderTexture.active = previousActive;
if (rt) UnityEngine.Object.DestroyImmediate(rt);
if (tex) UnityEngine.Object.DestroyImmediate(tex);
}
Expand Down