Report the result of a connection attempt to ConnectAsync's caller (AI-assisted)

CHANGELOG.md

@@ -32,6 +32,8 @@ Making a new release? Simply add the new header with the version and date undern
 ## Unreleased
 
 * Added headless API for Studio companion plugins. ([#639])
+* `Rojo:ConnectAsync` now returns whether the session was established, along with the reason it was not. ([#639])
+* Granted headless API permissions are no longer forgotten when Studio restarts, since a plugin source is no longer stored as a table key. ([#639])
 * Fixed `$path` values that point outside the project folder failing to match `syncRule`s on Windows, which broke `rojo sourcemap` with a "could not be turned into a Roblox Instance" error. ([#1290])
 * Fixed `rojo serve` silently stopping syncing file changes on Windows when the served project path was a verbatim (`\\?\`) path, because tree paths and file-watcher event paths were canonicalized to different forms. ([#1290])
 * Fixed `rojo sourcemap --absolute` emitting verbatim (`\\?\`) paths on Windows, which broke require types in luau-lsp. ([#1290])

plugin/src/App/init.lua

@@ -759,6 +759,8 @@ function App:startSession(host: string?, port: string?)
 			toolbarIcon = Assets.Images.PluginButtonWarning,
 		})
 
+		self.headlessAPI:_settleConnectAttempt(false, msg)
+
 		return
 	end
 
@@ -891,6 +893,17 @@ function App:startSession(host: string?, port: string?)
 			self.headlessAPI:_updateProperty("Address", nil)
 			self.headlessAPI:_updateProperty("ProjectName", nil)
 		end
+
+		-- Settled last so that a caller waiting on ConnectAsync sees the API
+		-- properties already describing the session it asked for.
+		if status == ServeSession.Status.Connected then
+			self.headlessAPI:_settleConnectAttempt(true)
+		elseif status == ServeSession.Status.Disconnected then
+			self.headlessAPI:_settleConnectAttempt(
+				false,
+				if details ~= nil then tostring(details) else "Disconnected from session"
+			)
+		end
 	end)
 
 	serveSession:setConfirmCallback(function(instanceMap, patch, serverInfo)

plugin/src/HeadlessAPI.lua

@@ -17,6 +17,11 @@ local apiPermissionAllowlist = {
 	RequestAccess = true,
 }
 
+-- Code that has no plugin of its own to identify it in a traceback is
+-- attributed to this source. The command bar is the usual case, but so is any
+-- chunk a plugin compiles at runtime, so this source is not trusted by itself.
+local COMMAND_BAR_SOURCE = "RobloxStudio_CommandBar"
+
 export type CallerInfo = {
 	Source: string,
 	Type: "Local" | "Cloud" | "Studio",
@@ -31,12 +36,57 @@ export type CallerInfo = {
 
 local API = {}
 
+-- plugin:SetSetting stores tables as JSON and mangles any '.' in a key into
+-- '_'. Sources are plugin file names like 'user_Companion.rbxmx', so a source
+-- is never used as a key on disk: permissions are stored as a list of records
+-- and only turned into a map in memory.
+local function loadPermissions(): { [string]: { [string]: boolean } }
+	local permissions = {}
+	local stored = Settings:get("apiPermissions")
+	if type(stored) ~= "table" then
+		return permissions
+	end
+
+	for _, record in ipairs(stored) do
+		if type(record) == "table" and type(record.source) == "string" and type(record.apis) == "table" then
+			local apis = {}
+			for _, api in record.apis do
+				if type(api) == "string" then
+					apis[api] = true
+				end
+			end
+
+			permissions[record.source] = apis
+		end
+	end
+
+	return permissions
+end
+
+local function savePermissions(permissions: { [string]: { [string]: boolean } })
+	local stored = {}
+	for source, apis in permissions do
+		local list = {}
+		for api in apis do
+			table.insert(list, api)
+		end
+		table.sort(list)
+
+		table.insert(stored, { source = source, apis = list })
+	end
+	table.sort(stored, function(a, b)
+		return a.source < b.source
+	end)
+
+	Settings:set("apiPermissions", stored)
+end
+
 function API.new(app)
 	local Rojo = {}
 
 	Rojo._rateLimit = {}
 	Rojo._sourceToPlugin = {}
-	Rojo._permissions = Settings:get("apiPermissions") or {}
+	Rojo._permissions = loadPermissions()
 	Rojo._activePermissionRequests = {}
 	Rojo._changedEvent = Instance.new("BindableEvent")
 	Rojo._apiDescriptions = {}
@@ -97,7 +147,7 @@ function API.new(app)
 			return cloudPlugin
 		end
 
-		return "RobloxStudio_CommandBar"
+		return COMMAND_BAR_SOURCE
 	end
 
 	function Rojo:_getCallerName()
@@ -273,7 +323,7 @@ function API.new(app)
 
 		-- Update stored permissions
 		Rojo._permissions[source] = sourcePermissions
-		Settings:set("apiPermissions", Rojo._permissions)
+		savePermissions(Rojo._permissions)
 
 		-- Share changes
 		Rojo._permissionsChangedEvent:Fire(source, sourcePermissions)
@@ -284,7 +334,7 @@ function API.new(app)
 		Log.info(string.format("Denying access to Rojo APIs for '%s'", name))
 
 		-- Update stored permissions
-		Settings:set("apiPermissions", Rojo._permissions)
+		savePermissions(Rojo._permissions)
 
 		-- Share changes
 		Rojo._permissionsChangedEvent:Fire(source, nil)
@@ -384,19 +434,54 @@ function API.new(app)
 		)
 	end
 
+	-- Threads waiting on the outcome of a connection attempt they started.
+	Rojo._connectWaiters = {}
+
+	function Rojo:_settleConnectAttempt(success: boolean, message: string?)
+		local waiters = Rojo._connectWaiters
+		if #waiters == 0 then
+			return
+		end
+
+		Rojo._connectWaiters = {}
+		for _, waiter in waiters do
+			task.spawn(waiter, success, message)
+		end
+	end
+
 	Rojo._apiDescriptions.ConnectAsync = {
 		Type = "Method",
-		Description = "Connects to a Rojo server",
+		Description = "Connects to a Rojo server and returns whether the session was established",
 	}
-	function Rojo:ConnectAsync(host: string?, port: string?)
+	function Rojo:ConnectAsync(host: string?, port: string?): (boolean, string?)
 		assert(type(host) == "string" or host == nil, "Host must be type `string?`")
 		assert(type(port) == "string" or port == nil, "Port must be type `string?`")
 
 		if Rojo:_checkRateLimit("ConnectAsync") then
-			return
+			return false, "Rojo:ConnectAsync is being rate limited"
 		end
 
+		-- The attempt can settle before startSession returns, as it does when the
+		-- sync lock is already held, so the waiter is registered beforehand and
+		-- the thread only yields if it is still pending afterwards.
+		local thread = coroutine.running()
+		local settled, success, message = false, false, nil :: string?
+
+		table.insert(Rojo._connectWaiters, function(attemptSuccess: boolean, attemptMessage: string?)
+			settled, success, message = true, attemptSuccess, attemptMessage
+
+			if coroutine.status(thread) == "suspended" then
+				task.spawn(thread)
+			end
+		end)
+
 		app:startSession(host, port)
+
+		if not settled then
+			coroutine.yield()
+		end
+
+		return success, message
 	end
 
 	Rojo._apiDescriptions.DisconnectAsync = {

plugin/src/HeadlessAPI.spec.lua

@@ -0,0 +1,134 @@
+return function()
+	local HeadlessAPI = require(script.Parent.HeadlessAPI)
+	local Settings = require(script.Parent.Settings)
+
+	-- The API only reaches back into the app to start and stop sessions, so a
+	-- stub is enough to observe what a caller is told about the attempt.
+	local function newStubApp()
+		local stub = {
+			serveSession = nil,
+			startedWith = nil,
+		}
+
+		function stub:startSession(host, port)
+			stub.startedWith = { host = host, port = port }
+
+			if stub.onStartSession then
+				stub.onStartSession()
+			end
+		end
+
+		function stub:endSession() end
+
+		return stub
+	end
+
+	it("should tell the caller when a session connects", function()
+		local app = newStubApp()
+		local api = HeadlessAPI.new(app)
+
+		app.onStartSession = function()
+			api:_settleConnectAttempt(true)
+		end
+
+		local success, message = api:ConnectAsync("localhost", "34872")
+
+		expect(success).to.equal(true)
+		expect(message).to.equal(nil)
+		expect(app.startedWith.host).to.equal("localhost")
+		expect(app.startedWith.port).to.equal("34872")
+	end)
+
+	it("should tell the caller when a session fails before it starts", function()
+		local app = newStubApp()
+		local api = HeadlessAPI.new(app)
+
+		-- Failures such as a held sync lock settle inside startSession, before
+		-- the calling thread ever has a chance to yield.
+		app.onStartSession = function()
+			api:_settleConnectAttempt(false, "Could not sync because user 'someone' is already syncing")
+		end
+
+		local success, message = api:ConnectAsync()
+
+		expect(success).to.equal(false)
+		expect(message).to.equal("Could not sync because user 'someone' is already syncing")
+	end)
+
+	it("should wait for a session that settles later", function()
+		local app = newStubApp()
+		local api = HeadlessAPI.new(app)
+
+		app.onStartSession = function()
+			task.delay(0.05, function()
+				api:_settleConnectAttempt(false, "Connection refused")
+			end)
+		end
+
+		local success, message = api:ConnectAsync()
+
+		expect(success).to.equal(false)
+		expect(message).to.equal("Connection refused")
+	end)
+
+	it("should only settle the attempts that are still waiting", function()
+		local app = newStubApp()
+		local api = HeadlessAPI.new(app)
+
+		app.onStartSession = function()
+			api:_settleConnectAttempt(true)
+		end
+
+		expect(api:ConnectAsync()).to.equal(true)
+
+		-- A later disconnect has nobody to report to and must not resume the
+		-- thread that already got its answer.
+		api:_settleConnectAttempt(false, "Disconnected from session")
+
+		expect(api:ConnectAsync()).to.equal(true)
+	end)
+
+	-- Which source a call comes from is read off the traceback, and the test
+	-- runner is itself a plugin, so the source is stubbed to describe the caller
+	-- the test is about.
+	local function newApiAsSource(source: string)
+		local api, readOnlyApi = HeadlessAPI.new(newStubApp())
+
+		function api:_getCallerSource()
+			return source
+		end
+
+		return api, readOnlyApi
+	end
+
+	it("should refuse callers that have not been granted access", function()
+		local _, readOnlyApi = newApiAsSource("user_Impostor.rbxmx")
+
+		expect(function()
+			return readOnlyApi.ConnectAsync
+		end).to.throw()
+
+		-- Reading the version is how a caller checks compatibility, so it never
+		-- needs a grant.
+		expect(readOnlyApi.Version).to.be.ok()
+	end)
+
+	it("should store permissions without putting a source in a table key", function()
+		local api = newApiAsSource("user_Companion.rbxmx")
+		api:_setPermissions("user_Companion.rbxmx", "Companion", { "ConnectAsync" })
+
+		-- plugin:SetSetting rewrites '.' in keys, which would make the stored
+		-- grant unmatchable the next time Studio opens.
+		local stored = Settings:get("apiPermissions")
+		expect(#stored).to.equal(1)
+		expect(stored[1].source).to.equal("user_Companion.rbxmx")
+		expect(stored[1].apis[1]).to.equal("ConnectAsync")
+
+		-- A fresh API reads that back, which is what a new Studio session does.
+		local _, readOnlyApi = newApiAsSource("user_Companion.rbxmx")
+		expect(readOnlyApi.ConnectAsync).to.be.a("function")
+
+		api:_removePermissions("user_Companion.rbxmx", "Companion")
+		expect(#Settings:get("apiPermissions")).to.equal(0)
+	end)
+end