@quenty/studio-bridge 0.21.0 → 0.22.0

package/src/commands/viewport/screenshot/capture-screenshot.lua

@@ -1,205 +1,163 @@
 --!optimize 2
---[[
-	CaptureScreenshot action handler for the studio-bridge plugin.
+--[=[
+	Action that captures the Studio viewport and answers with it as base64.
 
-	Captures the Studio viewport using CaptureService, loads the result into
-	an EditableImage to read raw RGBA pixels, encodes as PNG via png-luau,
-	base64-encodes the PNG via EncodingService, and sends it over the WebSocket.
+	Captures with CaptureService, reads the RGBA pixels through an EditableImage,
+	encodes them as PNG when the plugin supplied an encoder (raw RGBA otherwise),
+	and base64-encodes the result with EncodingService.
 
-	Protocol:
-	  Request:  { type: "captureScreenshot", payload: { format?: "png" } }
-	  Response: { type: "screenshotResult", payload: { data, format, width, height } }
-]]
+	Request: `{ type: "captureScreenshot", payload: { format?: "png" } }`
+	Response: `{ type: "screenshotResult", payload: { data, format, width, height } }`
+
+	@class CaptureScreenshotAction
+]=]
 
--- PNG encoder is injected via router._vendorPng (set by the plugin before
--- actions are pushed). Falls back gracefully to raw RGBA if unavailable.
-local _png = nil
-local _disposed = false
+-- ReadPixelsBuffer reads at most 1024x1024
+local MAX_READ_SIZE = 1024
 
 local CaptureScreenshotAction = {}
 
--- ---------------------------------------------------------------------------
--- Helpers
--- ---------------------------------------------------------------------------
+type SendMessage = (message: { [string]: any }) -> ()
 
-local function buildResponse(sessionId: string, requestId: string?, payload: { [string]: any }): { [string]: any }
-	local msg: { [string]: any } = {
-		type = "screenshotResult",
-		sessionId = sessionId,
-		payload = payload,
-	}
-	if requestId ~= nil and requestId ~= "" then
-		msg.requestId = requestId
-	end
-	return msg
+export type Services = {
+	captureService: any,
+	assetService: any,
+	encodingService: any,
+}
+
+-- The plugin hands its PNG encoder over through router._vendorPng
+local png: any = nil
+local isTornDown = false
+
+--[=[
+	Registers the `captureScreenshot` handler. Capturing replies asynchronously,
+	so it needs `sendMessage`.
+]=]
+function CaptureScreenshotAction.register(router: any, sendMessage: SendMessage?)
+	isTornDown = false
+	png = router._vendorPng
+
+	router:setResponseType("captureScreenshot", "screenshotResult")
+	router:register(
+		"captureScreenshot",
+		function(_payload: { [string]: any }, requestId: string, sessionId: string): { [string]: any }?
+			if not sendMessage then
+				return {
+					code = "CAPABILITY_NOT_SUPPORTED",
+					message = "Screenshot capture requires a sendMessage callback",
+				}
+			end
+
+			-- Lune has no task library, so fall back to running inline there
+			local spawnOk = pcall(function()
+				task.spawn(CaptureScreenshotAction._captureAsync, sendMessage, requestId, sessionId)
+			end)
+			if not spawnOk then
+				CaptureScreenshotAction._captureAsync(sendMessage, requestId, sessionId)
+			end
+
+			-- Replied through sendMessage, so the router must not reply as well
+			return nil
+		end
+	)
 end
 
-local function buildError(sessionId: string, requestId: string?, code: string, message: string): { [string]: any }
-	local msg: { [string]: any } = {
-		type = "error",
-		sessionId = sessionId,
-		payload = { code = code, message = message },
-	}
-	if requestId ~= nil and requestId ~= "" then
-		msg.requestId = requestId
-	end
-	return msg
+--[=[
+	Stops an in-flight capture from replying.
+]=]
+function CaptureScreenshotAction.teardown()
+	isTornDown = true
 end
 
--- ---------------------------------------------------------------------------
--- Core capture logic (runs in a spawned thread)
--- ---------------------------------------------------------------------------
+function CaptureScreenshotAction._captureAsync(
+	sendMessage: SendMessage,
+	requestId: string?,
+	sessionId: string,
+	services: Services?
+)
+	local function fail(code: string, message: string)
+		sendMessage(CaptureScreenshotAction._buildMessage("error", sessionId, requestId, {
+			code = code,
+			message = message,
+		}))
+	end
 
-local function captureAsync(sendMessage: (msg: { [string]: any }) -> (), requestId: string?, sessionId: string)
-	-- Resolve services (deferred so the module loads in Lune test context)
-	local getServiceOk, captureServiceOrErr, assetServiceOrErr, encodingServiceOrErr = pcall(function()
-		return game:GetService("CaptureService"), game:GetService("AssetService"), game:GetService("EncodingService")
+	-- Resolved here rather than at load, so the module still loads under Lune
+	local servicesOk, servicesOrErr = pcall(function()
+		return services
+			or {
+				captureService = game:GetService("CaptureService"),
+				assetService = game:GetService("AssetService"),
+				encodingService = game:GetService("EncodingService"),
+			}
 	end)
-	if not getServiceOk then
-		sendMessage(
-			buildError(sessionId, requestId, "CAPTURE_FAILED", `Required services not available: {captureServiceOrErr}`)
-		)
+	if not servicesOk then
+		fail("CAPTURE_FAILED", `Required services not available: {servicesOrErr}`)
 		return
 	end
-	local CaptureService = captureServiceOrErr
-	local AssetService = assetServiceOrErr
-	local EncodingService = encodingServiceOrErr
-
-	-- Step 1: Capture the viewport via CaptureService
-	local callerThread = coroutine.running()
-	local contentId: string? = nil
-	local captureOk, captureErr = pcall(function()
-		CaptureService:CaptureScreenshot(function(id)
-			contentId = id
-			task.spawn(callerThread)
-		end)
-	end)
+	local assetService = servicesOrErr.assetService
+	local encodingService = servicesOrErr.encodingService
 
+	local captureOk, contentIdOrErr = CaptureScreenshotAction._captureContentIdAsync(servicesOrErr.captureService)
+	if isTornDown then
+		return
+	end
 	if not captureOk then
-		sendMessage(buildError(sessionId, requestId, "CAPTURE_FAILED", `CaptureService failed: {captureErr}`))
+		fail("CAPTURE_FAILED", `CaptureService failed: {contentIdOrErr}`)
 		return
 	end
 
-	-- Wait for the callback
-	coroutine.yield()
-
-	if _disposed then return end
-
-	if not contentId or contentId == "" then
-		sendMessage(buildError(sessionId, requestId, "CAPTURE_FAILED", "CaptureService returned no content ID"))
+	local contentId = contentIdOrErr
+	if contentId == nil or contentId == "" then
+		fail("CAPTURE_FAILED", "CaptureService returned no content ID")
 		return
 	end
 
-	-- Step 2: Load into an EditableImage to read pixels
 	local editableOk, editableImage = pcall(function()
-		return AssetService:CreateEditableImageAsync(Content.fromUri(contentId :: string))
+		return assetService:CreateEditableImageAsync(Content.fromUri(contentId))
 	end)
-
-	if not editableOk or not editableImage then
-		sendMessage(
-			buildError(
-				sessionId,
-				requestId,
-				"EDITABLE_IMAGE_FAILED",
-				`Failed to create EditableImage from capture: {editableImage or "unknown error"}`
-			)
-		)
+	if not (editableOk and editableImage) then
+		fail("EDITABLE_IMAGE_FAILED", `Failed to create EditableImage from capture: {editableImage or "unknown error"}`)
 		return
 	end
 
-	-- Step 3: Read all pixels as RGBA buffer
-	-- ReadPixelsBuffer has a 1024x1024 limit. If the screenshot is larger,
-	-- scale it down into a new EditableImage that fits within the limit.
-	local MAX_SIZE = 1024
-	local srcWidth = editableImage.Size.X
-	local srcHeight = editableImage.Size.Y
-	local width = srcWidth
-	local height = srcHeight
-
-	local readTarget = editableImage
-	local scaledImage = nil
-
-	if srcWidth > MAX_SIZE or srcHeight > MAX_SIZE then
-		local scale = math.min(MAX_SIZE / srcWidth, MAX_SIZE / srcHeight)
-		width = math.floor(srcWidth * scale)
-		height = math.floor(srcHeight * scale)
-
-		local scaleOk, scaleResult = pcall(function()
-			local scaled = AssetService:CreateEditableImage({ Size = Vector2.new(width, height) })
-			scaled:DrawImageTransformed(
-				Vector2.new(width / 2, height / 2), -- center position
-				Vector2.new(scale, scale), -- scale
-				0, -- rotation
-				editableImage
-			)
-			return scaled
-		end)
-
-		if scaleOk and scaleResult then
-			scaledImage = scaleResult
-			readTarget = scaleResult
-		else
-			-- If scaling fails, just read the top-left 1024x1024 corner
-			width = math.min(srcWidth, MAX_SIZE)
-			height = math.min(srcHeight, MAX_SIZE)
-		end
-	end
+	local readTarget, scaledImage, width, height = CaptureScreenshotAction._fitToReadLimit(assetService, editableImage)
 
-	local pixelsOk, pixelsBuf = pcall(function()
+	local pixelsOk, pixels = pcall(function()
 		return readTarget:ReadPixelsBuffer(Vector2.zero, Vector2.new(width, height))
 	end)
 
-	-- Clean up EditableImages immediately to free memory
-	pcall(function()
-		editableImage:Destroy()
-	end)
+	-- Free the images right away, they hold the whole capture in memory
+	pcall(editableImage.Destroy, editableImage)
 	if scaledImage then
-		pcall(function()
-			scaledImage:Destroy()
-		end)
+		pcall(scaledImage.Destroy, scaledImage)
 	end
 
-	if not pixelsOk or not pixelsBuf then
-		sendMessage(
-			buildError(
-				sessionId,
-				requestId,
-				"PIXEL_READ_FAILED",
-				`Failed to read pixels: {pixelsBuf or "unknown error"}`
-			)
-		)
+	if not (pixelsOk and pixels) then
+		fail("PIXEL_READ_FAILED", `Failed to read pixels: {pixels or "unknown error"}`)
 		return
 	end
 
-	-- Step 4: Encode as PNG (if available) or raw RGBA, then base64
-	local dataBuf = pixelsBuf
+	local data = pixels
 	local format = "rgba"
-
-	if _png then
-		local pngOk, pngBuf = pcall(_png.encode, pixelsBuf, { width = width, height = height })
-		if pngOk and pngBuf then
-			dataBuf = pngBuf
+	if png then
+		local pngOk, pngData = pcall(png.encode, pixels, { width = width, height = height })
+		if pngOk and pngData then
+			data = pngData
 			format = "png"
 		end
 	end
 
 	local encodeOk, encoded = pcall(function()
-		return EncodingService:Base64Encode(dataBuf)
+		-- JSONEncode turns a buffer into a { t = "buffer", zbase64 } object, so send a string
+		return buffer.tostring(encodingService:Base64Encode(data))
 	end)
-	if not encodeOk or not encoded then
-		sendMessage(
-			buildError(
-				sessionId,
-				requestId,
-				"ENCODE_FAILED",
-				`Failed to base64 encode data: {encoded or "unknown error"}`
-			)
-		)
+	if not (encodeOk and encoded) then
+		fail("ENCODE_FAILED", `Failed to base64 encode data: {encoded or "unknown error"}`)
 		return
 	end
 
-	-- Step 5: Send the result
-	sendMessage(buildResponse(sessionId, requestId, {
+	sendMessage(CaptureScreenshotAction._buildMessage("screenshotResult", sessionId, requestId, {
 		data = encoded,
 		format = format,
 		width = width,
@@ -207,41 +165,76 @@ local function captureAsync(sendMessage: (msg: { [string]: any }) -> (), request
 	}))
 end
 
--- ---------------------------------------------------------------------------
--- Public API
--- ---------------------------------------------------------------------------
-
-function CaptureScreenshotAction.register(router: any, sendMessage: ((msg: { [string]: any }) -> ())?)
-	_disposed = false
-	router:setResponseType("captureScreenshot", "screenshotResult")
-	-- Capture the PNG encoder from the router's shared deps (set by the plugin)
-	_png = router._vendorPng or nil
-
-	router:register("captureScreenshot", function(_payload: { [string]: any }, requestId: string, sessionId: string): { [string]: any }?
-		if not sendMessage then
-			return {
-				code = "CAPABILITY_NOT_SUPPORTED",
-				message = "Screenshot capture requires a sendMessage callback",
-			}
-		end
+--[[
+	Returns true and the capture's content id, or false and the error
+	CaptureService threw. Yields until the capture callback fires.
+]]
+function CaptureScreenshotAction._captureContentIdAsync(captureService: any): (boolean, any)
+	local callerThread = coroutine.running()
+	local isWaiting = false
+	local isReady = false
+	local contentId: string? = nil
 
-		-- Run asynchronously since CaptureService uses a callback.
-		-- pcall guards against Lune test context where task is not available.
-		local spawnOk = pcall(function()
-			task.spawn(captureAsync, sendMessage, requestId, sessionId)
+	local captureOk, captureErr = pcall(function()
+		captureService:CaptureScreenshot(function(id)
+			contentId = id
+			isReady = true
+			if isWaiting then
+				task.spawn(callerThread)
+			end
 		end)
-		if not spawnOk then
-			-- Fallback: call directly (will fail on GetService in Lune)
-			captureAsync(sendMessage, requestId, sessionId)
-		end
+	end)
+	if not captureOk then
+		return false, captureErr
+	end
+
+	if not isReady then
+		isWaiting = true
+		coroutine.yield()
+	end
+
+	return true, contentId
+end
+
+--[[
+	Returns an image no larger than the read limit, scaled down from the capture
+	when needed. If scaling fails, falls back to the capture's top-left corner.
+]]
+function CaptureScreenshotAction._fitToReadLimit(assetService: any, editableImage: any): (any, any?, number, number)
+	local sourceWidth = editableImage.Size.X
+	local sourceHeight = editableImage.Size.Y
+	if sourceWidth <= MAX_READ_SIZE and sourceHeight <= MAX_READ_SIZE then
+		return editableImage, nil, sourceWidth, sourceHeight
+	end
+
+	local scale = math.min(MAX_READ_SIZE / sourceWidth, MAX_READ_SIZE / sourceHeight)
+	local width = math.floor(sourceWidth * scale)
+	local height = math.floor(sourceHeight * scale)
 
-		-- Return nil so ActionRouter does not generate a wrapped response
-		return nil
+	local scaleOk, scaledImage = pcall(function()
+		local scaled = assetService:CreateEditableImage({ Size = Vector2.new(width, height) })
+		scaled:DrawImageTransformed(Vector2.new(width / 2, height / 2), Vector2.new(scale, scale), 0, editableImage)
+		return scaled
 	end)
+	if scaleOk and scaledImage then
+		return scaledImage, scaledImage, width, height
+	end
+
+	return editableImage, nil, math.min(sourceWidth, MAX_READ_SIZE), math.min(sourceHeight, MAX_READ_SIZE)
 end
 
-function CaptureScreenshotAction.teardown()
-	_disposed = true
+function CaptureScreenshotAction._buildMessage(
+	messageType: string,
+	sessionId: string,
+	requestId: string?,
+	payload: { [string]: any }
+): { [string]: any }
+	return {
+		type = messageType,
+		sessionId = sessionId,
+		requestId = if requestId ~= "" then requestId else nil,
+		payload = payload,
+	}
 end
 
 return CaptureScreenshotAction

package/src/commands/viewport/screenshot/capture-screenshot.spec.lua

@@ -0,0 +1,129 @@
+--!strict
+--[[
+	@class capture-screenshot.spec.lua
+]]
+
+local require = require(script.Parent.loader).load(script)
+
+local EncodingService = game:GetService("EncodingService")
+local HttpService = game:GetService("HttpService")
+
+local ActionRouter = require("ActionRouter")
+local CaptureScreenshotAction = require("capture-screenshot")
+local Jest = require("Jest")
+
+local describe = Jest.Globals.describe
+local expect = Jest.Globals.expect
+local it = Jest.Globals.it
+
+local function setup()
+	CaptureScreenshotAction.register(ActionRouter.new(), function() end)
+
+	local pixels = buffer.fromstring("\255\0\0\255\0\0\255\255")
+	local captureError: string? = nil
+	local callsBackImmediately = false
+
+	local services: CaptureScreenshotAction.Services = {
+		captureService = {
+			CaptureScreenshot = function(_self, onCaptureReady: (string) -> ())
+				if captureError then
+					error(captureError)
+				end
+				if callsBackImmediately then
+					onCaptureReady("rbxassetid://1")
+				else
+					task.defer(onCaptureReady, "rbxassetid://1")
+				end
+			end,
+		},
+		assetService = {
+			CreateEditableImageAsync = function()
+				return {
+					Size = Vector2.new(2, 1),
+					ReadPixelsBuffer = function()
+						return pixels
+					end,
+					Destroy = function() end,
+				}
+			end,
+		},
+		encodingService = EncodingService,
+	}
+
+	return {
+		pixels = pixels,
+		setCaptureError = function(message: string)
+			captureError = message
+		end,
+		setCallsBackImmediately = function()
+			callsBackImmediately = true
+		end,
+		captureAsync = function(): { [string]: any }
+			local sent: { [string]: any }? = nil
+			task.spawn(CaptureScreenshotAction._captureAsync, function(message)
+				sent = message
+			end, "request-1", "session-1", services)
+
+			local deadline = os.clock() + 5
+			while sent == nil and os.clock() < deadline do
+				task.wait()
+			end
+			assert(sent, "captureScreenshot never replied")
+
+			return HttpService:JSONDecode(HttpService:JSONEncode(sent))
+		end,
+	}
+end
+
+describe("captureScreenshot", function()
+	it("should answer CAPABILITY_NOT_SUPPORTED without sendMessage", function()
+		local router = ActionRouter.new()
+		CaptureScreenshotAction.register(router, nil)
+
+		local response = router:Dispatch({ type = "captureScreenshot", sessionId = "session-1", payload = {} })
+
+		expect(response).toEqual({
+			type = "screenshotResult",
+			sessionId = "session-1",
+			payload = {
+				code = "CAPABILITY_NOT_SUPPORTED",
+				message = "Screenshot capture requires a sendMessage callback",
+			},
+		})
+	end)
+
+	it("should send the pixels as a base64 string through JSON", function()
+		local controller = setup()
+
+		local message = controller.captureAsync()
+
+		expect(message.type).toBe("screenshotResult")
+		expect(message.requestId).toBe("request-1")
+		expect(message.payload.format).toBe("rgba")
+		expect(message.payload.width).toBe(2)
+		expect(message.payload.height).toBe(1)
+		expect(typeof(message.payload.data)).toBe("string")
+
+		local decoded = EncodingService:Base64Decode(buffer.fromstring(message.payload.data))
+		expect(buffer.tostring(decoded)).toBe(buffer.tostring(controller.pixels))
+	end)
+
+	it("should finish when CaptureService calls back before returning", function()
+		local controller = setup()
+		controller.setCallsBackImmediately()
+
+		local message = controller.captureAsync()
+
+		expect(message.type).toBe("screenshotResult")
+	end)
+
+	it("should answer CAPTURE_FAILED when CaptureService throws", function()
+		local controller = setup()
+		controller.setCaptureError("no capture permission")
+
+		local message = controller.captureAsync()
+
+		expect(message.type).toBe("error")
+		expect(message.payload.code).toBe("CAPTURE_FAILED")
+	end)
+end)

package/src/commands/viewport/screenshot/screenshot.test.ts

@@ -145,22 +145,19 @@ describe('captureScreenshotHandlerAsync', () => {
     );
   });
 
-  it('handles missing fields gracefully', async () => {
+  it('rejects data that is not a base64 string', async () => {
     const session = {
       captureScreenshotAsync: vi.fn().mockResolvedValue({
-        data: undefined,
+        data: { m: null, t: 'buffer', zbase64: 'KLUv' },
         format: 'png',
-        width: undefined,
-        height: undefined,
+        width: 1,
+        height: 1,
       }),
     } as any;
 
-    const result = await captureScreenshotHandlerAsync(session);
-
-    expect(result.data).toBe('');
-    expect(result.width).toBe(0);
-    expect(result.height).toBe(0);
-    expect(result.summary).toBe('Screenshot captured (0x0)');
+    await expect(captureScreenshotHandlerAsync(session)).rejects.toThrow(
+      'The plugin sent screenshot data as object, not a base64 string'
+    );
   });
 
   it('passes options without affecting capture', async () => {

package/src/commands/viewport/screenshot/screenshot.ts

@@ -25,6 +25,11 @@ export async function captureScreenshotHandlerAsync(
   _options: ScreenshotOptions = {}
 ): Promise<ScreenshotResult> {
   const result: BridgeScreenshotResult = await session.captureScreenshotAsync();
+  if (typeof result.data !== 'string') {
+    throw new Error(
+      `The plugin sent screenshot data as ${typeof result.data}, not a base64 string`
+    );
+  }
 
   let pngBase64: string;
   if (result.format === 'rgba') {
@@ -32,7 +37,7 @@ export async function captureScreenshotHandlerAsync(
     const pngBuffer = rgbaToPng(rgbaBuffer, result.width, result.height);
     pngBase64 = pngBuffer.toString('base64');
   } else {
-    pngBase64 = result.data ?? '';
+    pngBase64 = result.data;
   }
 
   return {