@quenty/studio-bridge 0.21.1 → 0.22.0

package/src/commands/console/exec/execute.lua

@@ -1,317 +1,433 @@
---[[
-	Execute action handler for the studio-bridge plugin.
-
-	Compiles and runs Luau code received from the server, returning a
-	scriptComplete response on success or an error response on failure.
-
-	Supports:
-	  - requestId correlation: if present in the request, echoed in all
-	    response messages (scriptComplete and output).
-	  - Distinct error codes: SCRIPT_LOAD_ERROR, SCRIPT_RUNTIME_ERROR.
-	  - Sequential queueing: concurrent execute requests are processed
-	    one at a time in FIFO order.
-	  - Return values: everything the script returns is marshalled onto
-	    `payload.returnValues`, so a caller can read a result as a value
-	    instead of scraping it back out of printed output.
-
-	This module has no Roblox dependencies and is testable under Lune.
-]]
+--!strict
+--[=[
+	Action that compiles and runs Luau code sent by the bridge, answering with a
+	`scriptComplete` message.
+
+	- Echoes the request's `requestId` on every response.
+	- Fails with `SCRIPT_LOAD_ERROR` when the code does not compile,
+	  `SCRIPT_RUNTIME_ERROR` with a `traceback` when it throws, and
+	  `SCRIPT_TIMEOUT` when it outlives `payload.timeoutMs`. A timed out script
+	  is cancelled, so the requests queued behind it still run.
+	- Passes `payload.args`, a list of strings, to the script as `...`.
+	- Runs concurrent requests one at a time, in arrival order.
+	- Returns everything printed while the script runs, read from LogService, so
+	  output from required modules and from threads the script starts comes back.
+	- Marshals everything the script returns onto `payload.returnValues`, so a
+	  caller reads a result as a value instead of scraping it out of the output.
+
+	Loaded standalone inside the plugin, so it cannot require anything.
+
+	Compiles through the router's `Compile`, which still works in a playtest's
+	client where `loadstring` does not exist. Routers from older plugins have no
+	`Compile`, and fall back to `loadstring`.
+
+	@class ExecuteAction
+]=]
 
-local ExecuteAction = {}
+local MAX_DEPTH = 64
+local CHUNK_NAME = "=exec"
+local PLUGIN_PREFIX = "[StudioBridge]"
+-- Tracebacks name this module's frames the way its first line does
+local OWN_FRAME_PREFIX = string.match(debug.traceback(), "^([^\n]-):%d+")
 
--- ---------------------------------------------------------------------------
--- Internal queue for sequential execution
--- ---------------------------------------------------------------------------
+local ExecuteAction = {}
 
-local _queue: { { payload: { [string]: any }, requestId: string?, sessionId: string, sendMessage: ((msg: { [string]: any }) -> ())? } } =
-	{}
-local _processing = false
+type SendMessage = (message: { [string]: any }) -> ()
 
--- ---------------------------------------------------------------------------
--- Return value serialization
--- ---------------------------------------------------------------------------
+type QueuedRequest = {
+	payload: { [string]: any },
+	requestId: string?,
+	sessionId: string,
+	sendMessage: SendMessage?,
+	receivedAt: number,
+}
+
+type LogEntry = {
+	level: string,
+	body: string,
+	timestamp: number,
+}
+
+local LOG_LEVELS: { [Enum.MessageType]: string } = {
+	[Enum.MessageType.MessageOutput] = "Print",
+	[Enum.MessageType.MessageInfo] = "Info",
+	[Enum.MessageType.MessageWarning] = "Warning",
+	[Enum.MessageType.MessageError] = "Error",
+}
+
+local queue: { QueuedRequest } = {}
+local isProcessing = false
+local compile: (source: string, chunkName: string?) -> (any, string?) = loadstring :: any
+
+--[=[
+	Registers the `execute` handler. With `sendMessage`, requests queue and each
+	reply is sent directly; without it, the reply goes back through the router.
+]=]
+function ExecuteAction.register(router: any, sendMessage: SendMessage?)
+	if type(router.Compile) == "function" then
+		compile = function(source: string, chunkName: string?)
+			return router:Compile(source, chunkName)
+		end
+	end
 
--- Produces the SerializedReturnValue shapes declared in
--- web-socket-protocol.ts. Action modules are loadstring'd standalone inside
--- the plugin and cannot require anything, so this cannot be shared with the
--- equivalent marshaller in query-data-model.lua.
+	router:setResponseType("execute", "scriptComplete")
+	router:register("execute", function(payload: { [string]: any }, requestId: string, sessionId: string)
+		if not sendMessage then
+			return ExecuteAction._handleExecute(payload, requestId, sessionId, nil)
+		end
 
-local MAX_DEPTH = 64
+		table.insert(queue, {
+			payload = payload,
+			requestId = requestId,
+			sessionId = sessionId,
+			sendMessage = sendMessage,
+			receivedAt = os.clock(),
+		})
+		ExecuteAction._processQueue()
 
-local function unsupported(typeName: string, text: string): { [string]: any }
-	return { type = "Unsupported", typeName = typeName, toString = text }
+		-- Replied through sendMessage, so the router must not reply as well
+		return nil
+	end)
 end
 
-local serializeValue: (value: any, depth: number, seen: { [any]: boolean }) -> any
-
--- An array-like table becomes a JSON array, anything else a JSON object with
--- string keys. Mixed tables take the object branch on purpose: JSONEncode
--- rejects them, and a rejected encode drops the whole scriptComplete message.
-local function serializeTable(value: { [any]: any }, depth: number, seen: { [any]: boolean }): any
-	local count = 0
-	for _ in value do
-		count += 1
-	end
-
-	if count == #value then
-		local array = {}
-		for index = 1, count do
-			array[index] = serializeValue(value[index], depth + 1, seen)
+--[=[
+	Cancels every queued request, replying `ACTION_REPLACED` to each.
+]=]
+function ExecuteAction.teardown()
+	for _, request in queue do
+		if request.sendMessage then
+			request.sendMessage({
+				type = "scriptComplete",
+				sessionId = request.sessionId,
+				requestId = request.requestId,
+				payload = {
+					success = false,
+					error = "Action re-registered, request cancelled",
+					code = "ACTION_REPLACED",
+				},
+			})
 		end
-		return array
 	end
 
-	local map: { [string]: any } = {}
-	for key, item in value do
-		map[if type(key) == "string" then key else tostring(key)] = serializeValue(item, depth + 1, seen)
-	end
-	return map
+	ExecuteAction._resetQueue()
 end
 
-function serializeValue(value: any, depth: number, seen: { [any]: boolean }): any
-	local valueType = typeof(value)
+--[=[
+	Runs an execute request and returns its response payload, bypassing the queue.
+]=]
+function ExecuteAction.handleExecute(
+	payload: { [string]: any },
+	requestId: string?,
+	sessionId: string
+): { [string]: any }?
+	return ExecuteAction._handleExecute(payload, requestId, sessionId, nil)
+end
 
-	if valueType == "nil" or valueType == "string" or valueType == "boolean" then
-		return value
-	elseif valueType == "number" then
-		-- inf and nan have no JSON spelling, and encoding one produces a
-		-- document the server cannot decode -- losing the entire message.
-		if value ~= value or value == math.huge or value == -math.huge then
-			return unsupported("number", tostring(value))
-		end
-		return value
-	elseif valueType == "table" then
-		if depth > MAX_DEPTH then
-			return unsupported("table", "<max depth exceeded>")
-		end
-		if seen[value] then
-			return unsupported("table", "<cycle>")
-		end
-		seen[value] = true
-		local serialized = serializeTable(value, depth, seen)
-		seen[value] = nil
-		return serialized
-	elseif valueType == "Vector3" then
-		return { type = "Vector3", value = { value.X, value.Y, value.Z } }
-	elseif valueType == "Vector2" then
-		return { type = "Vector2", value = { value.X, value.Y } }
-	elseif valueType == "CFrame" then
-		return { type = "CFrame", value = { value:GetComponents() } }
-	elseif valueType == "Color3" then
-		return { type = "Color3", value = { value.R, value.G, value.B } }
-	elseif valueType == "UDim2" then
-		return { type = "UDim2", value = { value.X.Scale, value.X.Offset, value.Y.Scale, value.Y.Offset } }
-	elseif valueType == "UDim" then
-		return { type = "UDim", value = { value.Scale, value.Offset } }
-	elseif valueType == "BrickColor" then
-		return { type = "BrickColor", name = value.Name, value = value.Number }
-	elseif valueType == "EnumItem" then
-		return { type = "EnumItem", enum = tostring(value.EnumType), name = value.Name, value = value.Value }
-	elseif valueType == "Instance" then
-		return { type = "Instance", className = value.ClassName, path = value:GetFullName() }
-	else
-		return unsupported(valueType, tostring(value))
-	end
+function ExecuteAction._resetQueue()
+	queue = {}
+	isProcessing = false
 end
 
--- Marshal the values a script returned, in order. A returned nil becomes an
--- explicit marker: a JSON array cannot hold a hole, and dropping the entry
--- would shift every later value into the wrong position.
-local function serializeReturnValues(packed: { n: number, [number]: any }): { any }
-	local returnValues = {}
-	for index = 2, packed.n do
-		local serialized = serializeValue(packed[index], 1, {})
-		if serialized == nil then
-			serialized = { type = "Nil" }
-		end
-		returnValues[index - 1] = serialized
+function ExecuteAction._processQueue()
+	if isProcessing then
+		return
+	end
+	isProcessing = true
+
+	while #queue > 0 do
+		local request = table.remove(queue, 1) :: QueuedRequest
+		ExecuteAction._handleExecute(
+			request.payload,
+			request.requestId,
+			request.sessionId,
+			request.sendMessage,
+			request.receivedAt
+		)
 	end
-	return returnValues
-end
 
--- ---------------------------------------------------------------------------
--- Core execution logic
--- ---------------------------------------------------------------------------
+	isProcessing = false
+end
 
--- Handle an execute request by compiling and running the provided code.
--- Captures print/warn output by temporarily intercepting the globals.
 function ExecuteAction._handleExecute(
 	payload: { [string]: any },
 	requestId: string?,
 	sessionId: string,
-	sendMessage: ((msg: { [string]: any }) -> ())?
+	sendMessage: SendMessage?,
+	receivedAt: number?
 ): { [string]: any }?
-	local code = payload and (payload.script or payload.code)
+	-- An empty requestId means none was sent
+	local effectiveRequestId = if requestId ~= "" then requestId else nil
 
-	-- Normalize requestId: treat empty string as absent (v1 compatibility)
-	local effectiveRequestId: string? = nil
-	if requestId ~= nil and requestId ~= "" then
-		effectiveRequestId = requestId
-	end
+	local function reply(responsePayload: { [string]: any }): { [string]: any }?
+		if not sendMessage then
+			return responsePayload
+		end
 
-	-- Helper to build response messages with optional requestId
-	local function buildResponse(responsePayload: { [string]: any }): { [string]: any }
-		local msg: { [string]: any } = {
+		sendMessage({
 			type = "scriptComplete",
 			sessionId = sessionId,
+			requestId = effectiveRequestId,
 			payload = responsePayload,
-		}
-		if effectiveRequestId then
-			msg.requestId = effectiveRequestId
-		end
-		return msg
-	end
-
-	-- Helper to send or return the result
-	local function sendResult(responsePayload: { [string]: any }): { [string]: any }?
-		if sendMessage then
-			sendMessage(buildResponse(responsePayload))
-			return nil
-		end
-		return responsePayload
+		})
+		return nil
 	end
 
-	if not code or type(code) ~= "string" then
-		return sendResult({ success = false, error = "Missing code in execute request", code = "SCRIPT_LOAD_ERROR" })
+	local code = payload and (payload.script or payload.code)
+	if type(code) ~= "string" then
+		return reply({ success = false, error = "Missing code in execute request", code = "SCRIPT_LOAD_ERROR" })
 	end
 
-	local compileOk, fnOrErr = pcall(loadstring, code)
+	local compileOk, fn, compileErr = pcall(compile, code, CHUNK_NAME)
 	if not compileOk then
-		return sendResult({ success = false, error = tostring(fnOrErr), code = "SCRIPT_LOAD_ERROR" })
+		return reply({ success = false, error = tostring(fn), code = "SCRIPT_LOAD_ERROR" })
 	end
-	if not fnOrErr then
-		return sendResult({ success = false, error = "Failed to compile script", code = "SCRIPT_LOAD_ERROR" })
+	if not fn then
+		return reply({
+			success = false,
+			error = if compileErr then tostring(compileErr) else "Failed to compile script",
+			code = "SCRIPT_LOAD_ERROR",
+		})
 	end
 
-	local fn = fnOrErr
-
-	-- Capture output by temporarily intercepting print/warn in the global
-	-- environment. Local variable shadowing doesn't work because loadstring'd
-	-- code accesses the global environment, not this module's locals.
-	local captured: { { level: string, body: string, timestamp: number } } = {}
-	local env = getfenv(fn)
-	local originalPrint = env.print
-	local originalWarn = env.warn
-
-	env.print = function(...)
-		local parts = {}
-		for i = 1, select("#", ...) do
-			parts[i] = tostring(select(i, ...))
-		end
-		local body = table.concat(parts, "\t")
-		table.insert(captured, { level = "Print", body = body, timestamp = os.clock() })
-		originalPrint(...)
+	local args = payload.args or {}
+	if type(args) ~= "table" then
+		return reply({ success = false, error = "Script args must be a list", code = "SCRIPT_LOAD_ERROR" })
 	end
 
-	env.warn = function(...)
-		local parts = {}
-		for i = 1, select("#", ...) do
-			parts[i] = tostring(select(i, ...))
+	local timeoutMs = payload.timeoutMs
+	local timeoutSeconds = if type(timeoutMs) == "number" then timeoutMs / 1000 else nil
+	-- The caller's timeout started when it sent the request, not when its turn came
+	if timeoutSeconds and receivedAt then
+		timeoutSeconds -= os.clock() - receivedAt
+		if timeoutSeconds <= 0 then
+			return reply({
+				success = false,
+				error = `Script timed out after {timeoutMs}ms waiting behind other scripts, and never ran`,
+				code = "SCRIPT_TIMEOUT",
+			})
 		end
-		local body = table.concat(parts, "\t")
-		table.insert(captured, { level = "Warning", body = body, timestamp = os.clock() })
-		originalWarn(...)
 	end
 
-	local returned = table.pack(pcall(fn))
-	local success = returned[1]
+	local captured: { LogEntry } = {}
+	local logConnection = ExecuteAction._captureOutput(captured)
+	local status, result = ExecuteAction._runAsync(fn, timeoutSeconds, args)
 
-	-- Restore originals
-	env.print = originalPrint
-	env.warn = originalWarn
+	-- MessageOut can fire a frame late, after the script's last print
+	task.wait()
+	logConnection:Disconnect()
 
-	if not success then
-		return sendResult({
+	if status == "timeout" then
+		return reply({
 			success = false,
-			error = tostring(returned[2]),
+			error = `Script timed out after {timeoutMs}ms and was cancelled`,
+			code = "SCRIPT_TIMEOUT",
+			output = captured,
+		})
+	elseif status == "error" then
+		return reply({
+			success = false,
+			error = result.message,
+			traceback = result.traceback,
 			code = "SCRIPT_RUNTIME_ERROR",
 			output = captured,
 		})
 	end
 
-	return sendResult({
+	return reply({
 		success = true,
 		output = captured,
-		returnValues = serializeReturnValues(returned),
+		returnValues = ExecuteAction._serializeReturnValues(result),
 	})
 end
 
--- ---------------------------------------------------------------------------
--- Queue processing
--- ---------------------------------------------------------------------------
-
-local function _processQueue()
-	if _processing then
-		return
+--[[
+	Drops the frames of this module's own wrapper from a traceback, leaving only
+	the script's.
+]]
+function ExecuteAction._dropOwnFrames(traceback: string): string
+	if not OWN_FRAME_PREFIX then
+		return traceback
 	end
-	_processing = true
 
-	while #_queue > 0 do
-		local item = table.remove(_queue, 1)
-		ExecuteAction._handleExecute(item.payload, item.requestId, item.sessionId, item.sendMessage)
+	local lines = {}
+	for _, line in string.split(traceback, "\n") do
+		if string.sub(line, 1, #OWN_FRAME_PREFIX + 1) ~= `{OWN_FRAME_PREFIX}:` then
+			table.insert(lines, line)
+		end
 	end
+	return table.concat(lines, "\n")
+end
 
-	_processing = false
+--[[
+	Collects every message logged until the connection is disconnected, other
+	than the plugin's own.
+]]
+function ExecuteAction._captureOutput(captured: { LogEntry }): RBXScriptConnection
+	return game:GetService("LogService").MessageOut:Connect(function(message: string, messageType: Enum.MessageType)
+		if string.sub(message, 1, #PLUGIN_PREFIX) == PLUGIN_PREFIX then
+			return
+		end
+
+		table.insert(captured, {
+			level = LOG_LEVELS[messageType] or "Print",
+			body = message,
+			timestamp = os.clock(),
+		})
+	end)
 end
 
--- ---------------------------------------------------------------------------
--- Public API
--- ---------------------------------------------------------------------------
+--[[
+	Runs `fn` with `args` in its own thread. Returns "ok" and everything it returned
+	(pcall-style, status first), "error" and its message and traceback, or
+	"timeout" once `timeoutSec` passes, after cancelling the thread.
+]]
+function ExecuteAction._runAsync(fn: (...any) -> ...any, timeoutSec: number?, args: { any }?): (string, any)
+	local callerThread = coroutine.running()
+	local isWaiting = false
+	local outcome: { status: string, result: any }? = nil
+
+	local function finish(status: string, result: any)
+		if outcome then
+			return
+		end
+		outcome = { status = status, result = result }
+		if isWaiting then
+			task.defer(callerThread)
+		end
+	end
 
--- Register this handler with the ActionRouter.
-function ExecuteAction.register(router: any, sendMessage: ((msg: { [string]: any }) -> ())?)
-	router:setResponseType("execute", "scriptComplete")
-	router:register("execute", function(payload: { [string]: any }, requestId: string, sessionId: string)
-		if sendMessage then
-			-- Queue the request for sequential processing
-			table.insert(_queue, {
-				payload = payload,
-				requestId = requestId,
-				sessionId = sessionId,
-				sendMessage = sendMessage,
-			})
-			_processQueue()
-			-- Return nil so the ActionRouter does not generate a response
-			return nil
+	local scriptThread = task.spawn(function()
+		local packed = table.pack(xpcall(fn, function(err)
+			return {
+				message = tostring(err),
+				traceback = ExecuteAction._dropOwnFrames(debug.traceback(nil, 2)),
+			}
+		end, table.unpack(args or {})))
+		if packed[1] then
+			finish("ok", packed)
 		else
-			-- Direct mode (no sendMessage): return payload for ActionRouter wrapping
-			return ExecuteAction._handleExecute(payload, requestId, sessionId, nil)
+			finish("error", packed[2])
 		end
 	end)
+
+	if not outcome then
+		local timeoutThread = if timeoutSec
+			then task.delay(timeoutSec, function()
+				if not outcome then
+					pcall(task.cancel, scriptThread)
+					finish("timeout", nil)
+				end
+			end)
+			else nil
+
+		isWaiting = true
+		coroutine.yield()
+		if timeoutThread then
+			pcall(task.cancel, timeoutThread)
+		end
+	end
+
+	local final = outcome :: { status: string, result: any }
+	return final.status, final.result
 end
 
--- Handle an execute request directly (for testing without ActionRouter).
-function ExecuteAction.handleExecute(
-	payload: { [string]: any },
-	requestId: string?,
-	sessionId: string
-): { [string]: any }?
-	return ExecuteAction._handleExecute(payload, requestId, sessionId, nil)
+--[[
+	Marshals what a script returned, in order, skipping the pcall status. A nil
+	becomes an explicit marker: a JSON array cannot hold a hole, and dropping it
+	would shift every later value into the wrong position.
+]]
+function ExecuteAction._serializeReturnValues(packed: { n: number, [number]: any }): { any }
+	local returnValues = {}
+	for index = 2, packed.n do
+		local serialized = ExecuteAction._serializeValue(packed[index], 1, {})
+		if serialized == nil then
+			serialized = { type = "Nil" }
+		end
+		returnValues[index - 1] = serialized
+	end
+	return returnValues
 end
 
--- Reset the internal queue state (for testing).
-function ExecuteAction._resetQueue()
-	_queue = {}
-	_processing = false
+--[[
+	Produces the SerializedReturnValue shapes declared in web-socket-protocol.ts.
+	Duplicates the marshaller in query-data-model.lua, since actions cannot
+	require each other.
+]]
+function ExecuteAction._serializeValue(value: any, depth: number, seen: { [any]: boolean }): any
+	local valueType = typeof(value)
+
+	if valueType == "nil" or valueType == "string" or valueType == "boolean" then
+		return value
+	elseif valueType == "number" then
+		-- inf and nan have no JSON spelling, and encoding one produces a
+		-- document the server cannot decode, losing the entire message
+		if value ~= value or value == math.huge or value == -math.huge then
+			return ExecuteAction._unsupported("number", tostring(value))
+		end
+		return value
+	elseif valueType == "table" then
+		if depth > MAX_DEPTH then
+			return ExecuteAction._unsupported("table", "<max depth exceeded>")
+		end
+		if seen[value] then
+			return ExecuteAction._unsupported("table", "<cycle>")
+		end
+
+		seen[value] = true
+		local serialized = ExecuteAction._serializeTable(value, depth, seen)
+		seen[value] = nil
+		return serialized
+	elseif valueType == "Vector3" then
+		return { type = "Vector3", value = { value.X, value.Y, value.Z } }
+	elseif valueType == "Vector2" then
+		return { type = "Vector2", value = { value.X, value.Y } }
+	elseif valueType == "CFrame" then
+		return { type = "CFrame", value = { value:GetComponents() } }
+	elseif valueType == "Color3" then
+		return { type = "Color3", value = { value.R, value.G, value.B } }
+	elseif valueType == "UDim2" then
+		return { type = "UDim2", value = { value.X.Scale, value.X.Offset, value.Y.Scale, value.Y.Offset } }
+	elseif valueType == "UDim" then
+		return { type = "UDim", value = { value.Scale, value.Offset } }
+	elseif valueType == "BrickColor" then
+		return { type = "BrickColor", name = value.Name, value = value.Number }
+	elseif valueType == "EnumItem" then
+		return { type = "EnumItem", enum = tostring(value.EnumType), name = value.Name, value = value.Value }
+	elseif valueType == "Instance" then
+		return { type = "Instance", className = value.ClassName, path = value:GetFullName() }
+	else
+		return ExecuteAction._unsupported(valueType, tostring(value))
+	end
 end
 
-function ExecuteAction.teardown()
-	for _, item in _queue do
-		if item.sendMessage then
-			item.sendMessage({
-				type = "scriptComplete",
-				sessionId = item.sessionId,
-				requestId = item.requestId,
-				payload = {
-					success = false,
-					error = "Action re-registered, request cancelled",
-					code = "ACTION_REPLACED",
-				},
-			})
+--[[
+	An array-like table becomes a JSON array, anything else a JSON object with
+	string keys. Mixed tables take the object branch on purpose: JSONEncode
+	rejects them, and a rejected encode drops the whole scriptComplete message.
+]]
+function ExecuteAction._serializeTable(value: { [any]: any }, depth: number, seen: { [any]: boolean }): any
+	local count = 0
+	for _ in value do
+		count += 1
+	end
+
+	if count == #value then
+		local array = {}
+		for index = 1, count do
+			array[index] = ExecuteAction._serializeValue(value[index], depth + 1, seen)
 		end
+		return array
 	end
-	_queue = {}
-	_processing = false
+
+	local map: { [string]: any } = {}
+	for key, item in value do
+		map[if type(key) == "string" then key else tostring(key)] = ExecuteAction._serializeValue(item, depth + 1, seen)
+	end
+	return map
+end
+
+function ExecuteAction._unsupported(typeName: string, text: string): { [string]: any }
+	return { type = "Unsupported", typeName = typeName, toString = text }
 end
 
 return ExecuteAction