diff --git a/plugin/src/ApiContext.lua b/plugin/src/ApiContext.lua index 0b064463..71dd58c4 100644 --- a/plugin/src/ApiContext.lua +++ b/plugin/src/ApiContext.lua @@ -15,6 +15,33 @@ local validateApiSerialize = Types.ifEnabled(Types.ApiSerializeResponse) local validateApiRefPatch = Types.ifEnabled(Types.ApiRefPatchResponse) local EXEC_COMPLETION_BODY_LIMIT_BYTES = 256 * 1024 +local AUTOMATION_COMPLETION_BODY_LIMIT_BYTES = 4 * 1024 * 1024 +local AUTOMATION_HANDLER_VERSION = 2 + +local function generatePluginSessionId() + return HttpService:GenerateGUID(false) +end + +local function beginPluginSession(apiContext, serverSessionId, generateId) + apiContext.__sessionId = serverSessionId + apiContext.__pluginSessionId = (generateId or generatePluginSessionId)() + return apiContext.__pluginSessionId +end + +local function withExecSession(payload, pluginSessionId, studioMode) + local body = table.clone(payload) + body.pluginSessionId = pluginSessionId + body.studioMode = studioMode + return body +end + +local function buildExecClaimUrl(baseUrl, pluginSessionId, studioMode) + return ("%s/api/exec/jobs/next?pluginSessionId=%s&studioMode=%s"):format(baseUrl, pluginSessionId, studioMode) +end + +local function buildAutomationClaimUrl(baseUrl, pluginSessionId, studioMode) + return ("%s/api/automation/jobs/next?pluginSessionId=%s&studioMode=%s"):format(baseUrl, pluginSessionId, studioMode) +end local function rejectFailedRequests(response) if response.code >= 400 then @@ -26,18 +53,19 @@ local function rejectFailedRequests(response) return response end -local function decodeExecResponse(response, validator, description) +local function decodeExecResponse(response, validator, description, protocolName) + protocolName = protocolName or "Prism exec" local decodeOk, body = pcall(response.msgpack, response) if not decodeOk then return Promise.reject( - string.format("Prism exec protocol error: could not decode %s: %s", description, tostring(body)) + string.format("%s protocol error: could not decode %s: %s", protocolName, description, tostring(body)) ) end local valid, validationError = validator(body) if not valid then return Promise.reject( - string.format("Prism exec protocol error: malformed %s: %s", description, tostring(validationError)) + string.format("%s protocol error: malformed %s: %s", protocolName, description, tostring(validationError)) ) end @@ -127,6 +155,8 @@ function ApiContext.new(baseUrl) local self = { __baseUrl = baseUrl, __sessionId = nil, + __pluginSessionId = nil, + __connectionGeneration = 0, __messageCursor = -1, __wsClient = nil, __connected = true, @@ -143,6 +173,7 @@ function ApiContext:__fmtDebug(output) output:writeLine("Connected: {}", self.__connected) output:writeLine("Base URL: {}", self.__baseUrl) output:writeLine("Session ID: {}", self.__sessionId) + output:writeLine("Plugin Session ID: {}", self.__pluginSessionId) output:writeLine("Message Cursor: {}", self.__messageCursor) output:unindent() @@ -151,6 +182,7 @@ end function ApiContext:disconnect() self.__connected = false + self.__connectionGeneration += 1 for request in self.__activeRequests do Log.trace("Cancelling request {}", request) request:cancel() @@ -162,6 +194,7 @@ function ApiContext:disconnect() self.__wsClient:Close() end self.__wsClient = nil + self.__pluginSessionId = nil end function ApiContext:setMessageCursor(index) @@ -169,6 +202,9 @@ function ApiContext:setMessageCursor(index) end function ApiContext:connect() + self.__connected = true + self.__connectionGeneration += 1 + local generation = self.__connectionGeneration local url = ("%s/api/rojo"):format(self.__baseUrl) return Http.get(url) @@ -182,12 +218,19 @@ function ApiContext:connect() end) :andThen(rejectWrongPlaceId) :andThen(function(body) - self.__sessionId = body.sessionId + if not self.__connected or self.__connectionGeneration ~= generation then + return Promise.reject("Connection was stopped before server verification completed") + end + beginPluginSession(self, body.sessionId) return body end) end +function ApiContext:getPluginSessionId() + return self.__pluginSessionId +end + function ApiContext:read(ids) local url = ("%s/api/read/%s"):format(self.__baseUrl, table.concat(ids, ",")) @@ -356,8 +399,16 @@ function ApiContext:refPatch(ids: { string }) end) end -function ApiContext:claimNextExecJob() - local url = ("%s/api/exec/jobs/next"):format(self.__baseUrl) +function ApiContext:claimNextExecJob(studioMode) + local validMode, modeError = Types.StudioMode(studioMode) + if not validMode then + return Promise.reject(string.format("Cannot claim exec job: %s", tostring(modeError))) + end + if self.__pluginSessionId == nil then + return Promise.reject("Cannot claim exec job without an active plugin session") + end + + local url = buildExecClaimUrl(self.__baseUrl, self.__pluginSessionId, studioMode) return Http.get(url):andThen(function(response) if response.code == 204 then @@ -372,11 +423,20 @@ function ApiContext:claimNextExecJob() end) end -function ApiContext:completeExecJob(jobId, payload) +function ApiContext:completeExecJob(jobId, payload, studioMode) local validJobId, jobIdError = Types.Uuid(jobId) if not validJobId then return Promise.reject(string.format("Cannot complete exec job: %s", tostring(jobIdError))) end + local validMode, modeError = Types.StudioMode(studioMode) + if not validMode then + return Promise.reject(string.format("Cannot complete exec job: %s", tostring(modeError))) + end + if self.__pluginSessionId == nil then + return Promise.reject("Cannot complete exec job without an active plugin session") + end + + payload = withExecSession(payload, self.__pluginSessionId, studioMode) local encodeOk, body = pcall(Http.msgpackEncode, payload) if not encodeOk then @@ -421,4 +481,121 @@ function ApiContext:completeExecJob(jobId, payload) end) end +function ApiContext:claimNextAutomationJob(studioMode) + local validMode, modeError = Types.StudioMode(studioMode) + if not validMode then + return Promise.reject(string.format("Cannot claim automation job: %s", tostring(modeError))) + end + if self.__pluginSessionId == nil then + return Promise.reject("Cannot claim automation job without an active plugin session") + end + + local url = buildAutomationClaimUrl(self.__baseUrl, self.__pluginSessionId, studioMode) + return Http.get(url):andThen(function(response) + if response.code == 204 then + return nil + end + if response.code ~= 200 then + return Promise.reject( + string.format("Unexpected automation claim response status %s", tostring(response.code)) + ) + end + return decodeExecResponse( + response, + Types.ApiAutomationClaimResponse, + "claimed-job response", + "Prism automation" + ) + end) +end + +function ApiContext:completeAutomationJob(jobId, payload, studioMode) + local validJobId, jobIdError = Types.Uuid(jobId) + if not validJobId then + return Promise.reject(string.format("Cannot complete automation job: %s", tostring(jobIdError))) + end + local validMode, modeError = Types.StudioMode(studioMode) + if not validMode then + return Promise.reject(string.format("Cannot complete automation job: %s", tostring(modeError))) + end + if self.__pluginSessionId == nil then + return Promise.reject("Cannot complete automation job without an active plugin session") + end + + payload = withExecSession(payload, self.__pluginSessionId, studioMode) + local encodeOk, body = pcall(Http.msgpackEncode, payload) + if not encodeOk then + return Promise.reject("Could not encode automation completion payload: " .. tostring(body)) + end + if #body > AUTOMATION_COMPLETION_BODY_LIMIT_BYTES then + return Promise.reject( + string.format( + "Automation completion payload is %d bytes, exceeding the %d-byte server limit", + #body, + AUTOMATION_COMPLETION_BODY_LIMIT_BYTES + ) + ) + end + + local url = ("%s/api/automation/jobs/%s/complete"):format(self.__baseUrl, jobId) + return Http.post(url, body) + :andThen(function(response) + if response.code ~= 200 then + return Promise.reject( + string.format("Unexpected automation completion response status %s", tostring(response.code)) + ) + end + local responseBody = + decodeExecResponse(response, Types.ApiAutomationJobResponse, "completion response", "Prism automation") + return { status = "accepted", job = responseBody } + end) + :catch(function(errorValue) + if isCompletionConflict(errorValue) then + return { status = "conflict" } + end + return Promise.reject(errorValue) + end) +end + +function ApiContext:updateAutomationStatus(studioMode) + local validMode, modeError = Types.StudioMode(studioMode) + if not validMode then + return Promise.reject(string.format("Cannot update automation status: %s", tostring(modeError))) + end + if self.__pluginSessionId == nil or self.__sessionId == nil then + return Promise.reject("Cannot update automation status without an active plugin session") + end + + local body = Http.msgpackEncode({ + pluginSessionId = self.__pluginSessionId, + serverSessionId = self.__sessionId, + studioMode = studioMode, + execHandlerAvailable = true, + automationHandlerVersion = AUTOMATION_HANDLER_VERSION, + pluginVersion = Version.display(Config.version), + }) + local url = ("%s/api/automation/status"):format(self.__baseUrl) + + return Http.post(url, body):andThen(function(response) + if response.code ~= 200 then + return Promise.reject(string.format("Unexpected automation status response %s", tostring(response.code))) + end + + return decodeExecResponse( + response, + Types.ApiAutomationHeartbeatResponse, + "automation status response", + "Prism automation" + ) + end) +end + +ApiContext._test = { + beginPluginSession = beginPluginSession, + buildAutomationClaimUrl = buildAutomationClaimUrl, + buildExecClaimUrl = buildExecClaimUrl, + generatePluginSessionId = generatePluginSessionId, + withExecSession = withExecSession, +} + return ApiContext diff --git a/plugin/src/ApiContext.spec.lua b/plugin/src/ApiContext.spec.lua new file mode 100644 index 00000000..334603e2 --- /dev/null +++ b/plugin/src/ApiContext.spec.lua @@ -0,0 +1,58 @@ +return function() + local ApiContext = require(script.Parent.ApiContext) + local Types = require(script.Parent.Types) + local helpers = ApiContext._test + + describe("automation plugin session identity", function() + it("retains one generated ID for an active session", function() + local context = ApiContext.new("http://127.0.0.1:34872") + local generateCount = 0 + local id = helpers.beginPluginSession(context, "server-one", function() + generateCount += 1 + return "00000000-0000-4000-8000-000000000001" + end) + + expect(id).to.equal("00000000-0000-4000-8000-000000000001") + expect(context:getPluginSessionId()).to.equal(id) + expect(generateCount).to.equal(1) + end) + + it("creates a new ID for a new connection lifecycle", function() + local context = ApiContext.new("http://127.0.0.1:34872") + local nextId = 0 + local function generate() + nextId += 1 + return ("00000000-0000-4000-8000-%012d"):format(nextId) + end + local first = helpers.beginPluginSession(context, "server-one", generate) + context:disconnect() + local second = helpers.beginPluginSession(context, "server-one", generate) + + expect(first).never.to.equal(second) + expect(context:getPluginSessionId()).to.equal(second) + end) + + it("adds identity and mode to claims and completions", function() + local id = "00000000-0000-4000-8000-000000000001" + local url = helpers.buildExecClaimUrl("http://127.0.0.1:34872", id, "edit") + local payload = helpers.withExecSession({ outcome = "success" }, id, "run") + + expect(url:find("pluginSessionId=" .. id, 1, true)).to.be.ok() + expect(url:find("studioMode=edit", 1, true)).to.be.ok() + expect(payload.pluginSessionId).to.equal(id) + expect(payload.studioMode).to.equal("run") + expect(payload.outcome).to.equal("success") + + local automationUrl = helpers.buildAutomationClaimUrl("http://127.0.0.1:34872", id, "edit") + expect(automationUrl:find("/api/automation/jobs/next", 1, true)).to.be.ok() + expect(automationUrl:find("pluginSessionId=" .. id, 1, true)).to.be.ok() + end) + + it("validates all supported Studio modes", function() + for _, mode in { "edit", "play", "run", "unknown" } do + expect(Types.StudioMode(mode)).to.equal(true) + end + expect(Types.StudioMode("server")).to.equal(false) + end) + end) +end diff --git a/plugin/src/Automation.lua b/plugin/src/Automation.lua new file mode 100644 index 00000000..043fd8d7 --- /dev/null +++ b/plugin/src/Automation.lua @@ -0,0 +1,201 @@ +local RunService = game:GetService("RunService") + +local Packages = script.Parent.Parent.Packages +local Log = require(Packages.Log) +local Promise = require(Packages.Promise) + +local Inspect = require(script.Parent.AutomationHandlers.Inspect) +local InstanceReferences = require(script.Parent.InstanceReferences) + +local POLL_INTERVAL_SECONDS = 0.25 +local COMPLETION_RETRY_SECONDS = 0.25 +local COMPLETION_ATTEMPTS = 3 + +local function currentStudioMode() + if RunService:IsEdit() then + return "edit" + elseif RunService:IsRunMode() then + return "run" + elseif RunService:IsRunning() then + return "play" + else + return "unknown" + end +end + +local function dispatch(job, references) + if job.request.kind ~= "inspect" then + return { outcome = "failure", error = "Unsupported automation job kind: " .. tostring(job.request.kind) } + end + local ok, result, handlerError = pcall(Inspect.run, job.request, { references = references }) + if not ok then + return { outcome = "failure", error = "Inspect handler failed: " .. tostring(result) } + end + if result == nil then + return { outcome = "failure", error = tostring(handlerError) } + end + return { + outcome = "success", + result = { + kind = "inspect", + root = result.root, + visitedInstances = result.visitedInstances, + truncated = result.truncated, + truncationReason = result.truncationReason, + }, + } +end + +local Automation = {} +Automation.__index = Automation + +function Automation.new(options) + local dependencies = options.dependencies or {} + return setmetatable({ + __apiContext = options.apiContext, + __delay = dependencies.delay or Promise.delay, + __studioMode = dependencies.studioMode or currentStudioMode, + __dispatch = dependencies.dispatch or dispatch, + __makeReferences = dependencies.makeReferences or InstanceReferences.new, + __onError = options.onError or function(errorValue) + Log.error("Prism automation poller failed: {}", errorValue) + end, + __running = false, + __generation = 0, + __busy = false, + __scheduledPromise = nil, + __references = nil, + }, Automation) +end + +function Automation:__isCurrent(generation) + return self.__running and self.__generation == generation +end + +function Automation:__fail(generation, errorValue) + if self:__isCurrent(generation) then + self.__onError(errorValue) + end +end + +function Automation:__schedule(generation, delaySeconds) + if not self:__isCurrent(generation) then + return + end + self.__scheduledPromise = self.__delay(delaySeconds):andThen(function() + if self:__isCurrent(generation) then + self:__poll(generation) + end + end) +end + +function Automation:__release(generation) + if self:__isCurrent(generation) then + self.__busy = false + self:__schedule(generation, POLL_INTERVAL_SECONDS) + end +end + +function Automation:__complete(generation, jobId, payload, attempt) + if not self:__isCurrent(generation) then + return + end + local ok, promise = + pcall(self.__apiContext.completeAutomationJob, self.__apiContext, jobId, payload, self.__studioMode()) + if not ok then + self:__fail(generation, promise) + return + end + promise + :andThen(function(response) + if not self:__isCurrent(generation) then + return + end + if response.status == "conflict" then + Log.warn("Prism automation completion for job {} returned HTTP 409", jobId) + end + self:__release(generation) + end) + :catch(function(errorValue) + if not self:__isCurrent(generation) then + return + end + if attempt >= COMPLETION_ATTEMPTS then + self:__fail(generation, errorValue) + return + end + self.__delay(COMPLETION_RETRY_SECONDS):andThen(function() + if self:__isCurrent(generation) then + self:__complete(generation, jobId, payload, attempt + 1) + end + end) + end) +end + +function Automation:__poll(generation) + if not self:__isCurrent(generation) or self.__busy then + return + end + local mode = self.__studioMode() + if mode ~= "edit" then + self:__schedule(generation, POLL_INTERVAL_SECONDS) + return + end + self.__busy = true + local ok, promise = pcall(self.__apiContext.claimNextAutomationJob, self.__apiContext, mode) + if not ok then + self:__fail(generation, promise) + return + end + promise + :andThen(function(job) + if not self:__isCurrent(generation) then + return + end + if job == nil then + self:__release(generation) + return + end + local payload = self.__dispatch(job, self.__references) + self:__complete(generation, job.jobId, payload, 1) + end) + :catch(function(errorValue) + self:__fail(generation, errorValue) + end) +end + +function Automation:start() + if self.__running then + return + end + local sessionId = self.__apiContext:getPluginSessionId() + if sessionId == nil then + self.__onError("Cannot start automation poller without a plugin session ID") + return + end + self.__references = self.__makeReferences(sessionId) + self.__running = true + self.__generation += 1 + self:__schedule(self.__generation, 0) +end + +function Automation:stop() + if not self.__running then + return + end + self.__running = false + self.__generation += 1 + self.__busy = false + if self.__scheduledPromise ~= nil then + self.__scheduledPromise:cancel() + self.__scheduledPromise = nil + end + if self.__references ~= nil then + self.__references:clear() + self.__references = nil + end +end + +Automation._test = { dispatch = dispatch, currentStudioMode = currentStudioMode } + +return Automation diff --git a/plugin/src/Automation.spec.lua b/plugin/src/Automation.spec.lua new file mode 100644 index 00000000..0b5dae26 --- /dev/null +++ b/plugin/src/Automation.spec.lua @@ -0,0 +1,130 @@ +return function() + local Packages = script.Parent.Parent.Packages + local Promise = require(Packages.Promise) + local Automation = require(script.Parent.Automation) + + local function scheduler() + local pending = {} + return { + delay = function() + return Promise.new(function(resolve) + table.insert(pending, resolve) + end) + end, + step = function() + local resolve = table.remove(pending, 1) + assert(resolve ~= nil, "no pending timer") + resolve() + end, + count = function() + return #pending + end, + } + end + + local function fakeReferences() + return { clear = function() end } + end + + describe("typed automation poller", function() + it("starts once, idles on empty claims, and stops", function() + local timers = scheduler() + local claims = 0 + local api = { + getPluginSessionId = function() + return "plugin-session" + end, + claimNextAutomationJob = function() + claims += 1 + return Promise.resolve(nil) + end, + } + local poller = Automation.new({ + apiContext = api, + dependencies = { + delay = timers.delay, + studioMode = function() + return "edit" + end, + makeReferences = fakeReferences, + }, + }) + poller:start() + poller:start() + expect(timers.count()).to.equal(1) + timers.step() + expect(claims).to.equal(1) + poller:stop() + end) + + it("does not claim outside edit mode", function() + local timers = scheduler() + local claims = 0 + local poller = Automation.new({ + apiContext = { + getPluginSessionId = function() + return "session" + end, + claimNextAutomationJob = function() + claims += 1 + return Promise.resolve(nil) + end, + }, + dependencies = { + delay = timers.delay, + studioMode = function() + return "play" + end, + makeReferences = fakeReferences, + }, + }) + poller:start() + timers.step() + expect(claims).to.equal(0) + poller:stop() + end) + + it("dispatches and completes one claimed inspect job", function() + local timers = scheduler() + local dispatches = 0 + local completions = 0 + local api = { + getPluginSessionId = function() + return "session" + end, + claimNextAutomationJob = function() + return Promise.resolve({ jobId = "job", request = { kind = "inspect" } }) + end, + completeAutomationJob = function(_, _, payload) + completions += 1 + expect(payload.outcome).to.equal("success") + return Promise.resolve({ status = "accepted" }) + end, + } + local poller = Automation.new({ + apiContext = api, + dependencies = { + delay = timers.delay, + studioMode = function() + return "edit" + end, + makeReferences = fakeReferences, + dispatch = function() + dispatches += 1 + return { outcome = "success", result = { kind = "inspect" } } + end, + }, + }) + poller:start() + timers.step() + expect(dispatches).to.equal(1) + expect(completions).to.equal(1) + poller:stop() + end) + + it("turns unknown kinds into typed failures", function() + local payload = Automation._test.dispatch({ request = { kind = "unknown" } }, {}) + expect(payload.outcome).to.equal("failure") + end) + end) +end diff --git a/plugin/src/AutomationHandlers/Inspect.lua b/plugin/src/AutomationHandlers/Inspect.lua new file mode 100644 index 00000000..9f728a24 --- /dev/null +++ b/plugin/src/AutomationHandlers/Inspect.lua @@ -0,0 +1,276 @@ +local CollectionService = game:GetService("CollectionService") + +local AutomationValue = require(script.Parent.Parent.AutomationValue) + +local Inspect = {} + +local ROOT_SERVICES = { + Lighting = true, + ReplicatedStorage = true, + ServerStorage = true, + ServerScriptService = true, + StarterGui = true, + StarterPlayer = true, + StarterPack = true, + SoundService = true, + Teams = true, + TextChatService = true, +} + +local BASE_PART_PROPERTIES = { + "Anchored", + "CanCollide", + "CanQuery", + "CanTouch", + "CastShadow", + "CFrame", + "Position", + "Orientation", + "Size", + "Color", + "Material", + "Transparency", + "Reflectance", + "Massless", + "CollisionGroup", +} +local GUI_OBJECT_PROPERTIES = { + "Visible", + "Position", + "Size", + "AnchorPoint", + "Rotation", + "ZIndex", + "BackgroundColor3", + "BackgroundTransparency", +} +local CAMERA_PROPERTIES = { "CFrame", "Focus", "FieldOfView", "CameraType", "CameraSubject", "ViewportSize" } + +local function quoteSegment(name) + if name:match("^[A-Za-z_][A-Za-z0-9_]*$") then + return name + end + return '"' .. name:gsub("\\", "\\\\"):gsub('"', '\\"'):gsub("\n", "\\n"):gsub("\r", "\\r"):gsub("\t", "\\t") .. '"' +end + +local function childPath(parentPath, child) + return parentPath .. "." .. quoteSegment(child.Name) +end + +local function instancePath(instance) + if instance == game then + return "game" + elseif instance == workspace then + return "Workspace" + end + local segments = {} + local current = instance + while current ~= nil and current ~= game do + table.insert(segments, 1, quoteSegment(current.Name)) + current = current.Parent + end + return table.concat(segments, ".") +end + +local function isInDataModel(instance) + local ok, result = pcall(function() + return instance == game or instance:IsDescendantOf(game) + end) + return ok and result +end + +local function resolveTarget(segments) + local rootName = segments[1] + local current + if rootName == "game" then + current = game + elseif rootName == "Workspace" then + current = workspace + elseif ROOT_SERVICES[rootName] then + local ok, service = pcall(game.GetService, game, rootName) + if not ok then + return nil, string.format("Could not resolve service '%s': %s", rootName, tostring(service)) + end + current = service + else + return nil, string.format("Unsupported inspect root '%s'", tostring(rootName)) + end + + for index = 2, #segments do + local found = nil + for _, child in current:GetChildren() do + if child.Name == segments[index] then + found = child + break + end + end + if found == nil then + return nil, string.format("Inspect target segment '%s' was not found", segments[index]) + end + current = found + end + return current +end + +local function appendProperties(instance, names, output, references, path) + for _, name in names do + local readOk, value = pcall(function() + return instance[name] + end) + if not readOk then + output[name] = AutomationValue.diagnostic(value) + else + local valuePath = if typeof(value) == "Instance" then instancePath(value) else path + local encoded, encodeError = AutomationValue.encode(value, references, valuePath) + output[name] = encoded or AutomationValue.diagnostic(encodeError) + end + end +end + +local function readProperties(instance, references, path) + local properties = {} + appendProperties(instance, { "Archivable", "Parent" }, properties, references, path) + if instance:IsA("BasePart") then + appendProperties(instance, BASE_PART_PROPERTIES, properties, references, path) + end + if instance:IsA("Model") then + appendProperties(instance, { "WorldPivot", "PrimaryPart" }, properties, references, path) + local ok, cframe, size = pcall(instance.GetBoundingBox, instance) + if ok then + local encoded, encodeError = AutomationValue.encode({ cframe = cframe, size = size }, references, path) + properties.BoundingBox = encoded or AutomationValue.diagnostic(encodeError) + else + properties.BoundingBox = AutomationValue.diagnostic(cframe) + end + end + if instance:IsA("ValueBase") then + appendProperties(instance, { "Value" }, properties, references, path) + end + if instance:IsA("GuiObject") then + appendProperties(instance, GUI_OBJECT_PROPERTIES, properties, references, path) + end + if instance:IsA("LayerCollector") then + appendProperties(instance, { "Enabled" }, properties, references, path) + end + if instance:IsA("ScreenGui") then + appendProperties(instance, { "DisplayOrder", "IgnoreGuiInset" }, properties, references, path) + end + if instance:IsA("Camera") then + appendProperties(instance, CAMERA_PROPERTIES, properties, references, path) + end + return properties +end + +local function readAttributes(instance, references, path) + local output = {} + local ok, attributes = pcall(instance.GetAttributes, instance) + if not ok then + output[""] = AutomationValue.diagnostic(attributes) + return output + end + local keys = {} + for key in attributes do + table.insert(keys, key) + end + table.sort(keys) + for _, key in keys do + local encoded, encodeError = AutomationValue.encode(attributes[key], references, path) + output[key] = encoded or AutomationValue.diagnostic(encodeError) + end + return output +end + +local function readTags(instance) + local ok, tags = pcall(CollectionService.GetTags, CollectionService, instance) + if not ok then + return {} + end + table.sort(tags) + return tags +end + +function Inspect.run(request, context) + if request.target == nil or request.target.kind ~= "path" or type(request.target.segments) ~= "table" then + return nil, "Malformed inspect target" + end + local target, targetError = resolveTarget(request.target.segments) + if target == nil then + return nil, targetError + end + if not isInDataModel(target) then + return nil, "Inspect target was destroyed before traversal" + end + + local pathSegments = table.create(#request.target.segments) + for index, segment in request.target.segments do + pathSegments[index] = if index == 1 then segment else quoteSegment(segment) + end + local rootPath = table.concat(pathSegments, ".") + local visited = 0 + local globallyTruncated = false + local truncationReason = nil + + local function visit(instance, path, depth) + if visited >= request.maxInstances then + globallyTruncated = true + truncationReason = "maxInstances" + return nil + end + if not isInDataModel(instance) then + return nil + end + visited += 1 + local reference, referenceError = context.references:reference(instance, path) + if reference == nil then + return nil, referenceError + end + local node = { + reference = reference, + name = instance.Name, + className = instance.ClassName, + path = path, + properties = if request.includeProperties then readProperties(instance, context.references, path) else {}, + attributes = if request.includeAttributes then readAttributes(instance, context.references, path) else {}, + tags = if request.includeTags then readTags(instance) else {}, + children = {}, + truncated = false, + } + if depth >= request.depth then + return node + end + local children = instance:GetChildren() + local count = math.min(#children, request.maxChildren) + if #children > count then + node.truncated = true + globallyTruncated = true + truncationReason = truncationReason or "maxChildren" + end + for index = 1, count do + local child = visit(children[index], childPath(path, children[index]), depth + 1) + if child == nil then + node.truncated = true + break + end + table.insert(node.children, child) + end + return node + end + + local root, traversalError = visit(target, rootPath, 0) + if root == nil then + return nil, traversalError or "Inspect target was destroyed during traversal" + end + return { + root = root, + visitedInstances = visited, + truncated = globallyTruncated, + truncationReason = truncationReason, + } +end + +Inspect._test = { + quoteSegment = quoteSegment, + resolveTarget = resolveTarget, +} + +return Inspect diff --git a/plugin/src/AutomationHandlers/Inspect.spec.lua b/plugin/src/AutomationHandlers/Inspect.spec.lua new file mode 100644 index 00000000..26a48835 --- /dev/null +++ b/plugin/src/AutomationHandlers/Inspect.spec.lua @@ -0,0 +1,75 @@ +return function() + local CollectionService = game:GetService("CollectionService") + local Inspect = require(script.Parent.Inspect) + local InstanceReferences = require(script.Parent.Parent.InstanceReferences) + + local root + local function request(depth, maxChildren, maxInstances) + return { + kind = "inspect", + target = { kind = "path", segments = { "Workspace", root.Name } }, + depth = depth, + maxChildren = maxChildren or 100, + maxInstances = maxInstances or 2_000, + includeProperties = true, + includeAttributes = true, + includeTags = true, + } + end + + describe("Inspect handler", function() + beforeEach(function() + root = Instance.new("Folder") + root.Name = "PrismInspectSpec" + root.Parent = workspace + end) + + afterEach(function() + root:Destroy() + end) + + it("resolves roots, children, and quoted diagnostics", function() + local child = Instance.new("Part") + child.Name = "Name.With.Dots" + child.Parent = root + local result = Inspect.run(request(1), { references = InstanceReferences.new("session") }) + expect(result.root.children[1].path:find('"Name.With.Dots"', 1, true)).to.be.ok() + expect(result.root.children[1].className).to.equal("Part") + end) + + it("honors depth, child, and total limits", function() + for index = 1, 3 do + local child = Instance.new("Folder") + child.Name = tostring(index) + child.Parent = root + end + local depthZero = Inspect.run(request(0), { references = InstanceReferences.new("session") }) + expect(#depthZero.root.children).to.equal(0) + local childLimited = Inspect.run(request(1, 1), { references = InstanceReferences.new("session") }) + expect(#childLimited.root.children).to.equal(1) + expect(childLimited.truncated).to.equal(true) + local totalLimited = Inspect.run(request(1, 100, 1), { references = InstanceReferences.new("session") }) + expect(totalLimited.truncationReason).to.equal("maxInstances") + end) + + it("sorts attributes and tags and uses the property allowlist", function() + root:SetAttribute("Zed", 1) + root:SetAttribute("Alpha", true) + CollectionService:AddTag(root, "Zulu") + CollectionService:AddTag(root, "Alpha") + local result = Inspect.run(request(0), { references = InstanceReferences.new("session") }) + expect(result.root.attributes.Alpha.kind).to.equal("boolean") + expect(result.root.tags[1]).to.equal("Alpha") + expect(result.root.properties.Archivable.kind).to.equal("boolean") + expect(result.root.properties.Source).never.to.be.ok() + end) + + it("reports missing and destroyed targets", function() + local missing = request(0) + missing.target.segments[2] = "MissingPrismInspectTarget" + expect(select(1, Inspect.run(missing, { references = InstanceReferences.new("session") }))).never.to.be.ok() + root:Destroy() + expect(select(1, Inspect.run(request(0), { references = InstanceReferences.new("session") }))).never.to.be.ok() + end) + end) +end diff --git a/plugin/src/AutomationValue.lua b/plugin/src/AutomationValue.lua new file mode 100644 index 00000000..3f7758d9 --- /dev/null +++ b/plugin/src/AutomationValue.lua @@ -0,0 +1,143 @@ +local MAX_DEPTH = 32 +local MAX_ELEMENTS = 10_000 +local MAX_STRING_BYTES = 64 * 1024 + +local AutomationValue = {} + +local function finite(value) + return value == value and value ~= math.huge and value ~= -math.huge +end + +function AutomationValue.encode(value, references, diagnosticPath: string?) + local seen = {} + local elements = 0 + + local function encode(current, depth, path) + elements += 1 + if elements > MAX_ELEMENTS then + return nil, "Automation value exceeds the element limit" + end + if depth > MAX_DEPTH then + return nil, "Automation value exceeds the depth limit" + end + + local valueType = typeof(current) + if valueType == "nil" then + return { kind = "nil" } + elseif valueType == "boolean" then + return { kind = "boolean", value = current } + elseif valueType == "number" then + if not finite(current) then + return nil, "Automation value contains a non-finite number at " .. path + end + return { kind = "number", value = current } + elseif valueType == "string" then + if #current > MAX_STRING_BYTES then + return nil, "Automation string exceeds the 64 KiB limit at " .. path + end + return { kind = "string", value = current } + elseif valueType == "Vector2" then + return { kind = "vector2", x = current.X, y = current.Y } + elseif valueType == "Vector3" then + return { kind = "vector3", x = current.X, y = current.Y, z = current.Z } + elseif valueType == "Color3" then + return { kind = "color3", r = current.R, g = current.G, b = current.B } + elseif valueType == "CFrame" then + return { kind = "cFrame", components = { current:GetComponents() } } + elseif valueType == "UDim" then + return { kind = "uDim", scale = current.Scale, offset = current.Offset } + elseif valueType == "UDim2" then + return { + kind = "uDim2", + x = { scale = current.X.Scale, offset = current.X.Offset }, + y = { scale = current.Y.Scale, offset = current.Y.Offset }, + } + elseif valueType == "Rect" then + return { + kind = "rect", + min = { x = current.Min.X, y = current.Min.Y }, + max = { x = current.Max.X, y = current.Max.Y }, + } + elseif valueType == "EnumItem" then + return { kind = "enumItem", enumType = tostring(current.EnumType):gsub("^Enum%.", ""), name = current.Name } + elseif valueType == "BrickColor" then + return { kind = "brickColor", number = current.Number, name = current.Name } + elseif valueType == "NumberRange" then + return { kind = "numberRange", min = current.Min, max = current.Max } + elseif valueType == "Instance" then + if references == nil then + return nil, "Instance value cannot be encoded without a reference registry" + end + local reference, referenceError = references:reference(current, diagnosticPath or path) + if reference == nil then + return nil, referenceError + end + return { kind = "instanceReference", value = reference } + elseif valueType ~= "table" then + return nil, string.format("Unsupported automation value type '%s' at %s", valueType, path) + end + + if seen[current] then + return nil, "Automation value contains a cycle at " .. path + end + seen[current] = true + local numericCount = 0 + local maximumIndex = 0 + local stringKeys = {} + for key in current do + if type(key) == "number" and key >= 1 and key % 1 == 0 then + numericCount += 1 + maximumIndex = math.max(maximumIndex, key) + elseif type(key) == "string" then + table.insert(stringKeys, key) + else + seen[current] = nil + return nil, "Automation tables require positive integer or string keys" + end + end + if numericCount > 0 and #stringKeys > 0 then + seen[current] = nil + return nil, "Automation tables cannot mix array and dictionary keys" + end + + local encoded + if numericCount > 0 then + if maximumIndex ~= numericCount then + seen[current] = nil + return nil, "Automation arrays must be dense" + end + local values = table.create(numericCount) + for index = 1, numericCount do + local child, childError = encode(current[index], depth + 1, string.format("%s[%d]", path, index)) + if child == nil then + seen[current] = nil + return nil, childError + end + values[index] = child + end + encoded = { kind = "array", value = values } + else + table.sort(stringKeys) + local values = table.create(#stringKeys) + for index, key in stringKeys do + local child, childError = encode(current[key], depth + 1, path .. "." .. key) + if child == nil then + seen[current] = nil + return nil, childError + end + values[index] = { key = key, value = child } + end + encoded = { kind = "map", value = values } + end + seen[current] = nil + return encoded + end + + return encode(value, 1, "value") +end + +function AutomationValue.diagnostic(errorValue) + return { kind = "diagnostic", error = tostring(errorValue) } +end + +return AutomationValue diff --git a/plugin/src/AutomationValue.spec.lua b/plugin/src/AutomationValue.spec.lua new file mode 100644 index 00000000..04a65895 --- /dev/null +++ b/plugin/src/AutomationValue.spec.lua @@ -0,0 +1,40 @@ +return function() + local AutomationValue = require(script.Parent.AutomationValue) + + describe("AutomationValue", function() + it("encodes scalar and Roblox values", function() + expect(AutomationValue.encode(nil).kind).to.equal("nil") + expect(AutomationValue.encode(true).kind).to.equal("boolean") + expect(AutomationValue.encode(4).kind).to.equal("number") + expect(AutomationValue.encode("value").kind).to.equal("string") + expect(AutomationValue.encode(Vector3.new(1, 2, 3)).kind).to.equal("vector3") + expect(AutomationValue.encode(CFrame.new()).kind).to.equal("cFrame") + expect(AutomationValue.encode(Color3.new()).kind).to.equal("color3") + expect(AutomationValue.encode(UDim2.new()).kind).to.equal("uDim2") + expect(AutomationValue.encode(Rect.new(Vector2.zero, Vector2.one)).kind).to.equal("rect") + expect(AutomationValue.encode(Enum.Material.Plastic).kind).to.equal("enumItem") + expect(AutomationValue.encode(BrickColor.new("Bright red")).kind).to.equal("brickColor") + expect(AutomationValue.encode(NumberRange.new(1, 2)).kind).to.equal("numberRange") + end) + + it("encodes dense arrays and sorted maps", function() + local array = AutomationValue.encode({ "a", "b" }) + expect(array.kind).to.equal("array") + expect(array.value[2].value).to.equal("b") + + local map = AutomationValue.encode({ z = 1, a = 2 }) + expect(map.kind).to.equal("map") + expect(map.value[1].key).to.equal("a") + expect(map.value[2].key).to.equal("z") + end) + + it("rejects cycles, sparse values, unsupported values, and non-finite numbers", function() + local cyclic = {} + cyclic.self = cyclic + expect(select(1, AutomationValue.encode(cyclic))).never.to.be.ok() + expect(select(1, AutomationValue.encode({ [2] = true }))).never.to.be.ok() + expect(select(1, AutomationValue.encode(function() end))).never.to.be.ok() + expect(select(1, AutomationValue.encode(math.huge))).never.to.be.ok() + end) + end) +end diff --git a/plugin/src/Exec.lua b/plugin/src/Exec.lua index f74dbdd6..c7be58fe 100644 --- a/plugin/src/Exec.lua +++ b/plugin/src/Exec.lua @@ -565,14 +565,36 @@ end local Exec = {} Exec.__index = Exec +local function classifyStudioMode(isEdit, isRunMode, isRunning) + if isEdit then + return "edit" + elseif isRunMode then + return "run" + elseif isRunning then + return "play" + else + return "unknown" + end +end + +local function currentStudioMode() + return classifyStudioMode(RunService:IsEdit(), RunService:IsRunMode(), RunService:IsRunning()) +end + function Exec.new(options) assert(type(options) == "table", "Exec options must be a table") assert(type(options.apiContext) == "table", "Exec apiContext must be a table") local dependencies = options.dependencies or {} + local isEdit = dependencies.isEdit or defaultExecutionDependencies.isEdit local self = { __apiContext = options.apiContext, - __isEdit = dependencies.isEdit or defaultExecutionDependencies.isEdit, + __studioMode = dependencies.studioMode or function() + if dependencies.isEdit ~= nil then + return if isEdit() then "edit" else "play" + end + return currentStudioMode() + end, __delay = dependencies.delay or Promise.delay, __execute = dependencies.execute or function(job) return startExecution(job, defaultExecutionDependencies, EXECUTION_TIMEOUT_SECONDS) @@ -636,7 +658,8 @@ function Exec:__complete(generation, jobId, payload, attempt) return end - local requestOk, completionPromise = pcall(self.__apiContext.completeExecJob, self.__apiContext, jobId, payload) + local requestOk, completionPromise = + pcall(self.__apiContext.completeExecJob, self.__apiContext, jobId, payload, self.__studioMode()) if not requestOk then self:__fail(generation, completionPromise) return @@ -712,13 +735,14 @@ function Exec:__poll(generation) return end - if not self.__isEdit() then + local studioMode = self.__studioMode() + if studioMode ~= "edit" then self:__schedulePoll(generation, POLL_INTERVAL_SECONDS) return end self.__busy = true - local claimOk, claimPromise = pcall(self.__apiContext.claimNextExecJob, self.__apiContext) + local claimOk, claimPromise = pcall(self.__apiContext.claimNextExecJob, self.__apiContext, studioMode) if not claimOk then self:__fail(generation, claimPromise) return @@ -785,6 +809,7 @@ Exec._test = { runJob = runJob, sanitizeScriptName = sanitizeScriptName, startExecution = startExecution, + classifyStudioMode = classifyStudioMode, constants = { completionAttempts = COMPLETION_ATTEMPTS, maxCompletionBodyBytes = MAX_COMPLETION_BODY_BYTES, diff --git a/plugin/src/Exec.spec.lua b/plugin/src/Exec.spec.lua index 8f327ce7..c4d70029 100644 --- a/plugin/src/Exec.spec.lua +++ b/plugin/src/Exec.spec.lua @@ -394,6 +394,43 @@ return function() end) describe("Exec poller", function() + it("classifies edit, play, run, and unknown modes", function() + expect(helpers.classifyStudioMode(true, false, false)).to.equal("edit") + expect(helpers.classifyStudioMode(false, false, true)).to.equal("play") + expect(helpers.classifyStudioMode(false, true, true)).to.equal("run") + expect(helpers.classifyStudioMode(false, false, false)).to.equal("unknown") + end) + + it("includes Studio mode in claim and completion calls", function() + local claimMode + local completionMode + local poller = Exec.new({ + apiContext = { + claimNextExecJob = function(_, mode) + claimMode = mode + return Promise.resolve(claimedJob()) + end, + completeExecJob = function(_, _jobId, _payload, mode) + completionMode = mode + return Promise.resolve({ status = "accepted" }) + end, + }, + dependencies = { + studioMode = function() + return "edit" + end, + execute = function() + return Promise.resolve({ outcome = "success", logs = {} }) + end, + }, + }) + + poller:start() + expect(claimMode).to.equal("edit") + expect(completionMode).to.equal("edit") + poller:stop() + end) + it("starts only once", function() local scheduler = makeScheduler() local claimCount = 0 diff --git a/plugin/src/InstanceReferences.lua b/plugin/src/InstanceReferences.lua new file mode 100644 index 00000000..76e442e6 --- /dev/null +++ b/plugin/src/InstanceReferences.lua @@ -0,0 +1,57 @@ +local InstanceReferences = {} +InstanceReferences.__index = InstanceReferences + +function InstanceReferences.new(sessionId: string) + assert(type(sessionId) == "string" and sessionId ~= "", "sessionId must be a non-empty string") + return setmetatable({ + __sessionId = sessionId, + __nextId = 1, + __byId = setmetatable({}, { __mode = "v" }), + __byInstance = setmetatable({}, { __mode = "k" }), + }, InstanceReferences) +end + +local function isInDataModel(instance) + local ok, result = pcall(function() + return instance == game or instance:IsDescendantOf(game) + end) + return ok and result +end + +function InstanceReferences:reference(instance, path: string) + if typeof(instance) ~= "Instance" or not isInDataModel(instance) then + return nil, "Instance is destroyed or outside the current DataModel" + end + + local id = self.__byInstance[instance] + if id == nil then + id = string.format("pinst-%08d", self.__nextId) + self.__nextId += 1 + self.__byInstance[instance] = id + self.__byId[id] = instance + end + + return { + sessionId = self.__sessionId, + id = id, + path = path, + name = instance.Name, + className = instance.ClassName, + } +end + +function InstanceReferences:resolve(id: string) + local instance = self.__byId[id] + if instance == nil or not isInDataModel(instance) then + self.__byId[id] = nil + return nil, "Instance reference is stale" + end + return instance +end + +function InstanceReferences:clear() + table.clear(self.__byId) + table.clear(self.__byInstance) +end + +return InstanceReferences diff --git a/plugin/src/InstanceReferences.spec.lua b/plugin/src/InstanceReferences.spec.lua new file mode 100644 index 00000000..06d13715 --- /dev/null +++ b/plugin/src/InstanceReferences.spec.lua @@ -0,0 +1,29 @@ +return function() + local InstanceReferences = require(script.Parent.InstanceReferences) + + describe("InstanceReferences", function() + it("creates stable session-local references", function() + local folder = Instance.new("Folder") + folder.Parent = workspace + local references = InstanceReferences.new("plugin-session") + local first = references:reference(folder, "Workspace.Folder") + local second = references:reference(folder, "Workspace.Renamed") + expect(first.id).to.equal("pinst-00000001") + expect(second.id).to.equal(first.id) + expect(first.sessionId).to.equal("plugin-session") + expect(references:resolve(first.id)).to.equal(folder) + folder:Destroy() + expect(references:resolve(first.id)).never.to.be.ok() + end) + + it("clears the registry", function() + local folder = Instance.new("Folder") + folder.Parent = workspace + local references = InstanceReferences.new("plugin-session") + local reference = references:reference(folder, "Workspace.Folder") + references:clear() + expect(references:resolve(reference.id)).never.to.be.ok() + folder:Destroy() + end) + end) +end diff --git a/plugin/src/ServeSession.lua b/plugin/src/ServeSession.lua index 58aef5de..a6d56d96 100644 --- a/plugin/src/ServeSession.lua +++ b/plugin/src/ServeSession.lua @@ -13,6 +13,7 @@ local Timer = require(script.Parent.Timer) local ChangeBatcher = require(script.Parent.ChangeBatcher) local encodePatchUpdate = require(script.Parent.ChangeBatcher.encodePatchUpdate) +local Automation = require(script.Parent.Automation) local Exec = require(script.Parent.Exec) local InstanceMap = require(script.Parent.InstanceMap) local PatchSet = require(script.Parent.PatchSet) @@ -28,6 +29,121 @@ local Status = strict("Session.Status", { Disconnected = "Disconnected", }) +local AUTOMATION_HEARTBEAT_INTERVAL_SECONDS = 3 + +local function currentStudioMode() + if RunService:IsEdit() then + return "edit" + elseif RunService:IsRunMode() then + return "run" + elseif RunService:IsRunning() then + return "play" + else + return "unknown" + end +end + +local AutomationHeartbeat = {} +AutomationHeartbeat.__index = AutomationHeartbeat + +function AutomationHeartbeat.new(options) + return setmetatable({ + __apiContext = options.apiContext, + __delay = options.delay or Promise.delay, + __getStudioMode = options.getStudioMode or currentStudioMode, + __onError = options.onError, + __onReady = options.onReady or function() end, + __running = false, + __ready = false, + __generation = 0, + __requestPromise = nil, + __scheduledPromise = nil, + }, AutomationHeartbeat) +end + +function AutomationHeartbeat:__isCurrent(generation) + return self.__running and self.__generation == generation +end + +function AutomationHeartbeat:__schedule(generation) + if not self:__isCurrent(generation) then + return + end + + self.__scheduledPromise = self.__delay(AUTOMATION_HEARTBEAT_INTERVAL_SECONDS):andThen(function() + if self:__isCurrent(generation) then + self:__send(generation) + end + end) +end + +function AutomationHeartbeat:__send(generation) + if not self:__isCurrent(generation) then + return + end + + local ok, requestPromise = + pcall(self.__apiContext.updateAutomationStatus, self.__apiContext, self.__getStudioMode()) + if not ok then + self.__onError(requestPromise) + return + end + + self.__requestPromise = requestPromise + requestPromise + :andThen(function(response) + if not self:__isCurrent(generation) then + return + end + + self.__requestPromise = nil + if response.registration == "conflict" then + self.__onError("Another active Prism plugin session is already registered for automation") + return + end + if not self.__ready then + self.__ready = true + self.__onReady() + end + + self:__schedule(generation) + end) + :catch(function(errorValue) + if self:__isCurrent(generation) then + self.__onError(errorValue) + end + end) +end + +function AutomationHeartbeat:start() + if self.__running then + return + end + + self.__running = true + self.__ready = false + self.__generation += 1 + self:__send(self.__generation) +end + +function AutomationHeartbeat:stop() + if not self.__running then + return + end + + self.__running = false + self.__ready = false + self.__generation += 1 + if self.__scheduledPromise ~= nil then + self.__scheduledPromise:cancel() + self.__scheduledPromise = nil + end + if self.__requestPromise ~= nil then + self.__requestPromise:cancel() + self.__requestPromise = nil + end +end + local function debugPatch(object) return Fmt.debugify(object, function(patch, output) output:writeLine("Patch {{") @@ -122,6 +238,27 @@ function ServeSession.new(options) end end, }) + self.__automation = Automation.new({ + apiContext = self.__apiContext, + onError = function(errorValue) + if self.__status ~= Status.Disconnected then + self:__stopInternal(errorValue) + end + end, + }) + self.__automationHeartbeat = AutomationHeartbeat.new({ + apiContext = self.__apiContext, + onReady = function() + if self.__status == Status.Connected then + self.__automation:start() + end + end, + onError = function(errorValue) + if self.__status ~= Status.Disconnected then + self:__stopInternal(errorValue) + end + end, + }) return self end @@ -229,6 +366,7 @@ function ServeSession:start() -- connectWebSocket creates the stream client synchronously. Starting -- here keeps exec on the same stable, fully-synced session lifecycle. + self.__automationHeartbeat:start() self.__exec:start() return websocketPromise end) @@ -584,6 +722,8 @@ function ServeSession:__initialSync(serverInfo) end function ServeSession:__stopInternal(err) + self.__automationHeartbeat:stop() + self.__automation:stop() self.__exec:stop() self:__setStatus(Status.Disconnected, err) self.__apiContext:disconnect() @@ -604,4 +744,9 @@ function ServeSession:__setStatus(status, detail) end end +ServeSession._test = { + AutomationHeartbeat = AutomationHeartbeat, + currentStudioMode = currentStudioMode, +} + return ServeSession diff --git a/plugin/src/ServeSession.spec.lua b/plugin/src/ServeSession.spec.lua new file mode 100644 index 00000000..446576d1 --- /dev/null +++ b/plugin/src/ServeSession.spec.lua @@ -0,0 +1,153 @@ +return function() + local Packages = script.Parent.Parent.Packages + local Promise = require(Packages.Promise) + local ServeSession = require(script.Parent.ServeSession) + local AutomationHeartbeat = ServeSession._test.AutomationHeartbeat + + local function makeScheduler() + local queue = {} + return { + delay = function() + return Promise.new(function(resolve) + table.insert(queue, resolve) + end) + end, + step = function() + local resolve = table.remove(queue, 1) + assert(resolve ~= nil, "no heartbeat delay scheduled") + resolve() + end, + count = function() + return #queue + end, + } + end + + describe("automation heartbeat", function() + it("signals readiness once after registration", function() + local readyCount = 0 + local heartbeat = AutomationHeartbeat.new({ + apiContext = { + updateAutomationStatus = function() + return Promise.resolve({ registration = "registered" }) + end, + }, + delay = function() + return Promise.new(function() end) + end, + getStudioMode = function() + return "edit" + end, + onReady = function() + readyCount += 1 + end, + onError = function(errorValue) + error(errorValue) + end, + }) + heartbeat:start() + expect(readyCount).to.equal(1) + heartbeat:stop() + end) + + it("starts once and reports the current mode", function() + local scheduler = makeScheduler() + local updateCount = 0 + local reportedMode + local heartbeat = AutomationHeartbeat.new({ + apiContext = { + updateAutomationStatus = function(_, mode) + updateCount += 1 + reportedMode = mode + return Promise.resolve({ registration = "registered" }) + end, + }, + delay = scheduler.delay, + getStudioMode = function() + return "run" + end, + onError = function(errorValue) + error(errorValue) + end, + }) + + heartbeat:start() + heartbeat:start() + expect(updateCount).to.equal(1) + expect(reportedMode).to.equal("run") + expect(scheduler.count()).to.equal(1) + heartbeat:stop() + end) + + it("stops scheduling updates after disconnect", function() + local scheduler = makeScheduler() + local updateCount = 0 + local heartbeat = AutomationHeartbeat.new({ + apiContext = { + updateAutomationStatus = function() + updateCount += 1 + return Promise.resolve({ registration = "refreshed" }) + end, + }, + delay = scheduler.delay, + getStudioMode = function() + return "edit" + end, + onError = function(errorValue) + error(errorValue) + end, + }) + + heartbeat:start() + heartbeat:stop() + scheduler.step() + expect(updateCount).to.equal(1) + end) + + it("ignores a late response from a stopped generation", function() + local resolveHeartbeat + local errorCount = 0 + local heartbeat = AutomationHeartbeat.new({ + apiContext = { + updateAutomationStatus = function() + return Promise.new(function(resolve) + resolveHeartbeat = resolve + end) + end, + }, + getStudioMode = function() + return "edit" + end, + onError = function() + errorCount += 1 + end, + }) + + heartbeat:start() + heartbeat:stop() + resolveHeartbeat({ registration = "conflict" }) + expect(errorCount).to.equal(0) + end) + + it("surfaces duplicate-session conflicts clearly", function() + local capturedError + local heartbeat = AutomationHeartbeat.new({ + apiContext = { + updateAutomationStatus = function() + return Promise.resolve({ registration = "conflict" }) + end, + }, + getStudioMode = function() + return "edit" + end, + onError = function(errorValue) + capturedError = errorValue + end, + }) + + heartbeat:start() + expect(capturedError:find("already registered", 1, true)).to.be.ok() + heartbeat:stop() + end) + end) +end diff --git a/plugin/src/Types.lua b/plugin/src/Types.lua index 52db9886..61e7fe26 100644 --- a/plugin/src/Types.lua +++ b/plugin/src/Types.lua @@ -96,6 +96,54 @@ local ExecJobState = t.union( t.literal("timedOut") ) +local StudioMode = t.union(t.literal("edit"), t.literal("play"), t.literal("run"), t.literal("unknown")) + +local AutomationRegistration = t.union(t.literal("registered"), t.literal("refreshed"), t.literal("conflict")) + +local ApiAutomationHeartbeatResponse = t.strictInterface({ + registration = AutomationRegistration, + activePluginSessionId = t.optional(Uuid), +}) + +local AutomationJobState = t.union( + t.literal("pending"), + t.literal("claimed"), + t.literal("succeeded"), + t.literal("failed"), + t.literal("timedOut") +) + +local InspectTarget = t.strictInterface({ + kind = t.literal("path"), + segments = t.array(t.string), +}) + +local InspectRequest = t.strictInterface({ + kind = t.literal("inspect"), + target = InspectTarget, + depth = t.number, + maxChildren = t.number, + maxInstances = t.number, + includeProperties = t.boolean, + includeAttributes = t.boolean, + includeTags = t.boolean, +}) + +local ApiAutomationClaimResponse = t.strictInterface({ + jobId = Uuid, + state = t.literal("claimed"), + request = InspectRequest, + executionTimeoutMs = t.number, +}) + +local ApiAutomationJobResponse = t.strictInterface({ + jobId = Uuid, + state = AutomationJobState, + claimedByPluginSessionId = t.optional(Uuid), + result = t.optional(t.table), + error = t.optional(t.string), +}) + local ExecLogLevel = t.union(t.literal("print"), t.literal("warn")) local ExecLog = t.strictInterface({ @@ -201,6 +249,12 @@ return strict("Types", { ApiRefPatchResponse = ApiRefPatchResponse, ApiExecClaimResponse = ApiExecClaimResponse, ApiExecJobResponse = ApiExecJobResponse, + ApiAutomationHeartbeatResponse = ApiAutomationHeartbeatResponse, + ApiAutomationClaimResponse = ApiAutomationClaimResponse, + ApiAutomationJobResponse = ApiAutomationJobResponse, + AutomationJobState = AutomationJobState, + AutomationRegistration = AutomationRegistration, + StudioMode = StudioMode, ExecJobState = ExecJobState, ExecValue = ExecValue, ExecLog = ExecLog, diff --git a/src/automation.rs b/src/automation.rs new file mode 100644 index 00000000..2088a8e0 --- /dev/null +++ b/src/automation.rs @@ -0,0 +1,680 @@ +use std::{ + collections::{BTreeMap, HashMap, VecDeque}, + sync::Mutex, + time::{Duration, Instant}, +}; + +use serde::{Deserialize, Serialize}; +use thiserror::Error; +use uuid::Uuid; + +use crate::automation_status::PluginSessionId; + +pub const MAX_PENDING_AUTOMATION_JOBS: usize = 16; +pub const MAX_RETAINED_AUTOMATION_JOBS: usize = 64; +pub const AUTOMATION_CLAIM_TIMEOUT: Duration = Duration::from_secs(30); +pub const INSPECT_EXECUTION_TIMEOUT: Duration = Duration::from_secs(15); +pub const AUTOMATION_TERMINAL_RETENTION: Duration = Duration::from_secs(5 * 60); +pub const MAX_AUTOMATION_REQUEST_BODY_BYTES: usize = 256 * 1024; +pub const MAX_AUTOMATION_RESULT_BODY_BYTES: usize = 4 * 1024 * 1024; +pub const MAX_INSPECT_DEPTH: u8 = 8; +pub const MAX_INSPECT_CHILDREN: u32 = 1_000; +pub const MAX_INSPECT_INSTANCES: u32 = 10_000; +pub const MAX_AUTOMATION_STRING_BYTES: usize = 64 * 1024; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub enum AutomationJobState { + Pending, + Claimed, + Succeeded, + Failed, + TimedOut, +} + +impl AutomationJobState { + pub fn is_terminal(self) -> bool { + matches!(self, Self::Succeeded | Self::Failed | Self::TimedOut) + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(tag = "kind", rename_all = "camelCase")] +pub enum AutomationRequest { + Inspect(InspectRequest), +} + +impl AutomationRequest { + fn execution_timeout(&self) -> Duration { + match self { + Self::Inspect(_) => INSPECT_EXECUTION_TIMEOUT, + } + } + + fn accepts_result(&self, result: &AutomationResult) -> bool { + matches!( + (self, result), + (Self::Inspect(_), AutomationResult::Inspect(_)) + ) + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(tag = "kind", rename_all = "camelCase")] +pub enum InspectTarget { + Path { segments: Vec }, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct InspectRequest { + pub target: InspectTarget, + pub depth: u8, + pub max_children: u32, + pub max_instances: u32, + pub include_properties: bool, + pub include_attributes: bool, + pub include_tags: bool, +} + +impl InspectRequest { + pub fn validate(&self) -> Result<(), AutomationJobStoreError> { + let InspectTarget::Path { segments } = &self.target; + if segments.is_empty() || segments.iter().any(String::is_empty) { + return Err(AutomationJobStoreError::InvalidRequest( + "inspect target must contain at least one non-empty path segment".to_owned(), + )); + } + if self.depth > MAX_INSPECT_DEPTH { + return Err(AutomationJobStoreError::InvalidRequest(format!( + "inspect depth {} exceeds the maximum of {MAX_INSPECT_DEPTH}", + self.depth + ))); + } + if self.max_children == 0 || self.max_children > MAX_INSPECT_CHILDREN { + return Err(AutomationJobStoreError::InvalidRequest(format!( + "inspect maxChildren must be between 1 and {MAX_INSPECT_CHILDREN}" + ))); + } + if self.max_instances == 0 || self.max_instances > MAX_INSPECT_INSTANCES { + return Err(AutomationJobStoreError::InvalidRequest(format!( + "inspect maxInstances must be between 1 and {MAX_INSPECT_INSTANCES}" + ))); + } + Ok(()) + } +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(tag = "kind", rename_all = "camelCase")] +pub enum AutomationResult { + Inspect(InspectResult), +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct InspectResult { + pub root: InspectNode, + pub visited_instances: u32, + pub truncated: bool, + #[serde(skip_serializing_if = "Option::is_none")] + pub truncation_reason: Option, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct InspectNode { + pub reference: InstanceReference, + pub name: String, + pub class_name: String, + pub path: String, + #[serde(default, skip_serializing_if = "BTreeMap::is_empty")] + pub properties: BTreeMap, + #[serde(default, skip_serializing_if = "BTreeMap::is_empty")] + pub attributes: BTreeMap, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub tags: Vec, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub children: Vec, + pub truncated: bool, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct InstanceReference { + pub session_id: String, + pub id: String, + pub path: String, + pub name: String, + pub class_name: String, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(tag = "kind", rename_all = "camelCase")] +pub enum AutomationValue { + Nil, + Boolean { + value: bool, + }, + Number { + value: f64, + }, + String { + value: String, + }, + Array { + value: Vec, + }, + Map { + value: Vec, + }, + Vector2 { + x: f64, + y: f64, + }, + Vector3 { + x: f64, + y: f64, + z: f64, + }, + CFrame { + components: Vec, + }, + Color3 { + r: f64, + g: f64, + b: f64, + }, + UDim { + scale: f64, + offset: i32, + }, + UDim2 { + x: AutomationUdim, + y: AutomationUdim, + }, + Rect { + min: AutomationVector2, + max: AutomationVector2, + }, + EnumItem { + enum_type: String, + name: String, + }, + BrickColor { + number: u16, + name: String, + }, + NumberRange { + min: f64, + max: f64, + }, + InstanceReference { + value: InstanceReference, + }, + Diagnostic { + error: String, + }, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct AutomationMapEntry { + pub key: String, + pub value: AutomationValue, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct AutomationUdim { + pub scale: f64, + pub offset: i32, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct AutomationVector2 { + pub x: f64, + pub y: f64, +} + +#[derive(Debug, Clone, PartialEq)] +pub struct AutomationJob { + pub id: Uuid, + pub request: AutomationRequest, + pub state: AutomationJobState, + pub claimed_by: Option, + pub execution_timeout: Duration, + pub result: Option, + pub error: Option, + claim_deadline: Instant, + execution_deadline: Option, + terminal_at: Option, +} + +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct AutomationQueueCounts { + pub pending: usize, + pub claimed: usize, +} + +#[derive(Debug, Error, PartialEq, Eq)] +pub enum AutomationJobStoreError { + #[error("invalid automation request: {0}")] + InvalidRequest(String), + #[error("automation pending queue is full (limit: {limit})")] + PendingQueueFull { limit: usize }, + #[error("automation job {id} does not exist")] + UnknownJob { id: Uuid }, + #[error("automation job {id} cannot be completed from state {state:?}")] + JobNotClaimed { id: Uuid, state: AutomationJobState }, + #[error("automation job {id} was already completed with state {state:?}")] + DuplicateCompletion { id: Uuid, state: AutomationJobState }, + #[error("automation job {id} was claimed by a different plugin session")] + WrongClaimant { id: Uuid }, + #[error("automation result kind does not match job {id}")] + ResultKindMismatch { id: Uuid }, +} + +#[derive(Debug)] +pub struct AutomationJobStore { + inner: Mutex, + claim_timeout: Duration, + terminal_retention: Duration, +} + +#[derive(Debug, Default)] +struct AutomationJobStoreInner { + jobs: HashMap, + pending: VecDeque, + claimed: Option, + terminal: VecDeque, +} + +impl AutomationJobStore { + pub fn new() -> Self { + Self::with_durations(AUTOMATION_CLAIM_TIMEOUT, AUTOMATION_TERMINAL_RETENTION) + } + + pub fn submit( + &self, + request: AutomationRequest, + ) -> Result { + self.submit_at(request, Instant::now()) + } + + pub fn claim_next(&self, claimant: PluginSessionId) -> Option { + self.claim_next_at(claimant, Instant::now()) + } + + pub fn complete_success( + &self, + id: Uuid, + claimant: PluginSessionId, + result: AutomationResult, + ) -> Result { + self.complete_at(id, claimant, Some(result), None, Instant::now()) + } + + pub fn complete_failure( + &self, + id: Uuid, + claimant: PluginSessionId, + error: String, + ) -> Result { + self.complete_at(id, claimant, None, Some(error), Instant::now()) + } + + pub fn get(&self, id: Uuid) -> Result { + self.inner + .lock() + .unwrap() + .jobs + .get(&id) + .cloned() + .ok_or(AutomationJobStoreError::UnknownJob { id }) + } + + pub fn queue_counts(&self) -> AutomationQueueCounts { + let inner = self.inner.lock().unwrap(); + AutomationQueueCounts { + pending: inner.pending.len(), + claimed: usize::from(inner.claimed.is_some()), + } + } + + pub fn claimed_by(&self) -> Option { + let inner = self.inner.lock().unwrap(); + inner + .claimed + .and_then(|id| inner.jobs.get(&id)) + .and_then(|job| job.claimed_by) + } + + pub fn cleanup_expired(&self) { + self.cleanup_expired_at(Instant::now()); + } + + fn with_durations(claim_timeout: Duration, terminal_retention: Duration) -> Self { + Self { + inner: Mutex::new(AutomationJobStoreInner::default()), + claim_timeout, + terminal_retention, + } + } + + fn submit_at( + &self, + request: AutomationRequest, + now: Instant, + ) -> Result { + match &request { + AutomationRequest::Inspect(request) => request.validate()?, + } + let mut inner = self.inner.lock().unwrap(); + if inner.pending.len() + usize::from(inner.claimed.is_some()) >= MAX_PENDING_AUTOMATION_JOBS + { + return Err(AutomationJobStoreError::PendingQueueFull { + limit: MAX_PENDING_AUTOMATION_JOBS, + }); + } + let id = Uuid::new_v4(); + let execution_timeout = request.execution_timeout(); + let job = AutomationJob { + id, + request, + state: AutomationJobState::Pending, + claimed_by: None, + execution_timeout, + result: None, + error: None, + claim_deadline: now + self.claim_timeout, + execution_deadline: None, + terminal_at: None, + }; + inner.jobs.insert(id, job.clone()); + inner.pending.push_back(id); + Ok(job) + } + + fn claim_next_at(&self, claimant: PluginSessionId, now: Instant) -> Option { + let mut inner = self.inner.lock().unwrap(); + if inner.claimed.is_some() { + return None; + } + while let Some(id) = inner.pending.pop_front() { + let Some(job) = inner.jobs.get_mut(&id) else { + continue; + }; + if job.state != AutomationJobState::Pending { + continue; + } + job.state = AutomationJobState::Claimed; + job.claimed_by = Some(claimant); + job.execution_deadline = Some(now + job.execution_timeout); + let claimed = job.clone(); + inner.claimed = Some(id); + return Some(claimed); + } + None + } + + fn complete_at( + &self, + id: Uuid, + claimant: PluginSessionId, + result: Option, + error: Option, + now: Instant, + ) -> Result { + let mut inner = self.inner.lock().unwrap(); + let job = inner + .jobs + .get(&id) + .ok_or(AutomationJobStoreError::UnknownJob { id })?; + if job.state.is_terminal() { + return Err(AutomationJobStoreError::DuplicateCompletion { + id, + state: job.state, + }); + } + if job.state != AutomationJobState::Claimed { + return Err(AutomationJobStoreError::JobNotClaimed { + id, + state: job.state, + }); + } + if job.claimed_by != Some(claimant) { + return Err(AutomationJobStoreError::WrongClaimant { id }); + } + if let Some(result) = &result { + if !job.request.accepts_result(result) { + return Err(AutomationJobStoreError::ResultKindMismatch { id }); + } + } + + let job = inner.jobs.get_mut(&id).unwrap(); + job.state = if result.is_some() { + AutomationJobState::Succeeded + } else { + AutomationJobState::Failed + }; + job.result = result; + job.error = error; + job.terminal_at = Some(now); + let completed = job.clone(); + inner.claimed = None; + inner.terminal.push_back(id); + enforce_terminal_limit(&mut inner); + Ok(completed) + } + + fn cleanup_expired_at(&self, now: Instant) { + let mut inner = self.inner.lock().unwrap(); + let pending_count = inner.pending.len(); + for _ in 0..pending_count { + let id = inner.pending.pop_front().unwrap(); + let expired = inner.jobs.get(&id).is_some_and(|job| { + job.state == AutomationJobState::Pending && now >= job.claim_deadline + }); + if expired { + time_out(&mut inner, id, now); + } else if inner.jobs.contains_key(&id) { + inner.pending.push_back(id); + } + } + if let Some(id) = inner.claimed { + let expired = inner.jobs.get(&id).is_some_and(|job| { + job.execution_deadline + .is_some_and(|deadline| now >= deadline) + }); + if expired { + time_out(&mut inner, id, now); + inner.claimed = None; + } + } + while let Some(id) = inner.terminal.front().copied() { + let remove = inner.jobs.get(&id).is_none_or(|job| { + job.terminal_at.is_some_and(|terminal| { + now.saturating_duration_since(terminal) >= self.terminal_retention + }) + }); + if !remove { + break; + } + inner.terminal.pop_front(); + inner.jobs.remove(&id); + } + enforce_terminal_limit(&mut inner); + } +} + +impl Default for AutomationJobStore { + fn default() -> Self { + Self::new() + } +} + +fn time_out(inner: &mut AutomationJobStoreInner, id: Uuid, now: Instant) { + if let Some(job) = inner.jobs.get_mut(&id) { + job.state = AutomationJobState::TimedOut; + job.error = Some("automation job timed out".to_owned()); + job.terminal_at = Some(now); + inner.terminal.push_back(id); + enforce_terminal_limit(inner); + } +} + +fn enforce_terminal_limit(inner: &mut AutomationJobStoreInner) { + while inner.terminal.len() > MAX_RETAINED_AUTOMATION_JOBS { + if let Some(id) = inner.terminal.pop_front() { + inner.jobs.remove(&id); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn inspect(name: &str) -> AutomationRequest { + AutomationRequest::Inspect(InspectRequest { + target: InspectTarget::Path { + segments: vec![name.to_owned()], + }, + depth: 1, + max_children: 100, + max_instances: 2_000, + include_properties: false, + include_attributes: true, + include_tags: true, + }) + } + + fn result(name: &str) -> AutomationResult { + AutomationResult::Inspect(InspectResult { + root: InspectNode { + reference: InstanceReference { + session_id: "session".to_owned(), + id: "pinst-00000001".to_owned(), + path: name.to_owned(), + name: name.to_owned(), + class_name: "Workspace".to_owned(), + }, + name: name.to_owned(), + class_name: "Workspace".to_owned(), + path: name.to_owned(), + properties: BTreeMap::new(), + attributes: BTreeMap::new(), + tags: Vec::new(), + children: Vec::new(), + truncated: false, + }, + visited_instances: 1, + truncated: false, + truncation_reason: None, + }) + } + + #[test] + fn submit_claim_fifo_and_authoritative_completion() { + let store = AutomationJobStore::new(); + let first = store.submit(inspect("Workspace")).unwrap(); + let second = store.submit(inspect("Lighting")).unwrap(); + let claimant: PluginSessionId = Uuid::new_v4().to_string().parse().unwrap(); + let other: PluginSessionId = Uuid::new_v4().to_string().parse().unwrap(); + assert_eq!(store.claim_next(claimant).unwrap().id, first.id); + assert_eq!(store.claim_next(claimant), None); + assert_eq!( + store.complete_success(first.id, other, result("Workspace")), + Err(AutomationJobStoreError::WrongClaimant { id: first.id }) + ); + let completed = store + .complete_success(first.id, claimant, result("Workspace")) + .unwrap(); + assert_eq!(completed.state, AutomationJobState::Succeeded); + assert_eq!(store.claim_next(claimant).unwrap().id, second.id); + } + + #[test] + fn rejects_completion_before_claim_duplicate_and_unknown() { + let store = AutomationJobStore::new(); + let job = store.submit(inspect("Workspace")).unwrap(); + let claimant: PluginSessionId = Uuid::new_v4().to_string().parse().unwrap(); + assert!(matches!( + store.complete_failure(job.id, claimant, "no".to_owned()), + Err(AutomationJobStoreError::JobNotClaimed { .. }) + )); + store.claim_next(claimant).unwrap(); + store + .complete_failure(job.id, claimant, "failed".to_owned()) + .unwrap(); + assert!(matches!( + store.complete_failure(job.id, claimant, "again".to_owned()), + Err(AutomationJobStoreError::DuplicateCompletion { .. }) + )); + assert!(matches!( + store.get(Uuid::new_v4()), + Err(AutomationJobStoreError::UnknownJob { .. }) + )); + } + + #[test] + fn times_out_and_cleans_up_jobs() { + let store = + AutomationJobStore::with_durations(Duration::from_secs(1), Duration::from_secs(1)); + let now = Instant::now(); + let pending = store.submit_at(inspect("Workspace"), now).unwrap(); + store.cleanup_expired_at(now + Duration::from_secs(1)); + assert_eq!( + store.get(pending.id).unwrap().state, + AutomationJobState::TimedOut + ); + store.cleanup_expired_at(now + Duration::from_secs(2)); + assert!(matches!( + store.get(pending.id), + Err(AutomationJobStoreError::UnknownJob { .. }) + )); + } + + #[test] + fn times_out_claimed_jobs_and_rejects_late_completion() { + let store = + AutomationJobStore::with_durations(Duration::from_secs(30), Duration::from_secs(30)); + let now = Instant::now(); + let job = store.submit_at(inspect("Workspace"), now).unwrap(); + let claimant = PluginSessionId::new(); + store.claim_next_at(claimant, now).unwrap(); + store.cleanup_expired_at(now + INSPECT_EXECUTION_TIMEOUT); + assert_eq!( + store.get(job.id).unwrap().state, + AutomationJobState::TimedOut + ); + assert!(matches!( + store.complete_success(job.id, claimant, result("Workspace")), + Err(AutomationJobStoreError::DuplicateCompletion { + state: AutomationJobState::TimedOut, + .. + }) + )); + } + + #[test] + fn enforces_queue_and_request_limits() { + let store = AutomationJobStore::new(); + for _ in 0..MAX_PENDING_AUTOMATION_JOBS { + store.submit(inspect("Workspace")).unwrap(); + } + assert!(matches!( + store.submit(inspect("Workspace")), + Err(AutomationJobStoreError::PendingQueueFull { .. }) + )); + let mut invalid = inspect("Workspace"); + let AutomationRequest::Inspect(request) = &mut invalid; + request.depth = MAX_INSPECT_DEPTH + 1; + assert!(matches!( + AutomationJobStore::new().submit(invalid), + Err(AutomationJobStoreError::InvalidRequest(_)) + )); + } +} diff --git a/src/automation_status.rs b/src/automation_status.rs new file mode 100644 index 00000000..b8c63b41 --- /dev/null +++ b/src/automation_status.rs @@ -0,0 +1,305 @@ +use std::{ + fmt, + str::FromStr, + sync::Mutex, + time::{Duration, Instant, SystemTime, UNIX_EPOCH}, +}; + +use serde::{Deserialize, Serialize}; +use uuid::Uuid; + +use crate::SessionId; + +pub const AUTOMATION_HANDLER_VERSION: u32 = 2; +pub const AUTOMATION_SESSION_TIMEOUT: Duration = Duration::from_secs(10); + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)] +pub struct PluginSessionId(Uuid); + +impl PluginSessionId { + #[cfg(test)] + pub fn new() -> Self { + Self(Uuid::new_v4()) + } +} + +impl fmt::Display for PluginSessionId { + fn fmt(&self, writer: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(writer, "{}", self.0) + } +} + +impl FromStr for PluginSessionId { + type Err = uuid::Error; + + fn from_str(value: &str) -> Result { + Uuid::parse_str(value).map(Self) + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub enum StudioMode { + Edit, + Play, + Run, + Unknown, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AutomationSessionUpdate { + pub plugin_session_id: PluginSessionId, + pub server_session_id: SessionId, + pub studio_mode: StudioMode, + pub exec_handler_available: bool, + pub automation_handler_version: u32, + pub plugin_version: Option, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub enum AutomationRegistration { + Registered, + Refreshed, + Conflict, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AutomationPluginStatus { + pub plugin_session_id: PluginSessionId, + pub server_session_id: SessionId, + pub connected: bool, + pub last_seen_at_ms: u64, + pub studio_mode: StudioMode, + pub exec_handler_available: bool, + pub automation_handler_version: u32, + pub plugin_version: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AutomationStatusSnapshot { + pub plugin: Option, + pub duplicate_session_detected: bool, + pub exec_claimed_by_plugin_session_id: Option, +} + +#[derive(Debug)] +pub struct AutomationStatusStore { + inner: Mutex, + session_timeout: Duration, +} + +#[derive(Debug, Default)] +struct AutomationStatusInner { + active: Option, + duplicate_last_seen: Option, + exec_claimed_by_plugin_session_id: Option, +} + +#[derive(Debug, Clone)] +struct StoredAutomationPluginStatus { + update: AutomationSessionUpdate, + last_seen: Instant, + last_seen_at: SystemTime, +} + +impl AutomationStatusStore { + pub fn new() -> Self { + Self::with_timeout(AUTOMATION_SESSION_TIMEOUT) + } + + pub fn refresh(&self, update: AutomationSessionUpdate) -> AutomationRegistration { + self.refresh_at(update, Instant::now(), SystemTime::now()) + } + + pub fn snapshot(&self) -> AutomationStatusSnapshot { + self.snapshot_at(Instant::now()) + } + + pub fn note_exec_claim(&self, plugin_session_id: PluginSessionId) { + let mut inner = self.inner.lock().unwrap(); + inner.exec_claimed_by_plugin_session_id = Some(plugin_session_id); + } + + pub fn note_exec_completion(&self, plugin_session_id: PluginSessionId) { + let mut inner = self.inner.lock().unwrap(); + if inner.exec_claimed_by_plugin_session_id == Some(plugin_session_id) { + inner.exec_claimed_by_plugin_session_id = None; + } + } + + pub fn clear_exec_claim_if_idle(&self, has_claimed_job: bool) { + if !has_claimed_job { + self.inner.lock().unwrap().exec_claimed_by_plugin_session_id = None; + } + } + + fn with_timeout(session_timeout: Duration) -> Self { + Self { + inner: Mutex::new(AutomationStatusInner::default()), + session_timeout, + } + } + + fn refresh_at( + &self, + mut update: AutomationSessionUpdate, + now: Instant, + wall_clock: SystemTime, + ) -> AutomationRegistration { + let mut inner = self.inner.lock().unwrap(); + let active_is_current = inner.active.as_ref().is_some_and(|active| { + now.saturating_duration_since(active.last_seen) < self.session_timeout + }); + + if active_is_current { + let active = inner.active.as_mut().unwrap(); + if active.update.plugin_session_id != update.plugin_session_id { + inner.duplicate_last_seen = Some(now); + return AutomationRegistration::Conflict; + } + + if update.plugin_version.is_none() { + update.plugin_version = active.update.plugin_version.clone(); + } + active.update = update; + active.last_seen = now; + active.last_seen_at = wall_clock; + return AutomationRegistration::Refreshed; + } + + inner.active = Some(StoredAutomationPluginStatus { + update, + last_seen: now, + last_seen_at: wall_clock, + }); + inner.duplicate_last_seen = None; + AutomationRegistration::Registered + } + + fn snapshot_at(&self, now: Instant) -> AutomationStatusSnapshot { + let inner = self.inner.lock().unwrap(); + let plugin = inner.active.as_ref().map(|active| AutomationPluginStatus { + plugin_session_id: active.update.plugin_session_id, + server_session_id: active.update.server_session_id, + connected: now.saturating_duration_since(active.last_seen) < self.session_timeout, + last_seen_at_ms: active + .last_seen_at + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_millis() + .try_into() + .unwrap_or(u64::MAX), + studio_mode: active.update.studio_mode, + exec_handler_available: active.update.exec_handler_available, + automation_handler_version: active.update.automation_handler_version, + plugin_version: active.update.plugin_version.clone(), + }); + let duplicate_session_detected = inner.duplicate_last_seen.is_some_and(|last_seen| { + now.saturating_duration_since(last_seen) < self.session_timeout + }); + + AutomationStatusSnapshot { + plugin, + duplicate_session_detected, + exec_claimed_by_plugin_session_id: inner.exec_claimed_by_plugin_session_id, + } + } +} + +impl Default for AutomationStatusStore { + fn default() -> Self { + Self::new() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn update(plugin_session_id: PluginSessionId, mode: StudioMode) -> AutomationSessionUpdate { + AutomationSessionUpdate { + plugin_session_id, + server_session_id: SessionId::new(), + studio_mode: mode, + exec_handler_available: true, + automation_handler_version: AUTOMATION_HANDLER_VERSION, + plugin_version: Some("7.6.1".to_owned()), + } + } + + #[test] + fn reports_no_plugin_before_registration() { + let store = AutomationStatusStore::new(); + let status = store.snapshot(); + assert_eq!(status.plugin, None); + assert!(!status.duplicate_session_detected); + } + + #[test] + fn registers_and_refreshes_the_same_session() { + let store = AutomationStatusStore::new(); + let plugin_session_id = PluginSessionId::new(); + assert_eq!( + store.refresh(update(plugin_session_id, StudioMode::Edit)), + AutomationRegistration::Registered + ); + assert_eq!( + store.refresh(update(plugin_session_id, StudioMode::Play)), + AutomationRegistration::Refreshed + ); + assert_eq!( + store.snapshot().plugin.unwrap().studio_mode, + StudioMode::Play + ); + } + + #[test] + fn detects_a_second_active_session() { + let store = AutomationStatusStore::new(); + let first = PluginSessionId::new(); + let second = PluginSessionId::new(); + store.refresh(update(first, StudioMode::Edit)); + assert_eq!( + store.refresh(update(second, StudioMode::Edit)), + AutomationRegistration::Conflict + ); + let status = store.snapshot(); + assert_eq!(status.plugin.unwrap().plugin_session_id, first); + assert!(status.duplicate_session_detected); + } + + #[test] + fn a_stale_session_can_be_replaced() { + let timeout = Duration::from_secs(10); + let store = AutomationStatusStore::with_timeout(timeout); + let started = Instant::now(); + let first = PluginSessionId::new(); + let second = PluginSessionId::new(); + store.refresh_at(update(first, StudioMode::Edit), started, UNIX_EPOCH); + + let stale = store.snapshot_at(started + timeout); + assert!(!stale.plugin.unwrap().connected); + assert_eq!( + store.refresh_at( + update(second, StudioMode::Run), + started + timeout, + UNIX_EPOCH + Duration::from_secs(1), + ), + AutomationRegistration::Registered + ); + assert_eq!(store.snapshot().plugin.unwrap().plugin_session_id, second); + } + + #[test] + fn serializes_all_studio_modes_as_camel_case_strings() { + for (mode, expected) in [ + (StudioMode::Edit, "\"edit\""), + (StudioMode::Play, "\"play\""), + (StudioMode::Run, "\"run\""), + (StudioMode::Unknown, "\"unknown\""), + ] { + assert_eq!(serde_json::to_string(&mode).unwrap(), expected); + } + } +} diff --git a/src/cli/automation.rs b/src/cli/automation.rs new file mode 100644 index 00000000..8a0ec949 --- /dev/null +++ b/src/cli/automation.rs @@ -0,0 +1,373 @@ +use std::{ + io::Read, + net::IpAddr, + thread, + time::{Duration, Instant}, +}; + +use anyhow::{anyhow, bail, Context}; +use reqwest::{ + blocking::{Client, RequestBuilder, Response}, + header::{ACCEPT, CONTENT_TYPE}, + StatusCode, +}; +use serde::de::DeserializeOwned; +use serde::Serialize; + +use crate::automation::MAX_AUTOMATION_RESULT_BODY_BYTES; +use crate::web::{ + deserialize_msgpack, + interface::{ErrorResponse, ServerInfoResponse, PROTOCOL_VERSION}, + serialize_msgpack, +}; + +pub(super) const DEFAULT_ADDRESS: &str = "127.0.0.1"; +pub(super) const DEFAULT_PORT: u16 = 34872; +pub(super) const CONNECT_TIMEOUT: Duration = Duration::from_secs(5); +pub(super) const REQUEST_TIMEOUT: Duration = Duration::from_secs(10); +pub(super) const MSGPACK_CONTENT_TYPE: &str = "application/msgpack"; +const MAX_RESPONSE_SIZE_BYTES: usize = MAX_AUTOMATION_RESULT_BODY_BYTES + 64 * 1024; + +pub(super) fn build_client(operation: &str) -> anyhow::Result { + Client::builder() + .connect_timeout(CONNECT_TIMEOUT) + .timeout(REQUEST_TIMEOUT) + .build() + .with_context(|| format!("Could not create the HTTP client for {operation}")) +} + +pub(super) fn server_url(address: IpAddr, port: u16) -> String { + match address { + IpAddr::V4(address) => format!("http://{address}:{port}"), + IpAddr::V6(address) => format!("http://[{address}]:{port}"), + } +} + +pub(super) fn verify_rojo_server( + client: &Client, + server_url: &str, + client_description: &str, +) -> anyhow::Result<()> { + let response = send_request( + client + .get(format!("{server_url}/api/rojo")) + .header(ACCEPT, MSGPACK_CONTENT_TYPE), + server_url, + "checking the server", + )?; + let response = buffer_response(response, "checking the server")?; + + if response.status != StatusCode::OK { + let details = error_details(&response).unwrap_or_default(); + bail!( + "The HTTP service at {server_url} did not return Rojo server information (HTTP {}){details}.", + response.status + ); + } + + if !response.is_msgpack { + let content_type = response.content_type.as_deref().unwrap_or("missing"); + bail!( + "The HTTP service at {server_url} returned content type '{content_type}' from /api/rojo, not {MSGPACK_CONTENT_TYPE}; it does not appear to be a compatible Rojo server." + ); + } + + let info: ServerInfoResponse = deserialize_msgpack(&response.body).map_err(|error| { + anyhow!( + "The Rojo server at {server_url} returned malformed MessagePack from /api/rojo: {error}" + ) + })?; + + if info.protocol_version != PROTOCOL_VERSION { + bail!( + "The Rojo server at {server_url} uses protocol version {}, but this {client_description} requires version {}.", + info.protocol_version, + PROTOCOL_VERSION + ); + } + + Ok(()) +} + +pub(super) fn send_request( + request: RequestBuilder, + server_url: &str, + operation: &str, +) -> anyhow::Result { + request.send().map_err(|error| { + if error.is_connect() { + anyhow!("Could not connect to the Rojo server at {server_url} while {operation}: {error}") + } else if error.is_timeout() { + anyhow!("The request to the Rojo server at {server_url} timed out while {operation}: {error}") + } else { + anyhow!("The request to the Rojo server at {server_url} failed while {operation}: {error}") + } + }) +} + +pub(super) struct BufferedResponse { + pub status: StatusCode, + pub content_type: Option, + pub is_msgpack: bool, + pub body: Vec, +} + +pub(super) fn buffer_response( + response: Response, + operation: &str, +) -> anyhow::Result { + let status = response.status(); + let content_type = response + .headers() + .get(CONTENT_TYPE) + .and_then(|value| value.to_str().ok()) + .map(str::to_owned); + let is_msgpack = content_type + .as_deref() + .and_then(|value| value.split(';').next()) + .is_some_and(|value| value.trim().eq_ignore_ascii_case(MSGPACK_CONTENT_TYPE)); + let mut body = Vec::new(); + let mut limited = response.take((MAX_RESPONSE_SIZE_BYTES + 1) as u64); + limited + .read_to_end(&mut body) + .with_context(|| format!("Could not read the server response while {operation}"))?; + + if body.len() > MAX_RESPONSE_SIZE_BYTES { + bail!( + "The Rojo server response exceeded the {}-byte client limit while {operation}.", + MAX_RESPONSE_SIZE_BYTES + ); + } + + Ok(BufferedResponse { + status, + content_type, + is_msgpack, + body, + }) +} + +pub(super) fn decode_response( + response: Response, + expected_status: StatusCode, + server_url: &str, + operation: &str, + status_summary: fn(StatusCode) -> &'static str, +) -> anyhow::Result { + let response = buffer_response(response, operation)?; + decode_buffered_response( + &response, + expected_status, + server_url, + operation, + status_summary, + ) +} + +pub(super) fn get_msgpack( + client: &Client, + server_url: &str, + path: &str, + expected_status: StatusCode, + operation: &str, +) -> anyhow::Result { + let response = send_request( + client + .get(format!("{server_url}{path}")) + .header(ACCEPT, MSGPACK_CONTENT_TYPE), + server_url, + operation, + )?; + decode_response( + response, + expected_status, + server_url, + operation, + automation_http_summary, + ) +} + +pub(super) fn post_msgpack( + client: &Client, + server_url: &str, + path: &str, + request: &Request, + expected_status: StatusCode, + operation: &str, +) -> anyhow::Result { + let body = serialize_msgpack(request) + .with_context(|| format!("Could not encode the request while {operation}"))?; + let response = send_request( + client + .post(format!("{server_url}{path}")) + .header(ACCEPT, MSGPACK_CONTENT_TYPE) + .header(CONTENT_TYPE, MSGPACK_CONTENT_TYPE) + .body(body), + server_url, + operation, + )?; + decode_response( + response, + expected_status, + server_url, + operation, + automation_http_summary, + ) +} + +pub(super) fn decode_buffered_response( + response: &BufferedResponse, + expected_status: StatusCode, + server_url: &str, + operation: &str, + status_summary: fn(StatusCode) -> &'static str, +) -> anyhow::Result { + if response.status != expected_status { + return Err(api_http_error(response, operation, status_summary)); + } + + if !response.is_msgpack { + let content_type = response.content_type.as_deref().unwrap_or("missing"); + bail!( + "The Rojo server at {server_url} returned content type '{content_type}' while {operation}; expected {MSGPACK_CONTENT_TYPE}." + ); + } + + deserialize_msgpack(&response.body).map_err(|error| { + anyhow!( + "The Rojo server at {server_url} returned malformed MessagePack while {operation}: {error}" + ) + }) +} + +pub(super) struct PollOptions<'a> { + pub server_url: &'a str, + pub path: &'a str, + pub timeout: Duration, + pub interval: Duration, + pub operation: &'a str, +} + +pub(super) fn poll_status( + client: &Client, + options: PollOptions<'_>, + decode: impl Fn(Response) -> anyhow::Result, + is_terminal: impl Fn(&T) -> anyhow::Result, + timeout_error: impl Fn() -> anyhow::Error, +) -> anyhow::Result { + let started_at = Instant::now(); + + loop { + let elapsed = started_at.elapsed(); + if elapsed >= options.timeout { + return Err(timeout_error()); + } + let remaining = options.timeout - elapsed; + let response = send_request( + client + .get(format!("{}{}", options.server_url, options.path)) + .timeout(REQUEST_TIMEOUT.min(remaining)) + .header(ACCEPT, MSGPACK_CONTENT_TYPE), + options.server_url, + options.operation, + ) + .map_err(|error| { + if started_at.elapsed() >= options.timeout { + timeout_error() + } else { + error + } + })?; + let status = decode(response)?; + if is_terminal(&status)? { + return Ok(status); + } + + let elapsed = started_at.elapsed(); + if elapsed >= options.timeout { + return Err(timeout_error()); + } + thread::sleep(options.interval.min(options.timeout - elapsed)); + } +} + +fn api_http_error( + response: &BufferedResponse, + operation: &str, + status_summary: fn(StatusCode) -> &'static str, +) -> anyhow::Error { + let details = error_details(response).unwrap_or_default(); + anyhow!( + "{} while {operation} (HTTP {}){details}.", + status_summary(response.status), + response.status + ) +} + +fn error_details(response: &BufferedResponse) -> Option { + if !response.is_msgpack { + return None; + } + + let error: ErrorResponse = deserialize_msgpack(&response.body).ok()?; + if error.details().is_empty() { + Some(format!(" [{:?}]", error.kind())) + } else { + Some(format!(" [{:?}: {}]", error.kind(), error.details())) + } +} + +pub(super) fn automation_http_summary(status: StatusCode) -> &'static str { + match status { + StatusCode::BAD_REQUEST => "Prism automation request was malformed", + StatusCode::FORBIDDEN => "Prism automation API is not available from this peer", + StatusCode::NOT_FOUND => "Prism automation job disappeared or expired", + StatusCode::CONFLICT => "Prism automation request has a state or session conflict", + StatusCode::PAYLOAD_TOO_LARGE => "Prism automation request or result is too large", + StatusCode::TOO_MANY_REQUESTS => "Prism automation queue is full", + StatusCode::INTERNAL_SERVER_ERROR => "Prism automation server reported an internal error", + _ => "Prism automation server returned an unexpected HTTP status", + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::web::serialize_msgpack; + + fn summary(_status: StatusCode) -> &'static str { + "automation request failed" + } + + #[test] + fn constructs_ipv4_and_ipv6_server_urls() { + assert_eq!( + server_url("127.0.0.1".parse().unwrap(), 34872), + "http://127.0.0.1:34872" + ); + assert_eq!( + server_url("::1".parse().unwrap(), 34872), + "http://[::1]:34872" + ); + } + + #[test] + fn decodes_common_messagepack_error_envelopes() { + let response = BufferedResponse { + status: StatusCode::CONFLICT, + content_type: Some(MSGPACK_CONTENT_TYPE.to_owned()), + is_msgpack: true, + body: serialize_msgpack(ErrorResponse::conflict("duplicate session")).unwrap(), + }; + let error = decode_buffered_response::( + &response, + StatusCode::OK, + "http://127.0.0.1:34872", + "checking automation status", + summary, + ) + .unwrap_err(); + + assert!(error.to_string().contains("Conflict: duplicate session")); + } +} diff --git a/src/cli/exec.rs b/src/cli/exec.rs index 3b14419e..ae660076 100644 --- a/src/cli/exec.rs +++ b/src/cli/exec.rs @@ -1,43 +1,49 @@ use std::{ - io::{self, Read, Write}, + io::{self, Write}, net::IpAddr, path::{Path, PathBuf}, - thread, - time::{Duration, Instant}, + time::Duration, }; use anyhow::{anyhow, bail, Context}; use clap::Parser; use reqwest::{ - blocking::{Client, RequestBuilder, Response}, + blocking::Client, header::{ACCEPT, CONTENT_TYPE}, StatusCode, }; -use serde::de::DeserializeOwned; use uuid::Uuid; +#[cfg(test)] +use serde::de::DeserializeOwned; + use crate::{ exec::MAX_SOURCE_SIZE_BYTES, web::{ - deserialize_msgpack, interface::{ - ErrorResponse, ExecJobResponse, ExecJobState, ExecJobSubmissionRequest, ExecLog, - ExecLogLevel, ExecValue, ServerInfoResponse, PROTOCOL_VERSION, + ExecJobResponse, ExecJobState, ExecJobSubmissionRequest, ExecLog, ExecLogLevel, + ExecValue, }, serialize_msgpack, }, }; -use super::resolve_path; +#[cfg(test)] +use crate::web::interface::ErrorResponse; + +use super::{ + automation::{ + build_client, decode_response, poll_status, send_request, server_url, verify_rojo_server, + PollOptions, DEFAULT_ADDRESS, DEFAULT_PORT, MSGPACK_CONTENT_TYPE, + }, + resolve_path, +}; + +#[cfg(test)] +use super::automation::{decode_buffered_response, BufferedResponse}; -const DEFAULT_ADDRESS: &str = "127.0.0.1"; -const DEFAULT_PORT: u16 = 34872; const POLL_INTERVAL: Duration = Duration::from_millis(250); const LOCAL_TIMEOUT: Duration = Duration::from_secs(70); -const CONNECT_TIMEOUT: Duration = Duration::from_secs(5); -const REQUEST_TIMEOUT: Duration = Duration::from_secs(10); -const MAX_RESPONSE_SIZE_BYTES: usize = 2 * MAX_SOURCE_SIZE_BYTES; -const MSGPACK_CONTENT_TYPE: &str = "application/msgpack"; /// Runs a trusted Luau file through a connected Prism Studio plugin. #[derive(Debug, Parser)] @@ -58,13 +64,9 @@ impl ExecCommand { pub fn run(self) -> anyhow::Result<()> { let script = read_exec_script(&self.file)?; let server_url = server_url(self.address, self.port); - let client = Client::builder() - .connect_timeout(CONNECT_TIMEOUT) - .timeout(REQUEST_TIMEOUT) - .build() - .context("Could not create the HTTP client for Prism exec")?; + let client = build_client("Prism exec")?; - verify_rojo_server(&client, &server_url)?; + verify_rojo_server(&client, &server_url, "Prism exec client")?; let submitted = submit_job(&client, &server_url, &script)?; let job_id = parse_response_job_id(&submitted.job_id, "submission")?; log::debug!("Submitted exec job {}", job_id); @@ -148,55 +150,6 @@ fn script_name_from_path(path: &Path) -> anyhow::Result { Ok(file_name.to_owned()) } -fn server_url(address: IpAddr, port: u16) -> String { - match address { - IpAddr::V4(address) => format!("http://{address}:{port}"), - IpAddr::V6(address) => format!("http://[{address}]:{port}"), - } -} - -fn verify_rojo_server(client: &Client, server_url: &str) -> anyhow::Result<()> { - let response = send_request( - client - .get(format!("{server_url}/api/rojo")) - .header(ACCEPT, MSGPACK_CONTENT_TYPE), - server_url, - "checking the server", - )?; - let response = buffer_response(response, "checking the server")?; - - if response.status != StatusCode::OK { - let details = error_details(&response).unwrap_or_default(); - bail!( - "The HTTP service at {server_url} did not return Rojo server information (HTTP {}){details}.", - response.status - ); - } - - if !response.is_msgpack { - let content_type = response.content_type.as_deref().unwrap_or("missing"); - bail!( - "The HTTP service at {server_url} returned content type '{content_type}' from /api/rojo, not {MSGPACK_CONTENT_TYPE}; it does not appear to be a compatible Rojo server." - ); - } - - let info: ServerInfoResponse = deserialize_msgpack(&response.body).map_err(|error| { - anyhow!( - "The Rojo server at {server_url} returned malformed MessagePack from /api/rojo: {error}" - ) - })?; - - if info.protocol_version != PROTOCOL_VERSION { - bail!( - "The Rojo server at {server_url} uses protocol version {}, but this Prism exec client requires version {}.", - info.protocol_version, - PROTOCOL_VERSION - ); - } - - Ok(()) -} - fn submit_job( client: &Client, server_url: &str, @@ -217,11 +170,12 @@ fn submit_job( server_url, "submitting the exec job", )?; - let job: ExecJobResponse = decode_exec_response( + let job: ExecJobResponse = decode_response( response, StatusCode::CREATED, server_url, "submitting the exec job", + exec_http_summary, )?; if job.state != ExecJobState::Pending { @@ -240,54 +194,39 @@ fn poll_job( job_id: Uuid, timeout: Duration, ) -> anyhow::Result { - let started_at = Instant::now(); - - loop { - let elapsed = started_at.elapsed(); - if elapsed >= timeout { - return Err(local_timeout_error(job_id, timeout)); - } - let remaining = timeout - elapsed; - - let response = send_request( - client - .get(format!("{server_url}/api/exec/jobs/{job_id}")) - .timeout(REQUEST_TIMEOUT.min(remaining)) - .header(ACCEPT, MSGPACK_CONTENT_TYPE), + let path = format!("/api/exec/jobs/{job_id}"); + poll_status( + client, + PollOptions { server_url, - "polling the exec job", - ) - .map_err(|error| { - if started_at.elapsed() >= timeout { - local_timeout_error(job_id, timeout) - } else { - error + path: &path, + timeout, + interval: POLL_INTERVAL, + operation: "polling the exec job", + }, + |response| { + decode_response( + response, + StatusCode::OK, + server_url, + "polling the exec job", + exec_http_summary, + ) + }, + |job: &ExecJobResponse| { + let response_job_id = parse_response_job_id(&job.job_id, "status")?; + if response_job_id != job_id { + bail!( + "The Rojo server at {server_url} returned status for exec job {response_job_id} while job {job_id} was requested." + ); } - })?; - let job: ExecJobResponse = - decode_exec_response(response, StatusCode::OK, server_url, "polling the exec job")?; - let response_job_id = parse_response_job_id(&job.job_id, "status")?; - - if response_job_id != job_id { - bail!( - "The Rojo server at {server_url} returned status for exec job {response_job_id} while job {job_id} was requested." - ); - } - - if matches!( - job.state, - ExecJobState::Succeeded | ExecJobState::Failed | ExecJobState::TimedOut - ) { - return Ok(job); - } - - let elapsed = started_at.elapsed(); - if elapsed >= timeout { - return Err(local_timeout_error(job_id, timeout)); - } - - thread::sleep(POLL_INTERVAL.min(timeout - elapsed)); - } + Ok(matches!( + job.state, + ExecJobState::Succeeded | ExecJobState::Failed | ExecJobState::TimedOut + )) + }, + || local_timeout_error(job_id, timeout), + ) } fn parse_response_job_id(job_id: &str, response_kind: &str) -> anyhow::Result { @@ -303,97 +242,24 @@ fn local_timeout_error(job_id: Uuid, timeout: Duration) -> anyhow::Error { ) } -fn send_request( - request: RequestBuilder, - server_url: &str, - operation: &str, -) -> anyhow::Result { - request.send().map_err(|error| { - if error.is_connect() { - anyhow!("Could not connect to the Rojo server at {server_url} while {operation}: {error}") - } else if error.is_timeout() { - anyhow!("The request to the Rojo server at {server_url} timed out while {operation}: {error}") - } else { - anyhow!("The request to the Rojo server at {server_url} failed while {operation}: {error}") - } - }) -} - -struct BufferedResponse { - status: StatusCode, - content_type: Option, - is_msgpack: bool, - body: Vec, -} - -fn buffer_response(response: Response, operation: &str) -> anyhow::Result { - let status = response.status(); - let content_type = response - .headers() - .get(CONTENT_TYPE) - .and_then(|value| value.to_str().ok()) - .map(str::to_owned); - let is_msgpack = content_type - .as_deref() - .and_then(|value| value.split(';').next()) - .is_some_and(|value| value.trim().eq_ignore_ascii_case(MSGPACK_CONTENT_TYPE)); - let mut body = Vec::new(); - let mut limited = response.take((MAX_RESPONSE_SIZE_BYTES + 1) as u64); - limited - .read_to_end(&mut body) - .with_context(|| format!("Could not read the server response while {operation}"))?; - - if body.len() > MAX_RESPONSE_SIZE_BYTES { - bail!( - "The Rojo server response exceeded the {}-byte client limit while {operation}.", - MAX_RESPONSE_SIZE_BYTES - ); - } - - Ok(BufferedResponse { - status, - content_type, - is_msgpack, - body, - }) -} - -fn decode_exec_response( - response: Response, - expected_status: StatusCode, - server_url: &str, - operation: &str, -) -> anyhow::Result { - let response = buffer_response(response, operation)?; - decode_buffered_exec_response(&response, expected_status, server_url, operation) -} - +#[cfg(test)] fn decode_buffered_exec_response( response: &BufferedResponse, expected_status: StatusCode, server_url: &str, operation: &str, ) -> anyhow::Result { - if response.status != expected_status { - return Err(exec_http_error(response, operation)); - } - - if !response.is_msgpack { - let content_type = response.content_type.as_deref().unwrap_or("missing"); - bail!( - "The Rojo server at {server_url} returned content type '{content_type}' while {operation}; expected {MSGPACK_CONTENT_TYPE}." - ); - } - - deserialize_msgpack(&response.body).map_err(|error| { - anyhow!( - "The Rojo server at {server_url} returned malformed MessagePack while {operation}: {error}" - ) - }) + decode_buffered_response( + response, + expected_status, + server_url, + operation, + exec_http_summary, + ) } -fn exec_http_error(response: &BufferedResponse, operation: &str) -> anyhow::Error { - let summary = match response.status { +fn exec_http_summary(status: StatusCode) -> &'static str { + match status { StatusCode::BAD_REQUEST => "Rojo server rejected the request as malformed", StatusCode::FORBIDDEN => "Prism exec API is not available from this peer", StatusCode::NOT_FOUND => "Prism exec job disappeared or expired", @@ -402,25 +268,6 @@ fn exec_http_error(response: &BufferedResponse, operation: &str) -> anyhow::Erro StatusCode::TOO_MANY_REQUESTS => "Prism exec queue is full", StatusCode::INTERNAL_SERVER_ERROR => "Rojo server reported an internal error", _ => "Rojo server returned an unexpected HTTP status", - }; - let details = error_details(response).unwrap_or_default(); - - anyhow!( - "{summary} while {operation} (HTTP {}){details}.", - response.status - ) -} - -fn error_details(response: &BufferedResponse) -> Option { - if !response.is_msgpack { - return None; - } - - let error: ErrorResponse = deserialize_msgpack(&response.body).ok()?; - if error.details().is_empty() { - Some(format!(" [{:?}]", error.kind())) - } else { - Some(format!(" [{:?}: {}]", error.kind(), error.details())) } } diff --git a/src/cli/inspect.rs b/src/cli/inspect.rs new file mode 100644 index 00000000..e6a9a41b --- /dev/null +++ b/src/cli/inspect.rs @@ -0,0 +1,572 @@ +use std::{ + io::{self, Write}, + net::IpAddr, + time::Duration, +}; + +use anyhow::{anyhow, bail, Context}; +use clap::Parser; +use reqwest::{blocking::Client, StatusCode}; +use uuid::Uuid; + +use crate::{ + automation::{ + AutomationJobState, AutomationRequest, AutomationResult, AutomationValue, InspectNode, + InspectRequest, InspectResult, InspectTarget, MAX_INSPECT_CHILDREN, MAX_INSPECT_DEPTH, + }, + web::interface::{AutomationJobResponse, AutomationStatusResponse, StudioMode}, +}; + +use super::automation::{ + automation_http_summary, build_client, decode_response, get_msgpack, poll_status, post_msgpack, + server_url, verify_rojo_server, PollOptions, DEFAULT_ADDRESS, DEFAULT_PORT, +}; + +const DEFAULT_DEPTH: u8 = 1; +const DEFAULT_MAX_CHILDREN: u32 = 100; +const DEFAULT_MAX_INSTANCES: u32 = 2_000; +const POLL_INTERVAL: Duration = Duration::from_millis(250); +const LOCAL_TIMEOUT: Duration = Duration::from_secs(50); + +#[derive(Debug, Parser)] +pub struct InspectCommand { + /// DataModel path to inspect. + pub target: String, + + /// Number of child levels to include. + #[clap(long, default_value_t = DEFAULT_DEPTH, parse(try_from_str = parse_depth))] + pub depth: u8, + + /// Maximum number of children included at each node. + #[clap( + long, + default_value_t = DEFAULT_MAX_CHILDREN, + parse(try_from_str = parse_max_children) + )] + pub max_children: u32, + + /// Include conservative, class-aware properties. + #[clap(long)] + pub properties: bool, + + /// Include attributes. Enabled by default. + #[clap(long, action, default_value_t = true)] + pub attributes: bool, + + /// Include CollectionService tags. Enabled by default. + #[clap(long, action, default_value_t = true)] + pub tags: bool, + + /// Emit deterministic JSON instead of human-readable output. + #[clap(long)] + pub json: bool, + + /// IP address of the running Rojo server. + #[clap(long, default_value = DEFAULT_ADDRESS)] + pub address: IpAddr, + + /// Port of the running Rojo server. + #[clap(long, default_value_t = DEFAULT_PORT)] + pub port: u16, +} + +impl InspectCommand { + pub fn run(self) -> anyhow::Result<()> { + let segments = parse_target(&self.target)?; + let request = AutomationRequest::Inspect(InspectRequest { + target: InspectTarget::Path { segments }, + depth: self.depth, + max_children: self.max_children, + max_instances: DEFAULT_MAX_INSTANCES, + include_properties: self.properties, + include_attributes: self.attributes, + include_tags: self.tags, + }); + let server_url = server_url(self.address, self.port); + let client = build_client("Prism inspect")?; + verify_rojo_server(&client, &server_url, "Prism inspect client")?; + let status: AutomationStatusResponse = get_msgpack( + &client, + &server_url, + "/api/automation/status", + StatusCode::OK, + "checking automation availability", + )?; + validate_automation_status(&status)?; + let submitted: AutomationJobResponse = post_msgpack( + &client, + &server_url, + "/api/automation/jobs", + &request, + StatusCode::CREATED, + "submitting the inspect job", + )?; + if submitted.state != AutomationJobState::Pending { + bail!("Prism returned an unexpected state for a newly submitted inspect job"); + } + let job_id = Uuid::parse_str(&submitted.job_id) + .context("Prism returned a malformed inspect job ID")?; + let completed = poll_automation_job(&client, &server_url, job_id)?; + let result = match completed.state { + AutomationJobState::Succeeded => match completed.result { + Some(AutomationResult::Inspect(result)) => result, + _ => bail!("Prism returned a malformed or mismatched inspect result"), + }, + AutomationJobState::Failed | AutomationJobState::TimedOut => { + bail!( + "Prism inspect failed: {}", + completed + .error + .unwrap_or_else(|| "automation job failed without details".to_owned()) + ) + } + state => bail!("Prism inspect polling ended in non-terminal state {state:?}"), + }; + + let stdout = io::stdout(); + let mut stdout = stdout.lock(); + if self.json { + serde_json::to_writer_pretty(&mut stdout, &result) + .context("Could not render inspect JSON")?; + writeln!(stdout).context("Could not write inspect JSON")?; + } else { + render_human(&result, &mut stdout)?; + } + Ok(()) + } +} + +fn parse_depth(value: &str) -> Result { + let value = value + .parse::() + .map_err(|_| "depth must be a non-negative integer".to_owned())?; + if value > MAX_INSPECT_DEPTH { + return Err(format!("depth must not exceed {MAX_INSPECT_DEPTH}")); + } + Ok(value) +} + +fn parse_max_children(value: &str) -> Result { + let value = value + .parse::() + .map_err(|_| "max-children must be a positive integer".to_owned())?; + if value == 0 || value > MAX_INSPECT_CHILDREN { + return Err(format!( + "max-children must be between 1 and {MAX_INSPECT_CHILDREN}" + )); + } + Ok(value) +} + +fn parse_target(input: &str) -> anyhow::Result> { + if input.is_empty() { + bail!("Inspect target must not be empty"); + } + let bytes = input.as_bytes(); + let mut index = 0; + let mut segments = Vec::new(); + while index < bytes.len() { + let segment = if bytes[index] == b'"' { + let start = index; + index += 1; + let mut escaped = false; + while index < bytes.len() { + let byte = bytes[index]; + if escaped { + escaped = false; + } else if byte == b'\\' { + escaped = true; + } else if byte == b'"' { + index += 1; + break; + } + index += 1; + } + if index > bytes.len() || bytes.get(index.saturating_sub(1)) != Some(&b'"') { + bail!("Inspect target contains an unterminated quoted segment"); + } + let quoted = &input[start..index]; + let decoded: String = serde_json::from_str(quoted) + .with_context(|| format!("Invalid quoted path segment {quoted}"))?; + if decoded.is_empty() { + bail!("Inspect target contains an empty path segment"); + } + decoded + } else { + let start = index; + let first = input[index..].chars().next().unwrap(); + if !(first == '_' || first.is_ascii_alphabetic()) { + bail!("Unquoted path segments must be identifiers"); + } + index += first.len_utf8(); + while index < bytes.len() { + let character = input[index..].chars().next().unwrap(); + if character == '_' || character.is_ascii_alphanumeric() { + index += character.len_utf8(); + } else { + break; + } + } + input[start..index].to_owned() + }; + segments.push(segment); + if index == bytes.len() { + break; + } + if bytes[index] != b'.' { + bail!("Inspect target contains an unexpected character at byte {index}"); + } + index += 1; + if index == bytes.len() { + bail!("Inspect target must not end with a dot"); + } + } + + normalize_root(&mut segments[0])?; + Ok(segments) +} + +fn normalize_root(root: &mut String) -> anyhow::Result<()> { + *root = match root.as_str() { + "game" | "DataModel" => "game".to_owned(), + "workspace" | "Workspace" => "Workspace".to_owned(), + "Lighting" + | "ReplicatedStorage" + | "ServerStorage" + | "ServerScriptService" + | "StarterGui" + | "StarterPlayer" + | "StarterPack" + | "SoundService" + | "Teams" + | "TextChatService" => root.clone(), + _ => bail!("Unsupported inspect root '{root}'"), + }; + Ok(()) +} + +fn validate_automation_status(status: &AutomationStatusResponse) -> anyhow::Result<()> { + if status.duplicate_session_detected { + bail!("Prism detected multiple active Studio plugin sessions; close or disconnect the duplicate session"); + } + let plugin = status + .plugin + .as_ref() + .context("No Prism Studio plugin is connected for automation")?; + if !plugin.connected { + bail!("The Prism Studio automation session is stale or disconnected"); + } + if plugin.automation_handler_version != status.automation_handler_version { + bail!( + "Prism automation handler version mismatch: server {}, plugin {}", + status.automation_handler_version, + plugin.automation_handler_version + ); + } + if !status.typed_automation_available { + bail!("The connected Prism plugin does not provide typed automation handlers"); + } + if plugin.studio_mode != StudioMode::Edit { + bail!("Prism inspect is only available while Studio is in edit mode"); + } + Ok(()) +} + +fn poll_automation_job( + client: &Client, + server_url: &str, + job_id: Uuid, +) -> anyhow::Result { + let path = format!("/api/automation/jobs/{job_id}"); + poll_status( + client, + PollOptions { + server_url, + path: &path, + timeout: LOCAL_TIMEOUT, + interval: POLL_INTERVAL, + operation: "polling the inspect job", + }, + |response| { + decode_response( + response, + StatusCode::OK, + server_url, + "polling the inspect job", + automation_http_summary, + ) + }, + |job: &AutomationJobResponse| { + let response_id = Uuid::parse_str(&job.job_id) + .context("Prism returned a malformed inspect job ID")?; + if response_id != job_id { + bail!("Prism returned status for the wrong automation job"); + } + Ok(job.state.is_terminal()) + }, + || { + anyhow!( + "Timed out after {} seconds waiting for inspect job {job_id}; the server may retain it briefly", + LOCAL_TIMEOUT.as_secs() + ) + }, + ) +} + +fn render_human(result: &InspectResult, output: &mut impl Write) -> anyhow::Result<()> { + render_node(&result.root, 0, output)?; + if result.truncated { + writeln!( + output, + "[truncated: {} after {} instances]", + result.truncation_reason.as_deref().unwrap_or("limit"), + result.visited_instances + )?; + } + Ok(()) +} + +fn render_node(node: &InspectNode, depth: usize, output: &mut impl Write) -> anyhow::Result<()> { + let indent = " ".repeat(depth); + let label = if depth == 0 { &node.path } else { &node.name }; + writeln!(output, "{indent}{label} [{}]", node.class_name)?; + if !node.properties.is_empty() { + writeln!(output, "{indent} properties:")?; + for (key, value) in &node.properties { + writeln!(output, "{indent} {key} = {}", render_value(value))?; + } + } + if !node.attributes.is_empty() { + writeln!(output, "{indent} attributes:")?; + for (key, value) in &node.attributes { + writeln!(output, "{indent} {key} = {}", render_value(value))?; + } + } + if !node.tags.is_empty() { + writeln!(output, "{indent} tags:")?; + for tag in &node.tags { + writeln!(output, "{indent} - {tag}")?; + } + } + if !node.children.is_empty() { + writeln!(output, "{indent} children:")?; + for child in &node.children { + render_node(child, depth + 2, output)?; + } + } + if node.truncated { + writeln!(output, "{indent} [children truncated]")?; + } + Ok(()) +} + +fn render_value(value: &AutomationValue) -> String { + match value { + AutomationValue::Nil => "nil".to_owned(), + AutomationValue::Boolean { value } => value.to_string(), + AutomationValue::Number { value } => value.to_string(), + AutomationValue::String { value } => serde_json::to_string(value).unwrap(), + AutomationValue::Vector2 { x, y } => format!("Vector2({x}, {y})"), + AutomationValue::Vector3 { x, y, z } => format!("Vector3({x}, {y}, {z})"), + AutomationValue::Color3 { r, g, b } => format!("Color3({r}, {g}, {b})"), + AutomationValue::CFrame { components } => format!("CFrame({})", join_numbers(components)), + AutomationValue::UDim { scale, offset } => format!("UDim({scale}, {offset})"), + AutomationValue::UDim2 { x, y } => format!( + "UDim2({}, {}, {}, {})", + x.scale, x.offset, y.scale, y.offset + ), + AutomationValue::Rect { min, max } => { + format!("Rect({}, {}, {}, {})", min.x, min.y, max.x, max.y) + } + AutomationValue::EnumItem { enum_type, name } => format!("Enum.{enum_type}.{name}"), + AutomationValue::BrickColor { name, .. } => format!("BrickColor({name:?})"), + AutomationValue::NumberRange { min, max } => format!("NumberRange({min}, {max})"), + AutomationValue::InstanceReference { value } => { + format!("{} [{}]", value.path, value.class_name) + } + AutomationValue::Diagnostic { error } => format!(""), + AutomationValue::Array { .. } | AutomationValue::Map { .. } => { + serde_json::to_string(value).unwrap_or_else(|_| "".to_owned()) + } + } +} + +fn join_numbers(values: &[f64]) -> String { + values + .iter() + .map(ToString::to_string) + .collect::>() + .join(", ") +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::cli::Options; + use std::collections::BTreeMap; + + #[test] + fn parses_command_defaults_and_custom_values() { + let options = Options::try_parse_from(["rojo", "inspect", "Workspace"]).unwrap(); + let crate::cli::Subcommand::Inspect(command) = options.subcommand else { + panic!("expected inspect command") + }; + assert_eq!(command.depth, 1); + assert_eq!(command.max_children, 100); + assert_eq!(command.address, "127.0.0.1".parse::().unwrap()); + assert_eq!(command.port, 34872); + assert!(command.attributes); + assert!(command.tags); + + let options = Options::try_parse_from([ + "rojo", + "inspect", + "Workspace", + "--depth", + "3", + "--max-children", + "25", + "--address", + "127.0.0.2", + "--port", + "4000", + "--properties", + "--json", + ]) + .unwrap(); + let crate::cli::Subcommand::Inspect(command) = options.subcommand else { + panic!("expected inspect command") + }; + assert_eq!(command.depth, 3); + assert_eq!(command.max_children, 25); + assert!(command.properties); + assert!(command.json); + } + + #[test] + fn validates_depth_and_child_limits() { + assert!(Options::try_parse_from(["rojo", "inspect", "Workspace", "--depth", "9"]).is_err()); + assert!( + Options::try_parse_from(["rojo", "inspect", "Workspace", "--max-children", "0"]) + .is_err() + ); + } + + #[test] + fn parses_target_grammar_and_escapes() { + assert_eq!(parse_target("workspace").unwrap(), vec!["Workspace"]); + assert_eq!( + parse_target("Workspace.Map").unwrap(), + vec!["Workspace", "Map"] + ); + assert_eq!( + parse_target("Workspace.\"Name.With.Dots\"").unwrap(), + vec!["Workspace", "Name.With.Dots"] + ); + assert_eq!( + parse_target("Workspace.\"Quote \\\" Name\"").unwrap(), + vec!["Workspace", "Quote \" Name"] + ); + assert_eq!( + parse_target("Workspace.\"Backslash \\\\ Name\"").unwrap(), + vec!["Workspace", "Backslash \\ Name"] + ); + } + + #[test] + fn rejects_malformed_or_expression_like_targets() { + for target in [ + "Workspace.", + "Workspace..Map", + "Workspace[\"Map\"]", + "Workspace/Map", + "Workspace.\"bad\\q\"", + "NotAService.Map", + "inspect.lua", + ] { + assert!(parse_target(target).is_err(), "accepted {target}"); + } + } + + #[test] + fn rejects_missing_target() { + assert!(Options::try_parse_from(["rojo", "inspect"]).is_err()); + } + + fn sample_result() -> InspectResult { + InspectResult { + root: InspectNode { + reference: crate::automation::InstanceReference { + session_id: "plugin-session".to_owned(), + id: "pinst-00000001".to_owned(), + path: "Workspace.Map".to_owned(), + name: "Map".to_owned(), + class_name: "Model".to_owned(), + }, + name: "Map".to_owned(), + class_name: "Model".to_owned(), + path: "Workspace.Map".to_owned(), + properties: BTreeMap::from([( + "Archivable".to_owned(), + AutomationValue::Boolean { value: true }, + )]), + attributes: BTreeMap::from([( + "Theme".to_owned(), + AutomationValue::String { + value: "Night".to_owned(), + }, + )]), + tags: vec!["Generated".to_owned()], + children: Vec::new(), + truncated: true, + }, + visited_instances: 1, + truncated: true, + truncation_reason: Some("maxChildren".to_owned()), + } + } + + #[test] + fn renders_human_and_json_output() { + let result = sample_result(); + let mut output = Vec::new(); + render_human(&result, &mut output).unwrap(); + let output = String::from_utf8(output).unwrap(); + assert!(output.contains("Workspace.Map [Model]")); + assert!(output.contains("Theme = \"Night\"")); + assert!(output.contains("[truncated:")); + + let json = serde_json::to_string(&result).unwrap(); + let decoded: serde_json::Value = serde_json::from_str(&json).unwrap(); + assert_eq!(decoded["root"]["reference"]["sessionId"], "plugin-session"); + } + + #[test] + fn reports_status_failures() { + let mut status = AutomationStatusResponse { + server_session_id: crate::SessionId::new(), + server_version: "test".to_owned(), + protocol_version: 5, + automation_handler_version: 2, + automation_available: false, + exec_available: false, + typed_automation_available: false, + plugin: None, + duplicate_session_detected: false, + queues: crate::web_api::AutomationQueueStatusResponse { + exec_pending: 0, + exec_claimed: 0, + exec_claimed_by_plugin_session_id: None, + automation_pending: 0, + automation_claimed: 0, + automation_claimed_by_plugin_session_id: None, + }, + }; + assert!(validate_automation_status(&status).is_err()); + status.duplicate_session_detected = true; + assert!(validate_automation_status(&status) + .unwrap_err() + .to_string() + .contains("multiple")); + } +} diff --git a/src/cli/mod.rs b/src/cli/mod.rs index 39bbc809..60837654 100644 --- a/src/cli/mod.rs +++ b/src/cli/mod.rs @@ -1,10 +1,12 @@ //! Defines Rojo's CLI through clap types. +mod automation; mod build; mod doc; mod exec; mod fmt_project; mod init; +mod inspect; mod plugin; mod serve; mod sourcemap; @@ -22,6 +24,7 @@ pub use self::doc::DocCommand; pub use self::exec::ExecCommand; pub use self::fmt_project::FmtProjectCommand; pub use self::init::{InitCommand, InitKind}; +pub use self::inspect::InspectCommand; pub use self::plugin::{PluginCommand, PluginSubcommand}; pub use self::serve::ServeCommand; pub use self::sourcemap::SourcemapCommand; @@ -47,6 +50,7 @@ impl Options { Subcommand::Serve(subcommand) => subcommand.run(self.global), Subcommand::Build(subcommand) => subcommand.run(), Subcommand::Exec(subcommand) => subcommand.run(), + Subcommand::Inspect(subcommand) => subcommand.run(), Subcommand::Upload(subcommand) => subcommand.run(), Subcommand::Sourcemap(subcommand) => subcommand.run(), Subcommand::FmtProject(subcommand) => subcommand.run(), @@ -122,6 +126,7 @@ pub enum Subcommand { Serve(ServeCommand), Build(BuildCommand), Exec(ExecCommand), + Inspect(InspectCommand), Upload(UploadCommand), Sourcemap(SourcemapCommand), FmtProject(FmtProjectCommand), diff --git a/src/exec.rs b/src/exec.rs index e0c558e5..3d9309d0 100644 --- a/src/exec.rs +++ b/src/exec.rs @@ -100,6 +100,12 @@ pub struct ExecCleanupStats { pub removed_terminal: usize, } +#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)] +pub struct ExecQueueCounts { + pub pending: usize, + pub claimed: usize, +} + #[derive(Debug)] pub struct ExecJobStore { inner: Mutex, @@ -181,6 +187,14 @@ impl ExecJobStore { self.cleanup_expired_at(Instant::now()) } + pub fn queue_counts(&self) -> ExecQueueCounts { + let inner = self.inner.lock().unwrap(); + ExecQueueCounts { + pending: inner.pending_jobs.len(), + claimed: usize::from(inner.claimed_job.is_some()), + } + } + fn with_durations( claim_timeout: Duration, execution_timeout: Duration, diff --git a/src/lib.rs b/src/lib.rs index ca574581..b6d99d18 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -9,6 +9,8 @@ pub mod exec; mod tree_view; mod auth_cookie; +pub mod automation; +mod automation_status; mod change_processor; mod glob; mod json; diff --git a/src/serve_session.rs b/src/serve_session.rs index 601464ae..e1aa46c7 100644 --- a/src/serve_session.rs +++ b/src/serve_session.rs @@ -12,6 +12,8 @@ use memofs::Vfs; use thiserror::Error; use crate::{ + automation::AutomationJobStore, + automation_status::AutomationStatusStore, change_processor::ChangeProcessor, exec::ExecJobStore, message_queue::MessageQueue, @@ -86,6 +88,12 @@ pub struct ServeSession { /// Execution jobs submitted to this serve session. exec_job_store: ExecJobStore, + /// Typed one-shot automation jobs submitted to this serve session. + automation_job_store: AutomationJobStore, + + /// Studio plugin identity and availability for the automation API. + automation_status: AutomationStatusStore, + /// A channel to send mutation requests on. These will be handled by the /// ChangeProcessor and trigger changes in the tree. tree_mutation_sender: Sender, @@ -148,6 +156,8 @@ impl ServeSession { tree, message_queue, exec_job_store: ExecJobStore::new(), + automation_job_store: AutomationJobStore::new(), + automation_status: AutomationStatusStore::new(), tree_mutation_sender, vfs, }) @@ -178,6 +188,14 @@ impl ServeSession { &self.exec_job_store } + pub fn automation_job_store(&self) -> &AutomationJobStore { + &self.automation_job_store + } + + pub fn automation_status(&self) -> &AutomationStatusStore { + &self.automation_status + } + pub fn session_id(&self) -> SessionId { self.session_id } diff --git a/src/web/api.rs b/src/web/api.rs index 33d30803..b2a0c9a1 100644 --- a/src/web/api.rs +++ b/src/web/api.rs @@ -18,15 +18,28 @@ use rbx_dom_weak::{ use uuid::Uuid; use crate::{ + automation::{ + AutomationJobStoreError, AutomationRequest, MAX_AUTOMATION_REQUEST_BODY_BYTES, + MAX_AUTOMATION_RESULT_BODY_BYTES, + }, + automation_status::{ + AutomationRegistration, AutomationSessionUpdate, PluginSessionId, + AUTOMATION_HANDLER_VERSION, + }, exec::{ExecJobStoreError, MAX_SOURCE_SIZE_BYTES}, serve_session::ServeSession, snapshot::{InstanceWithMeta, PatchSet, PatchUpdate}, web::{ interface::{ - ErrorResponse, ExecJobClaimResponse, ExecJobCompletionRequest, ExecJobResponse, + AutomationHeartbeatRequest, AutomationHeartbeatResponse, AutomationJobClaimResponse, + AutomationJobCompletion, AutomationJobCompletionRequest, AutomationJobResponse, + AutomationPluginStatusResponse, AutomationQueueStatusResponse, + AutomationStatusResponse, ErrorResponse, ExecJobClaimResponse, + ExecJobCompletionEnvelope, ExecJobCompletionRequest, ExecJobResponse, ExecJobSubmissionRequest, ExecLog, Instance, MessagesPacket, OpenResponse, ReadResponse, ServerInfoResponse, SocketPacket, SocketPacketBody, SocketPacketType, - SubscribeMessage, WriteRequest, WriteResponse, PROTOCOL_VERSION, SERVER_VERSION, + StudioMode, SubscribeMessage, WriteRequest, WriteResponse, PROTOCOL_VERSION, + SERVER_VERSION, }, origin::canonical, util::{deserialize_msgpack, msgpack, msgpack_ok, response, serialize_msgpack}, @@ -38,6 +51,7 @@ use crate::{ const EXEC_SUBMISSION_BODY_LIMIT_BYTES: usize = MAX_SOURCE_SIZE_BYTES + 64 * 1024; const EXEC_COMPLETION_BODY_LIMIT_BYTES: usize = 256 * 1024; +const AUTOMATION_STATUS_BODY_LIMIT_BYTES: usize = 16 * 1024; enum ExecRoute { Submit, @@ -46,6 +60,35 @@ enum ExecRoute { Complete(String), } +enum AutomationRoute { + Submit, + ClaimNext, + Status(String), + Complete(String), +} + +impl AutomationRoute { + fn parse(method: &Method, path: &str) -> Option { + match (method, path) { + (&Method::POST, "/api/automation/jobs") => return Some(Self::Submit), + (&Method::GET, "/api/automation/jobs/next") => return Some(Self::ClaimNext), + _ => {} + } + let remainder = path.strip_prefix("/api/automation/jobs/")?; + let mut segments = remainder.split('/'); + let id = segments.next()?; + let suffix = segments.next(); + if id.is_empty() || segments.next().is_some() { + return None; + } + match (method, suffix) { + (&Method::GET, None) => Some(Self::Status(id.to_owned())), + (&Method::POST, Some("complete")) => Some(Self::Complete(id.to_owned())), + _ => None, + } + } +} + impl ExecRoute { fn parse(method: &Method, path: &str) -> Option { match (method, path) { @@ -87,14 +130,46 @@ pub async fn call( return match route { ExecRoute::Submit => service.handle_api_exec_submit(request).await, - ExecRoute::ClaimNext => service.handle_api_exec_claim_next(), + ExecRoute::ClaimNext => service.handle_api_exec_claim_next(&request), ExecRoute::Status(id) => service.handle_api_exec_status(&id), ExecRoute::Complete(id) => service.handle_api_exec_complete(&id, request).await, }; } + if let Some(route) = AutomationRoute::parse(request.method(), request.uri().path()) { + if let Some(response) = service.reject_non_local("/api/automation/jobs") { + return response; + } + service + .serve_session + .automation_job_store() + .cleanup_expired(); + return match route { + AutomationRoute::Submit => service.handle_api_automation_submit(request).await, + AutomationRoute::ClaimNext => service.handle_api_automation_claim_next(&request), + AutomationRoute::Status(id) => service.handle_api_automation_job_status(&id), + AutomationRoute::Complete(id) => { + service.handle_api_automation_complete(&id, request).await + } + }; + } + match (request.method(), request.uri().path()) { (&Method::GET, "/api/rojo") => service.handle_api_rojo().await, + (&Method::GET, "/api/automation/status") => { + if let Some(response) = service.reject_non_local("/api/automation/status") { + response + } else { + service.handle_api_automation_status() + } + } + (&Method::POST, "/api/automation/status") => { + if let Some(response) = service.reject_non_local("/api/automation/status") { + response + } else { + service.handle_api_automation_heartbeat(request).await + } + } (&Method::GET, path) if path.starts_with("/api/read/") => { service.handle_api_read(request).await } @@ -150,7 +225,12 @@ impl ApiService { } async fn handle_api_exec_submit(&self, request: Request) -> Response { - let body = match read_exec_body(request.into_body(), EXEC_SUBMISSION_BODY_LIMIT_BYTES).await + let body = match read_bounded_body( + request.into_body(), + EXEC_SUBMISSION_BODY_LIMIT_BYTES, + "Exec", + ) + .await { Ok(body) => body, Err(response) => return response, @@ -189,9 +269,27 @@ impl ApiService { } } - fn handle_api_exec_claim_next(&self) -> Response { + fn handle_api_exec_claim_next(&self, request: &Request) -> Response { + let plugin_session = match parse_exec_session_query(request.uri().query()) { + Ok(plugin_session) => plugin_session, + Err(response) => return *response, + }; + + if let Some((plugin_session_id, studio_mode)) = plugin_session { + if let Err(response) = self.refresh_automation_session(plugin_session_id, studio_mode) { + return *response; + } + } + match self.serve_session.exec_job_store().claim_next() { - Some(job) => no_store(msgpack_ok(ExecJobClaimResponse::from(job))), + Some(job) => { + if let Some((plugin_session_id, _)) = plugin_session { + self.serve_session + .automation_status() + .note_exec_claim(plugin_session_id); + } + no_store(msgpack_ok(ExecJobClaimResponse::from(job))) + } None => no_store(response( StatusCode::NO_CONTENT, "application/msgpack", @@ -228,12 +326,17 @@ impl ApiService { } }; - let body = match read_exec_body(request.into_body(), EXEC_COMPLETION_BODY_LIMIT_BYTES).await + let body = match read_bounded_body( + request.into_body(), + EXEC_COMPLETION_BODY_LIMIT_BYTES, + "Exec", + ) + .await { Ok(body) => body, Err(response) => return response, }; - let request: ExecJobCompletionRequest = match deserialize_msgpack(&body) { + let request: ExecJobCompletionEnvelope = match deserialize_msgpack(&body) { Ok(request) => request, Err(err) => { return msgpack( @@ -243,8 +346,21 @@ impl ApiService { } }; + let plugin_session = match parse_exec_completion_session( + request.plugin_session_id.as_deref(), + request.studio_mode, + ) { + Ok(plugin_session) => plugin_session, + Err(response) => return *response, + }; + if let Some((plugin_session_id, studio_mode)) = plugin_session { + if let Err(response) = self.refresh_automation_session(plugin_session_id, studio_mode) { + return *response; + } + } + let store = self.serve_session.exec_job_store(); - let result = match request { + let result = match request.completion { ExecJobCompletionRequest::Success { result, logs } => { store.complete_success(id, result.map(Into::into), into_stored_logs(logs)) } @@ -271,11 +387,338 @@ impl ApiService { }; match result { - Ok(job) => no_store(msgpack_ok(ExecJobResponse::from(job))), + Ok(job) => { + if let Some((plugin_session_id, _)) = plugin_session { + self.serve_session + .automation_status() + .note_exec_completion(plugin_session_id); + } + no_store(msgpack_ok(ExecJobResponse::from(job))) + } Err(err) => exec_store_error(err), } } + async fn handle_api_automation_heartbeat(&self, request: Request) -> Response { + let body = match read_bounded_body( + request.into_body(), + AUTOMATION_STATUS_BODY_LIMIT_BYTES, + "Automation status", + ) + .await + { + Ok(body) => body, + Err(response) => return response, + }; + let request: AutomationHeartbeatRequest = match deserialize_msgpack(&body) { + Ok(request) => request, + Err(err) => { + return msgpack( + ErrorResponse::bad_request(format!("Invalid body: {err}")), + StatusCode::BAD_REQUEST, + ); + } + }; + + if request.server_session_id != self.serve_session.session_id() { + return msgpack( + ErrorResponse::conflict("Automation heartbeat targets a different server session"), + StatusCode::CONFLICT, + ); + } + + let plugin_session_id = match request.plugin_session_id.parse::() { + Ok(plugin_session_id) => plugin_session_id, + Err(err) => { + return msgpack( + ErrorResponse::bad_request(format!("Malformed plugin session ID: {err}")), + StatusCode::BAD_REQUEST, + ); + } + }; + let registration = + self.serve_session + .automation_status() + .refresh(AutomationSessionUpdate { + plugin_session_id, + server_session_id: request.server_session_id, + studio_mode: request.studio_mode, + exec_handler_available: request.exec_handler_available, + automation_handler_version: request.automation_handler_version, + plugin_version: request.plugin_version, + }); + let active_plugin_session_id = self + .serve_session + .automation_status() + .snapshot() + .plugin + .map(|plugin| plugin.plugin_session_id.to_string()); + + no_store(msgpack_ok(AutomationHeartbeatResponse { + registration, + active_plugin_session_id, + })) + } + + async fn handle_api_automation_submit(&self, request: Request) -> Response { + let body = match read_bounded_body( + request.into_body(), + MAX_AUTOMATION_REQUEST_BODY_BYTES, + "Automation job", + ) + .await + { + Ok(body) => body, + Err(response) => return response, + }; + let request: AutomationRequest = match deserialize_msgpack(&body) { + Ok(request) => request, + Err(err) => { + return msgpack( + ErrorResponse::bad_request(format!("Invalid automation request: {err}")), + StatusCode::BAD_REQUEST, + ); + } + }; + match self.serve_session.automation_job_store().submit(request) { + Ok(job) => no_store(msgpack( + AutomationJobResponse::from(job), + StatusCode::CREATED, + )), + Err(error) => automation_store_error(error), + } + } + + fn handle_api_automation_claim_next(&self, request: &Request) -> Response { + let (plugin_session_id, studio_mode) = match parse_exec_session_query(request.uri().query()) + { + Ok(Some(metadata)) => metadata, + Ok(None) => { + return msgpack( + ErrorResponse::bad_request( + "Automation claim requires pluginSessionId and studioMode", + ), + StatusCode::BAD_REQUEST, + ); + } + Err(response) => return *response, + }; + if studio_mode != StudioMode::Edit { + return msgpack( + ErrorResponse::conflict("Typed automation jobs may only be claimed in edit mode"), + StatusCode::CONFLICT, + ); + } + if let Err(response) = self.refresh_typed_automation_session(plugin_session_id, studio_mode) + { + return *response; + } + match self + .serve_session + .automation_job_store() + .claim_next(plugin_session_id) + { + Some(job) => no_store(msgpack_ok(AutomationJobClaimResponse::from(job))), + None => no_store(response( + StatusCode::NO_CONTENT, + "application/msgpack", + Body::empty(), + )), + } + } + + fn handle_api_automation_job_status(&self, id: &str) -> Response { + let id = match Uuid::parse_str(id) { + Ok(id) => id, + Err(err) => { + return msgpack( + ErrorResponse::bad_request(format!("Malformed automation job ID: {err}")), + StatusCode::BAD_REQUEST, + ); + } + }; + match self.serve_session.automation_job_store().get(id) { + Ok(job) => no_store(msgpack_ok(AutomationJobResponse::from(job))), + Err(error) => automation_store_error(error), + } + } + + async fn handle_api_automation_complete( + &self, + id: &str, + request: Request, + ) -> Response { + let id = match Uuid::parse_str(id) { + Ok(id) => id, + Err(err) => { + return msgpack( + ErrorResponse::bad_request(format!("Malformed automation job ID: {err}")), + StatusCode::BAD_REQUEST, + ); + } + }; + let body = match read_bounded_body( + request.into_body(), + MAX_AUTOMATION_RESULT_BODY_BYTES, + "Automation completion", + ) + .await + { + Ok(body) => body, + Err(response) => return response, + }; + let request: AutomationJobCompletionRequest = match deserialize_msgpack(&body) { + Ok(request) => request, + Err(err) => { + return msgpack( + ErrorResponse::bad_request(format!("Invalid automation completion: {err}")), + StatusCode::BAD_REQUEST, + ); + } + }; + let plugin_session_id = match request.plugin_session_id.parse::() { + Ok(id) => id, + Err(err) => { + return msgpack( + ErrorResponse::bad_request(format!("Malformed plugin session ID: {err}")), + StatusCode::BAD_REQUEST, + ); + } + }; + if let Err(response) = + self.refresh_typed_automation_session(plugin_session_id, request.studio_mode) + { + return *response; + } + let store = self.serve_session.automation_job_store(); + let completed = match request.completion { + AutomationJobCompletion::Success { result } => { + store.complete_success(id, plugin_session_id, *result) + } + AutomationJobCompletion::Failure { error } => { + store.complete_failure(id, plugin_session_id, error) + } + }; + match completed { + Ok(job) => no_store(msgpack_ok(AutomationJobResponse::from(job))), + Err(error) => automation_store_error(error), + } + } + + fn handle_api_automation_status(&self) -> Response { + self.serve_session.exec_job_store().cleanup_expired(); + let queue_counts = self.serve_session.exec_job_store().queue_counts(); + self.serve_session.automation_job_store().cleanup_expired(); + let automation_queue_counts = self.serve_session.automation_job_store().queue_counts(); + self.serve_session + .automation_status() + .clear_exec_claim_if_idle(queue_counts.claimed != 0); + let status = self.serve_session.automation_status().snapshot(); + let automation_available = status.plugin.as_ref().is_some_and(|plugin| { + plugin.connected && plugin.automation_handler_version == AUTOMATION_HANDLER_VERSION + }); + let exec_available = automation_available + && status + .plugin + .as_ref() + .is_some_and(|plugin| plugin.exec_handler_available); + let plugin = status.plugin.map(|plugin| AutomationPluginStatusResponse { + connected: plugin.connected, + plugin_session_id: plugin.plugin_session_id.to_string(), + studio_mode: plugin.studio_mode, + plugin_version: plugin.plugin_version, + automation_handler_version: plugin.automation_handler_version, + last_seen_at_ms: plugin.last_seen_at_ms, + }); + + no_store(msgpack_ok(AutomationStatusResponse { + server_session_id: self.serve_session.session_id(), + server_version: SERVER_VERSION.to_owned(), + protocol_version: PROTOCOL_VERSION, + automation_handler_version: AUTOMATION_HANDLER_VERSION, + automation_available, + exec_available, + typed_automation_available: automation_available, + plugin, + duplicate_session_detected: status.duplicate_session_detected, + queues: AutomationQueueStatusResponse { + exec_pending: queue_counts.pending, + exec_claimed: queue_counts.claimed, + exec_claimed_by_plugin_session_id: status + .exec_claimed_by_plugin_session_id + .map(|id| id.to_string()), + automation_pending: automation_queue_counts.pending, + automation_claimed: automation_queue_counts.claimed, + automation_claimed_by_plugin_session_id: self + .serve_session + .automation_job_store() + .claimed_by() + .map(|id| id.to_string()), + }, + })) + } + + fn refresh_automation_session( + &self, + plugin_session_id: PluginSessionId, + studio_mode: StudioMode, + ) -> Result<(), Box>> { + let snapshot = self.serve_session.automation_status().snapshot(); + let existing = snapshot + .plugin + .filter(|plugin| plugin.plugin_session_id == plugin_session_id); + let registration = + self.serve_session + .automation_status() + .refresh(AutomationSessionUpdate { + plugin_session_id, + server_session_id: self.serve_session.session_id(), + studio_mode, + exec_handler_available: existing + .as_ref() + .is_none_or(|plugin| plugin.exec_handler_available), + automation_handler_version: existing + .as_ref() + .map_or(AUTOMATION_HANDLER_VERSION, |plugin| { + plugin.automation_handler_version + }), + plugin_version: None, + }); + + if registration == AutomationRegistration::Conflict { + return Err(Box::new(msgpack( + ErrorResponse::conflict( + "Another active Prism plugin session is already registered for automation", + ), + StatusCode::CONFLICT, + ))); + } + + Ok(()) + } + + fn refresh_typed_automation_session( + &self, + plugin_session_id: PluginSessionId, + studio_mode: StudioMode, + ) -> Result<(), Box>> { + let snapshot = self.serve_session.automation_status().snapshot(); + let authoritative = snapshot.plugin.as_ref().is_some_and(|plugin| { + plugin.connected + && plugin.plugin_session_id == plugin_session_id + && plugin.automation_handler_version == AUTOMATION_HANDLER_VERSION + }); + if !authoritative { + return Err(Box::new(msgpack( + ErrorResponse::conflict( + "The claimant is not the active compatible Prism automation plugin session", + ), + StatusCode::CONFLICT, + ))); + } + self.refresh_automation_session(plugin_session_id, studio_mode) + } + /// Get a summary of information about the server async fn handle_api_rojo(&self) -> Response { let tree = self.serve_session.tree(); @@ -648,13 +1091,110 @@ impl ApiService { } } -async fn read_exec_body(mut body: Body, limit: usize) -> Result, Response> { +fn parse_exec_session_query( + query: Option<&str>, +) -> Result, Box>> { + let Some(query) = query else { + return Ok(None); + }; + + let mut plugin_session_id = None; + let mut studio_mode = None; + for pair in query.split('&') { + let Some((key, value)) = pair.split_once('=') else { + return Err(Box::new(msgpack( + ErrorResponse::bad_request("Malformed exec claim query"), + StatusCode::BAD_REQUEST, + ))); + }; + match key { + "pluginSessionId" if plugin_session_id.is_none() => { + plugin_session_id = Some(value); + } + "studioMode" if studio_mode.is_none() => studio_mode = Some(value), + _ => { + return Err(Box::new(msgpack( + ErrorResponse::bad_request("Malformed exec claim query"), + StatusCode::BAD_REQUEST, + ))); + } + } + } + + parse_exec_session(plugin_session_id, studio_mode) +} + +fn parse_exec_completion_session( + plugin_session_id: Option<&str>, + studio_mode: Option, +) -> Result, Box>> { + match (plugin_session_id, studio_mode) { + (None, None) => Ok(None), + (Some(plugin_session_id), Some(studio_mode)) => { + let plugin_session_id = plugin_session_id.parse().map_err(|err| { + Box::new(msgpack( + ErrorResponse::bad_request(format!("Malformed plugin session ID: {err}")), + StatusCode::BAD_REQUEST, + )) + })?; + Ok(Some((plugin_session_id, studio_mode))) + } + _ => Err(Box::new(msgpack( + ErrorResponse::bad_request( + "Exec session metadata requires both pluginSessionId and studioMode", + ), + StatusCode::BAD_REQUEST, + ))), + } +} + +fn parse_exec_session( + plugin_session_id: Option<&str>, + studio_mode: Option<&str>, +) -> Result, Box>> { + match (plugin_session_id, studio_mode) { + (None, None) => Ok(None), + (Some(plugin_session_id), Some(studio_mode)) => { + let plugin_session_id = plugin_session_id.parse().map_err(|err| { + Box::new(msgpack( + ErrorResponse::bad_request(format!("Malformed plugin session ID: {err}")), + StatusCode::BAD_REQUEST, + )) + })?; + let studio_mode = match studio_mode { + "edit" => StudioMode::Edit, + "play" => StudioMode::Play, + "run" => StudioMode::Run, + "unknown" => StudioMode::Unknown, + _ => { + return Err(Box::new(msgpack( + ErrorResponse::bad_request("Malformed Studio mode"), + StatusCode::BAD_REQUEST, + ))); + } + }; + Ok(Some((plugin_session_id, studio_mode))) + } + _ => Err(Box::new(msgpack( + ErrorResponse::bad_request( + "Exec claim requires both pluginSessionId and studioMode query parameters", + ), + StatusCode::BAD_REQUEST, + ))), + } +} + +async fn read_bounded_body( + mut body: Body, + limit: usize, + request_name: &str, +) -> Result, Response> { if body .size_hint() .upper() .is_some_and(|size| size > limit as u64) { - return Err(exec_body_too_large(limit)); + return Err(body_too_large(request_name, limit)); } let capacity = body.size_hint().lower().min(limit as u64) as usize; @@ -672,7 +1212,7 @@ async fn read_exec_body(mut body: Body, limit: usize) -> Result, Respons }; if bytes.len().saturating_add(chunk.len()) > limit { - return Err(exec_body_too_large(limit)); + return Err(body_too_large(request_name, limit)); } bytes.extend_from_slice(&chunk); @@ -681,10 +1221,10 @@ async fn read_exec_body(mut body: Body, limit: usize) -> Result, Respons Ok(bytes) } -fn exec_body_too_large(limit: usize) -> Response { +fn body_too_large(request_name: &str, limit: usize) -> Response { msgpack( ErrorResponse::payload_too_large(format!( - "Exec request body exceeds the {limit}-byte limit" + "{request_name} request body exceeds the {limit}-byte limit" )), StatusCode::PAYLOAD_TOO_LARGE, ) @@ -711,6 +1251,28 @@ fn exec_store_error(error: ExecJobStoreError) -> Response { } } +fn automation_store_error(error: AutomationJobStoreError) -> Response { + let details = error.to_string(); + match error { + AutomationJobStoreError::InvalidRequest(_) => { + msgpack(ErrorResponse::bad_request(details), StatusCode::BAD_REQUEST) + } + AutomationJobStoreError::PendingQueueFull { .. } => msgpack( + ErrorResponse::too_many_requests(details), + StatusCode::TOO_MANY_REQUESTS, + ), + AutomationJobStoreError::UnknownJob { .. } => { + msgpack(ErrorResponse::not_found(details), StatusCode::NOT_FOUND) + } + AutomationJobStoreError::JobNotClaimed { .. } + | AutomationJobStoreError::DuplicateCompletion { .. } + | AutomationJobStoreError::WrongClaimant { .. } + | AutomationJobStoreError::ResultKindMismatch { .. } => { + msgpack(ErrorResponse::conflict(details), StatusCode::CONFLICT) + } + } +} + fn into_stored_logs(logs: Vec) -> Option> { if logs.is_empty() { None @@ -883,7 +1445,7 @@ mod test { use super::*; #[tokio::test] - async fn exec_routes_reject_non_loopback_peers() { + async fn local_automation_routes_reject_non_loopback_peers() { let dir = tempfile::tempdir().unwrap(); std::fs::write( dir.path().join("default.project.json"), @@ -900,6 +1462,12 @@ mod test { (Method::GET, "/api/exec/jobs/next".to_owned()), (Method::GET, format!("/api/exec/jobs/{id}")), (Method::POST, format!("/api/exec/jobs/{id}/complete")), + (Method::POST, "/api/automation/jobs".to_owned()), + (Method::GET, "/api/automation/jobs/next".to_owned()), + (Method::GET, format!("/api/automation/jobs/{id}")), + (Method::POST, format!("/api/automation/jobs/{id}/complete")), + (Method::GET, "/api/automation/status".to_owned()), + (Method::POST, "/api/automation/status".to_owned()), ]; for (method, route) in routes { diff --git a/src/web/interface.rs b/src/web/interface.rs index ba1db11f..82436984 100644 --- a/src/web/interface.rs +++ b/src/web/interface.rs @@ -14,6 +14,13 @@ use rbx_dom_weak::{ use serde::{Deserialize, Serialize}; use strum::Display; +pub use crate::automation::{ + AutomationJobState, AutomationMapEntry, AutomationRequest, AutomationResult, AutomationUdim, + AutomationValue, AutomationVector2, InspectNode, InspectRequest, InspectResult, InspectTarget, + InstanceReference, +}; +pub use crate::automation_status::{AutomationRegistration, StudioMode}; + use crate::{ exec::{ ExecJob as StoredExecJob, ExecJobState as StoredExecJobState, ExecLog as StoredExecLog, @@ -488,6 +495,143 @@ pub enum ExecJobCompletionRequest { }, } +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +#[serde(rename_all = "camelCase")] +pub struct ExecJobCompletionEnvelope { + #[serde(flatten)] + pub completion: ExecJobCompletionRequest, + #[serde(default)] + pub plugin_session_id: Option, + #[serde(default)] + pub studio_mode: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "camelCase")] +pub struct AutomationHeartbeatRequest { + pub plugin_session_id: String, + pub server_session_id: SessionId, + pub studio_mode: StudioMode, + pub exec_handler_available: bool, + pub automation_handler_version: u32, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub plugin_version: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "camelCase")] +pub struct AutomationHeartbeatResponse { + pub registration: AutomationRegistration, + #[serde(skip_serializing_if = "Option::is_none")] + pub active_plugin_session_id: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "camelCase")] +pub struct AutomationPluginStatusResponse { + pub connected: bool, + pub plugin_session_id: String, + pub studio_mode: StudioMode, + #[serde(skip_serializing_if = "Option::is_none")] + pub plugin_version: Option, + pub automation_handler_version: u32, + pub last_seen_at_ms: u64, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "camelCase")] +pub struct AutomationQueueStatusResponse { + pub exec_pending: usize, + pub exec_claimed: usize, + #[serde(skip_serializing_if = "Option::is_none")] + pub exec_claimed_by_plugin_session_id: Option, + pub automation_pending: usize, + pub automation_claimed: usize, + #[serde(skip_serializing_if = "Option::is_none")] + pub automation_claimed_by_plugin_session_id: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +#[serde(rename_all = "camelCase")] +pub struct AutomationStatusResponse { + pub server_session_id: SessionId, + pub server_version: String, + pub protocol_version: u64, + pub automation_handler_version: u32, + pub automation_available: bool, + pub exec_available: bool, + pub typed_automation_available: bool, + #[serde(skip_serializing_if = "Option::is_none")] + pub plugin: Option, + pub duplicate_session_detected: bool, + pub queues: AutomationQueueStatusResponse, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +#[serde(rename_all = "camelCase")] +pub struct AutomationJobResponse { + pub job_id: String, + pub state: AutomationJobState, + #[serde(skip_serializing_if = "Option::is_none")] + pub claimed_by_plugin_session_id: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub result: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub error: Option, +} + +impl From for AutomationJobResponse { + fn from(job: crate::automation::AutomationJob) -> Self { + Self { + job_id: job.id.to_string(), + state: job.state, + claimed_by_plugin_session_id: job.claimed_by.map(|id| id.to_string()), + result: job.result, + error: job.error, + } + } +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +#[serde(rename_all = "camelCase")] +pub struct AutomationJobClaimResponse { + pub job_id: String, + pub state: AutomationJobState, + pub request: AutomationRequest, + pub execution_timeout_ms: u64, +} + +impl From for AutomationJobClaimResponse { + fn from(job: crate::automation::AutomationJob) -> Self { + Self { + job_id: job.id.to_string(), + state: job.state, + request: job.request, + execution_timeout_ms: job + .execution_timeout + .as_millis() + .try_into() + .unwrap_or(u64::MAX), + } + } +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +#[serde(tag = "outcome", rename_all = "camelCase")] +pub enum AutomationJobCompletion { + Success { result: Box }, + Failure { error: String }, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +#[serde(rename_all = "camelCase")] +pub struct AutomationJobCompletionRequest { + #[serde(flatten)] + pub completion: AutomationJobCompletion, + pub plugin_session_id: String, + pub studio_mode: StudioMode, +} + /// General response type returned from all Rojo routes #[derive(Debug, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] diff --git a/tests/rojo_test/serve_util.rs b/tests/rojo_test/serve_util.rs index 905687b5..03f86454 100644 --- a/tests/rojo_test/serve_util.rs +++ b/tests/rojo_test/serve_util.rs @@ -15,8 +15,10 @@ use tempfile::{tempdir, TempDir}; use librojo::{ web_api::{ - ExecJobCompletionRequest, ExecJobSubmissionRequest, ReadResponse, SerializeRequest, - SerializeResponse, ServerInfoResponse, SocketPacket, SocketPacketType, + AutomationHeartbeatRequest, AutomationJobCompletionRequest, AutomationRequest, + ExecJobCompletionEnvelope, ExecJobCompletionRequest, ExecJobSubmissionRequest, + ReadResponse, SerializeRequest, SerializeResponse, ServerInfoResponse, SocketPacket, + SocketPacketType, StudioMode, }, SessionId, }; @@ -289,6 +291,87 @@ impl TestServeSession { self.api_request(reqwest::Method::GET, "/api/exec/jobs/next", None) } + pub fn get_api_exec_next_for_session( + &self, + plugin_session_id: &str, + studio_mode: StudioMode, + ) -> reqwest::blocking::Response { + let studio_mode = match studio_mode { + StudioMode::Edit => "edit", + StudioMode::Play => "play", + StudioMode::Run => "run", + StudioMode::Unknown => "unknown", + }; + self.api_request( + reqwest::Method::GET, + &format!( + "/api/exec/jobs/next?pluginSessionId={plugin_session_id}&studioMode={studio_mode}" + ), + None, + ) + } + + pub fn get_api_automation_status(&self) -> reqwest::blocking::Response { + self.api_request(reqwest::Method::GET, "/api/automation/status", None) + } + + pub fn post_api_automation_heartbeat( + &self, + request: &AutomationHeartbeatRequest, + ) -> reqwest::blocking::Response { + let body = serialize_msgpack(request).unwrap(); + self.api_request(reqwest::Method::POST, "/api/automation/status", Some(body)) + } + + pub fn post_api_automation_job( + &self, + request: &AutomationRequest, + ) -> reqwest::blocking::Response { + let body = serialize_msgpack(request).unwrap(); + self.api_request(reqwest::Method::POST, "/api/automation/jobs", Some(body)) + } + + pub fn get_api_automation_next( + &self, + plugin_session_id: &str, + studio_mode: StudioMode, + ) -> reqwest::blocking::Response { + let studio_mode = match studio_mode { + StudioMode::Edit => "edit", + StudioMode::Play => "play", + StudioMode::Run => "run", + StudioMode::Unknown => "unknown", + }; + self.api_request( + reqwest::Method::GET, + &format!( + "/api/automation/jobs/next?pluginSessionId={plugin_session_id}&studioMode={studio_mode}" + ), + None, + ) + } + + pub fn get_api_automation_job_status(&self, id: &str) -> reqwest::blocking::Response { + self.api_request( + reqwest::Method::GET, + &format!("/api/automation/jobs/{id}"), + None, + ) + } + + pub fn post_api_automation_complete( + &self, + id: &str, + request: &AutomationJobCompletionRequest, + ) -> reqwest::blocking::Response { + let body = serialize_msgpack(request).unwrap(); + self.api_request( + reqwest::Method::POST, + &format!("/api/automation/jobs/{id}/complete"), + Some(body), + ) + } + pub fn get_api_exec_status(&self, id: &str) -> reqwest::blocking::Response { self.api_request(reqwest::Method::GET, &format!("/api/exec/jobs/{id}"), None) } @@ -306,6 +389,19 @@ impl TestServeSession { ) } + pub fn post_api_exec_complete_for_session( + &self, + id: &str, + request: &ExecJobCompletionEnvelope, + ) -> reqwest::blocking::Response { + let body = serialize_msgpack(request).unwrap(); + self.api_request( + reqwest::Method::POST, + &format!("/api/exec/jobs/{id}/complete"), + Some(body), + ) + } + pub fn api_request( &self, method: reqwest::Method, diff --git a/tests/tests/serve.rs b/tests/tests/serve.rs index 688280de..74adac29 100644 --- a/tests/tests/serve.rs +++ b/tests/tests/serve.rs @@ -1,4 +1,4 @@ -use std::fs; +use std::{collections::BTreeMap, fs}; use insta::{assert_snapshot, assert_yaml_snapshot, with_settings}; use rbx_dom_weak::types::Ref; @@ -13,14 +13,398 @@ use crate::rojo_test::{ }; use librojo::{ + automation::{ + AutomationJobState, AutomationRequest, AutomationResult, InspectNode, InspectRequest, + InspectResult, InspectTarget, InstanceReference, MAX_AUTOMATION_REQUEST_BODY_BYTES, + MAX_PENDING_AUTOMATION_JOBS, + }, exec::{MAX_PENDING_JOBS, MAX_SOURCE_SIZE_BYTES}, web_api::{ - ExecJobClaimResponse, ExecJobCompletionRequest, ExecJobResponse, + AutomationHeartbeatRequest, AutomationHeartbeatResponse, AutomationJobClaimResponse, + AutomationJobCompletion, AutomationJobCompletionRequest, AutomationJobResponse, + AutomationRegistration, AutomationStatusResponse, ExecJobClaimResponse, + ExecJobCompletionEnvelope, ExecJobCompletionRequest, ExecJobResponse, ExecJobState as ApiExecJobState, ExecJobSubmissionRequest, ExecLog, ExecLogLevel, - ExecValue, SerializeResponse, SocketPacketType, + ExecValue, SerializeResponse, SocketPacketType, StudioMode, }, }; +fn inspect_request(target: &str) -> AutomationRequest { + AutomationRequest::Inspect(InspectRequest { + target: InspectTarget::Path { + segments: vec![target.to_owned()], + }, + depth: 1, + max_children: 100, + max_instances: 2_000, + include_properties: false, + include_attributes: true, + include_tags: true, + }) +} + +fn inspect_result(session_id: &str, target: &str) -> AutomationResult { + AutomationResult::Inspect(InspectResult { + root: InspectNode { + reference: InstanceReference { + session_id: session_id.to_owned(), + id: "pinst-00000001".to_owned(), + path: target.to_owned(), + name: target.to_owned(), + class_name: "Workspace".to_owned(), + }, + name: target.to_owned(), + class_name: "Workspace".to_owned(), + path: target.to_owned(), + properties: BTreeMap::new(), + attributes: BTreeMap::new(), + tags: Vec::new(), + children: Vec::new(), + truncated: false, + }, + visited_instances: 1, + truncated: false, + truncation_reason: None, + }) +} + +#[test] +fn typed_automation_jobs_enforce_fifo_session_and_limits() { + run_serve_test("empty", |session, _redactions| { + let server_info = session.get_api_rojo().unwrap(); + let plugin_id = Uuid::new_v4().to_string(); + let other_id = Uuid::new_v4().to_string(); + let heartbeat = AutomationHeartbeatRequest { + plugin_session_id: plugin_id.clone(), + server_session_id: server_info.session_id, + studio_mode: StudioMode::Edit, + exec_handler_available: true, + automation_handler_version: 2, + plugin_version: Some("test".to_owned()), + }; + assert_eq!( + session + .get_api_automation_next(&plugin_id, StudioMode::Edit) + .status(), + StatusCode::CONFLICT + ); + assert_eq!( + session.post_api_automation_heartbeat(&heartbeat).status(), + StatusCode::OK + ); + + let first_response = session.post_api_automation_job(&inspect_request("Workspace")); + assert_eq!(first_response.status(), StatusCode::CREATED); + let first: AutomationJobResponse = decode_response(first_response); + let second: AutomationJobResponse = + decode_response(session.post_api_automation_job(&inspect_request("Lighting"))); + assert_eq!(first.state, AutomationJobState::Pending); + + let counts: AutomationStatusResponse = decode_response(session.get_api_automation_status()); + assert_eq!(counts.queues.automation_pending, 2); + assert_eq!(counts.queues.automation_claimed, 0); + + let claim: AutomationJobClaimResponse = + decode_response(session.get_api_automation_next(&plugin_id, StudioMode::Edit)); + assert_eq!(claim.job_id, first.job_id); + assert_eq!(claim.state, AutomationJobState::Claimed); + assert_eq!(claim.execution_timeout_ms, 15_000); + let claimed_counts: AutomationStatusResponse = + decode_response(session.get_api_automation_status()); + assert_eq!(claimed_counts.queues.automation_claimed, 1); + assert_eq!( + claimed_counts + .queues + .automation_claimed_by_plugin_session_id + .as_deref(), + Some(plugin_id.as_str()) + ); + assert_eq!( + session + .get_api_automation_next(&plugin_id, StudioMode::Edit) + .status(), + StatusCode::NO_CONTENT + ); + assert_eq!( + session + .get_api_automation_next(&other_id, StudioMode::Edit) + .status(), + StatusCode::CONFLICT + ); + + let wrong_completion = AutomationJobCompletionRequest { + completion: AutomationJobCompletion::Success { + result: Box::new(inspect_result(&other_id, "Workspace")), + }, + plugin_session_id: other_id, + studio_mode: StudioMode::Edit, + }; + assert_eq!( + session + .post_api_automation_complete(&first.job_id, &wrong_completion) + .status(), + StatusCode::CONFLICT + ); + + let completion = AutomationJobCompletionRequest { + completion: AutomationJobCompletion::Success { + result: Box::new(inspect_result(&plugin_id, "Workspace")), + }, + plugin_session_id: plugin_id.clone(), + studio_mode: StudioMode::Edit, + }; + let completed: AutomationJobResponse = + decode_response(session.post_api_automation_complete(&first.job_id, &completion)); + assert_eq!(completed.state, AutomationJobState::Succeeded); + assert_eq!( + session + .post_api_automation_complete(&first.job_id, &completion) + .status(), + StatusCode::CONFLICT + ); + + let second_claim: AutomationJobClaimResponse = + decode_response(session.get_api_automation_next(&plugin_id, StudioMode::Edit)); + assert_eq!(second_claim.job_id, second.job_id); + let failed = AutomationJobCompletionRequest { + completion: AutomationJobCompletion::Failure { + error: "inspection failed".to_owned(), + }, + plugin_session_id: plugin_id.clone(), + studio_mode: StudioMode::Edit, + }; + assert_eq!( + session + .post_api_automation_complete(&second.job_id, &failed) + .status(), + StatusCode::OK + ); + let failed_status: AutomationJobResponse = + decode_response(session.get_api_automation_job_status(&second.job_id)); + assert_eq!(failed_status.state, AutomationJobState::Failed); + assert_eq!(failed_status.error.as_deref(), Some("inspection failed")); + assert_eq!( + session + .get_api_automation_next(&plugin_id, StudioMode::Edit) + .status(), + StatusCode::NO_CONTENT + ); + + let pending: AutomationJobResponse = + decode_response(session.post_api_automation_job(&inspect_request("Workspace"))); + assert_eq!( + session + .post_api_automation_complete(&pending.job_id, &failed) + .status(), + StatusCode::CONFLICT + ); + assert_eq!( + session + .api_request( + reqwest::Method::POST, + &format!("/api/automation/jobs/{}/complete", pending.job_id), + Some(vec![ + 0; + librojo::automation::MAX_AUTOMATION_RESULT_BODY_BYTES + + 1 + ]), + ) + .status(), + StatusCode::PAYLOAD_TOO_LARGE + ); + assert_eq!( + session.get_api_automation_job_status("not-a-uuid").status(), + StatusCode::BAD_REQUEST + ); + assert_eq!( + session + .api_request( + reqwest::Method::POST, + &format!("/api/automation/jobs/{}/complete", pending.job_id), + Some(vec![0xc1]), + ) + .status(), + StatusCode::BAD_REQUEST + ); + + for (method, path) in [ + (reqwest::Method::GET, "/api/automation/jobs/".to_owned()), + ( + reqwest::Method::GET, + "/api/automation/jobs/next/extra".to_owned(), + ), + ( + reqwest::Method::POST, + format!("/api/automation/jobs/{}/complete/extra", pending.job_id), + ), + ] { + assert_eq!( + session.api_request(method, &path, None).status(), + StatusCode::NOT_FOUND + ); + } + assert_eq!( + session + .get_api_automation_job_status(&Uuid::new_v4().to_string()) + .status(), + StatusCode::NOT_FOUND + ); + assert_eq!( + session + .api_request( + reqwest::Method::POST, + "/api/automation/jobs", + Some(vec![0; MAX_AUTOMATION_REQUEST_BODY_BYTES + 1]), + ) + .status(), + StatusCode::PAYLOAD_TOO_LARGE + ); + assert_eq!( + session + .api_request( + reqwest::Method::POST, + "/api/automation/jobs", + Some(vec![0xc1]), + ) + .status(), + StatusCode::BAD_REQUEST + ); + + for _ in 1..MAX_PENDING_AUTOMATION_JOBS { + assert_eq!( + session + .post_api_automation_job(&inspect_request("Workspace")) + .status(), + StatusCode::CREATED + ); + } + assert_eq!( + session + .post_api_automation_job(&inspect_request("Workspace")) + .status(), + StatusCode::TOO_MANY_REQUESTS + ); + }); +} + +#[test] +fn automation_status_tracks_plugin_identity_and_exec_queue() { + run_serve_test("empty", |session, _redactions| { + let server_info = session.get_api_rojo().unwrap(); + let initial: AutomationStatusResponse = + decode_response(session.get_api_automation_status()); + assert_eq!(initial.server_session_id, server_info.session_id); + assert_eq!(initial.protocol_version, server_info.protocol_version); + assert_eq!(initial.automation_handler_version, 2); + assert!(!initial.automation_available); + assert!(!initial.exec_available); + assert!(!initial.typed_automation_available); + assert_eq!(initial.plugin, None); + assert_eq!(initial.queues.exec_pending, 0); + assert_eq!(initial.queues.exec_claimed, 0); + + let first_plugin_session_id = Uuid::new_v4().to_string(); + let heartbeat = AutomationHeartbeatRequest { + plugin_session_id: first_plugin_session_id.clone(), + server_session_id: server_info.session_id, + studio_mode: StudioMode::Edit, + exec_handler_available: true, + automation_handler_version: 2, + plugin_version: Some("7.7.0-test".to_owned()), + }; + let registered: AutomationHeartbeatResponse = + decode_response(session.post_api_automation_heartbeat(&heartbeat)); + assert_eq!(registered.registration, AutomationRegistration::Registered); + + let mut refreshed_heartbeat = heartbeat.clone(); + refreshed_heartbeat.studio_mode = StudioMode::Play; + let refreshed: AutomationHeartbeatResponse = + decode_response(session.post_api_automation_heartbeat(&refreshed_heartbeat)); + assert_eq!(refreshed.registration, AutomationRegistration::Refreshed); + + let pending = submit_exec_job(&session, "status.lua", "return true"); + let pending_status: AutomationStatusResponse = + decode_response(session.get_api_automation_status()); + assert_eq!(pending_status.queues.exec_pending, 1); + + let claim: ExecJobClaimResponse = decode_response( + session.get_api_exec_next_for_session(&first_plugin_session_id, StudioMode::Edit), + ); + assert_eq!(claim.job_id, pending.job_id); + let connected: AutomationStatusResponse = + decode_response(session.get_api_automation_status()); + let plugin = connected.plugin.unwrap(); + assert!(plugin.connected); + assert_eq!(plugin.plugin_session_id, first_plugin_session_id); + assert_eq!(plugin.studio_mode, StudioMode::Edit); + assert_eq!(plugin.plugin_version.as_deref(), Some("7.7.0-test")); + assert!(connected.automation_available); + assert!(connected.exec_available); + assert!(connected.typed_automation_available); + assert_eq!(connected.queues.exec_pending, 0); + assert_eq!(connected.queues.exec_claimed, 1); + assert_eq!( + connected + .queues + .exec_claimed_by_plugin_session_id + .as_deref(), + Some(first_plugin_session_id.as_str()) + ); + + let completion = ExecJobCompletionEnvelope { + completion: ExecJobCompletionRequest::Success { + result: Some(ExecValue::Boolean { value: true }), + logs: Vec::new(), + }, + plugin_session_id: Some(first_plugin_session_id.clone()), + studio_mode: Some(StudioMode::Edit), + }; + assert_eq!( + session + .post_api_exec_complete_for_session(&pending.job_id, &completion) + .status(), + StatusCode::OK + ); + let completed_status: AutomationStatusResponse = + decode_response(session.get_api_automation_status()); + assert_eq!(completed_status.queues.exec_claimed, 0); + assert_eq!( + completed_status.queues.exec_claimed_by_plugin_session_id, + None + ); + + let second_heartbeat = AutomationHeartbeatRequest { + plugin_session_id: Uuid::new_v4().to_string(), + ..heartbeat + }; + let conflict: AutomationHeartbeatResponse = + decode_response(session.post_api_automation_heartbeat(&second_heartbeat)); + assert_eq!(conflict.registration, AutomationRegistration::Conflict); + let duplicate: AutomationStatusResponse = + decode_response(session.get_api_automation_status()); + assert!(duplicate.duplicate_session_detected); + + let malformed = AutomationHeartbeatRequest { + plugin_session_id: "not-a-uuid".to_owned(), + ..second_heartbeat + }; + assert_eq!( + session.post_api_automation_heartbeat(&malformed).status(), + StatusCode::BAD_REQUEST + ); + assert_eq!( + session + .api_request( + reqwest::Method::GET, + "/api/exec/jobs/next?pluginSessionId=not-a-uuid&studioMode=edit", + None, + ) + .status(), + StatusCode::BAD_REQUEST + ); + }); +} + #[test] fn exec_jobs_submit_claim_complete_and_release_slot() { run_serve_test("empty", |session, _redactions| {