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

plugin/src/App/init.lua

@@ -743,7 +743,25 @@ function App:useRunningConnectionInfo()
 	self.setPort(port)
 end
 
-function App:startSession(host: string?, port: string?)
+function App:startSession(host: string?, port: string?, onSettled: ((boolean, string?) -> ())?)
+	-- A caller that waits on this session is answered exactly once, and only
+	-- about the session it started.
+	local report = onSettled
+	local function settle(success: boolean, message: string?)
+		if report == nil then
+			return
+		end
+
+		local reportSettled = report
+		report = nil
+		reportSettled(success, message)
+	end
+
+	if self.serveSession ~= nil then
+		settle(false, "A sync session is already running")
+		return
+	end
+
 	local claimedLock, priorOwner = self:claimSyncLock()
 	if not claimedLock then
 		local msg = string.format("Could not sync because user '%s' is already syncing", tostring(priorOwner))
@@ -759,6 +777,8 @@ function App:startSession(host: string?, port: string?)
 			toolbarIcon = Assets.Images.PluginButtonWarning,
 		})
 
+		settle(false, msg)
+
 		return
 	end
 
@@ -891,6 +911,14 @@ 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 this session sees the API
+		-- properties already describing it.
+		if status == ServeSession.Status.Connected then
+			settle(true)
+		elseif status == ServeSession.Status.Disconnected then
+			settle(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,10 @@ local apiPermissionAllowlist = {
 	RequestAccess = true,
 }
 
+-- Callers with no plugin of their own in the traceback are attributed to this
+-- source: the command bar, and any chunk a plugin compiles at runtime.
+local COMMAND_BAR_SOURCE = "RobloxStudio_CommandBar"
+
 export type CallerInfo = {
 	Source: string,
 	Type: "Local" | "Cloud" | "Studio",
@@ -31,12 +35,58 @@ export type CallerInfo = {
 
 local API = {}
 
+-- Permissions are stored as a list of records rather than a map because
+-- source strings contain '.' and cannot be used as JSON keys.
+local function loadPermissions(): { [string]: { [string]: boolean } }
+	local permissions = {}
+	local stored = Settings:get("apiPermissions")
+	if type(stored) ~= "table" then
+		return permissions
+	end
+
+	for _, record in stored do
+		if type(record) ~= "table" or type(record.source) ~= "string" or type(record.apis) ~= "table" then
+			continue
+		end
+
+		local apis = {}
+		for _, api in record.apis do
+			if type(api) ~= "string" then
+				continue
+			end
+			apis[api] = true
+		end
+
+		permissions[record.source] = apis
+	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)
@@ -386,17 +436,34 @@ function API.new(app)
 
 	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
 
-		app:startSession(host, port)
+		-- The session can settle before startSession returns, as it does when one
+		-- is already running, so the thread yields only while still waiting.
+		local thread = coroutine.running()
+		local settled, success, message = false, false, nil :: string?
+
+		app:startSession(host, port, function(attemptSuccess: boolean, attemptMessage: string?)
+			settled, success, message = true, attemptSuccess, attemptMessage
+
+			if coroutine.status(thread) == "suspended" then
+				task.spawn(thread)
+			end
+		end)
+
+		while not settled do
+			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 starts and stops sessions, so a stub app is enough here.
+	-- onStartSession decides how the session this stub was asked for settles.
+	local function newStubApp()
+		local stub = { settleCount = 0 }
+
+		function stub:startSession(host, port, onSettled)
+			self.startedWith = { host = host, port = port }
+			self.settle = function(success, message)
+				self.settleCount += 1
+				onSettled(success, message)
+			end
+
+			if self.onStartSession then
+				self.onStartSession()
+			end
+		end
+
+		function stub:endSession() end
+
+		return stub
+	end
+
+	-- A caller's source is read off the traceback, and the test runner is itself
+	-- a plugin, so it is stubbed to name the caller each 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 tell the caller when a session connects", function()
+		local app = newStubApp()
+		local api = HeadlessAPI.new(app)
+
+		app.onStartSession = function()
+			app.settle(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)
+
+		-- A held sync lock settles inside startSession, before the caller yields.
+		app.onStartSession = function()
+			app.settle(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()
+				app.settle(false, "Connection refused")
+			end)
+		end
+
+		local success, message = api:ConnectAsync()
+
+		expect(success).to.equal(false)
+		expect(message).to.equal("Connection refused")
+	end)
+
+	it("should answer a caller once, even if its session settles again", function()
+		local app = newStubApp()
+		local api = HeadlessAPI.new(app)
+
+		app.onStartSession = function()
+			app.settle(true)
+		end
+
+		expect(api:ConnectAsync()).to.equal(true)
+		expect(app.settleCount).to.equal(1)
+
+		-- The session reporting again, as it does when it later disconnects, must
+		-- not reach a caller that already has its answer.
+		app.settle(false, "Disconnected from session")
+
+		expect(app.settleCount).to.equal(2)
+		expect(api:ConnectAsync()).to.equal(true)
+	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()
+
+		expect(readOnlyApi.Version).to.be.ok()
+	end)
+
+	it("should store permissions without putting a source in a table key", function()
+		local priorPermissions = Settings:get("apiPermissions")
+		local api = newApiAsSource("user_Companion.rbxmx")
+
+		api:_setPermissions("user_Companion.rbxmx", "Companion", { "ConnectAsync" })
+
+		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")
+
+		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)
+
+		Settings:set("apiPermissions", priorPermissions)
+	end)
+end