dominus-cli 0.6.0 → 2.1.0

package/plugin/src/init.server.lua

@@ -1,412 +1,225 @@
---[[
-	Dominus - Roblox Studio Plugin
-	AI-powered development agent bridge via WebSocket
+local HttpService = game:GetService("HttpService")
 
-	This plugin connects to the Dominus CLI running on your machine,
-	enabling real-time bidirectional communication for AI-assisted development.
-]]
-
-local toolbar = plugin:CreateToolbar("Dominus")
-local connectButton = toolbar:CreateButton("Connect", "Connect to Dominus CLI agent", "rbxassetid://12974519994")
-
-local WsClient = require(script.WsClient)
+local CommandRouter = require(script.CommandRouter)
 local Protocol = require(script.Protocol)
-local Explorer = require(script.Explorer)
-local Executor = require(script.Executor)
 local Watcher = require(script.Watcher)
-local Properties = require(script.Properties)
-local Reflection = require(script.Reflection)
-local TestRunner = require(script.TestRunner)
-local UIBuilder = require(script.UIBuilder)
+local WsClient = require(script.WsClient)
+
+local PROTOCOL_VERSION = 2
+local DEFAULT_PORT = 18088
+local MAX_MESSAGE_BYTES = 1024 * 1024
+
+local toolbar = plugin:CreateToolbar("Dominus 2")
+local connectButton = toolbar:CreateButton(
+	"Dominus 2",
+	"Connect this Studio session to the Dominus 2 MCP server",
+	"rbxassetid://12974519994"
+)
+
+local authenticated = false
+local transportOpen = false
+local shouldReconnect = true
+local reconnectScheduled = false
+local reconnectDelay = 1
+local shuttingDown = false
+
+local function getSetting(name, default)
+	local ok, value = pcall(function()
+		return plugin:GetSetting(name)
+	end)
+	if ok and value ~= nil then
+		return value
+	end
+	return default
+end
+
+local clientId = getSetting("Dominus2ClientId", nil)
+if type(clientId) ~= "string" then
+	clientId = HttpService:GenerateGUID(false)
+	pcall(function()
+		plugin:SetSetting("Dominus2ClientId", clientId)
+	end)
+end
+local bridgeToken = getSetting("Dominus2BridgeToken", nil)
+local port = tonumber(getSetting("Dominus2Port", DEFAULT_PORT)) or DEFAULT_PORT
 
-local BASE_PORT = 18088
-local MAX_PORT_SCAN = 10
-local activePort = BASE_PORT
-local connected = false
-local ws = nil
+local function log(message)
+	print("[Dominus 2] " .. message)
+end
+
+local function sendEnvelope(message)
+	local encodedOk, encoded = pcall(Protocol.encode, message)
+	if not encodedOk then
+		return false, tostring(encoded)
+	end
+	if #encoded > MAX_MESSAGE_BYTES then
+		return false, "Message exceeds one MiB"
+	end
+	return WsClient.send(encoded)
+end
 
-local function log(msg)
-	print("[Dominus]", msg)
+local function sendResponse(id, payload)
+	return sendEnvelope({ id = id, type = "studio:response", payload = payload, ts = os.time() })
+end
+
+local scheduleReconnect
+
+local function stopSession()
+	authenticated = false
+	transportOpen = false
+	Watcher.stop()
+	connectButton:SetActive(false)
 end
 
 local function handleMessage(rawMessage)
-	local ok, msg = pcall(Protocol.decode, rawMessage)
-	if not ok then
-		warn("[Dominus] Failed to decode message:", msg)
+	if #rawMessage > MAX_MESSAGE_BYTES then
+		warn("[Dominus 2] Rejected an oversized server message")
+		WsClient.disconnect()
+		return
+	end
+	local decodedOk, message = pcall(Protocol.decode, rawMessage)
+	if not decodedOk or type(message) ~= "table" then
+		warn("[Dominus 2] Rejected an invalid server message")
 		return
 	end
 
-	local msgType = msg.type
-	local payload = msg.payload
-	local msgId = msg.id
+	if message.type == "server:pairing_required" then
+		local payload = message.payload or {}
+		log(
+			"Pairing required. In your MCP client call dominus_pair_studio with connectionId "
+				.. tostring(payload.connectionId)
+				.. " and code "
+				.. tostring(payload.pairingCode)
+		)
+		return
+	end
 
-	-- Route message to appropriate handler
-	if msgType == "cli:get:explorer" then
-		local tree = Explorer.getTree(payload.rootPath, payload.maxDepth or 3)
-		WsClient.send(Protocol.encode({
-			id = msgId,
-			type = "studio:response",
-			payload = { tree = tree },
-			ts = os.clock(),
-		}))
-	elseif msgType == "cli:get:script" then
-		local result = Explorer.getScriptSource(payload.path)
-		WsClient.send(Protocol.encode({
-			id = msgId,
-			type = "studio:response",
-			payload = result,
-			ts = os.clock(),
-		}))
-	elseif msgType == "cli:set:script" then
-		local result = Explorer.setScriptSource(payload.path, payload.source)
-		WsClient.send(Protocol.encode({
-			id = msgId,
-			type = "studio:response",
-			payload = result,
-			ts = os.clock(),
-		}))
-	elseif msgType == "cli:run" then
-		local result = Executor.run(payload.code)
-		WsClient.send(Protocol.encode({
-			id = msgId,
-			type = "studio:response",
-			payload = result,
-			ts = os.clock(),
-		}))
-	elseif msgType == "cli:get:properties" then
-		local compact = payload.compact
-		if compact == nil then
-			compact = true
+	if message.type == "server:authenticated" then
+		local payload = message.payload or {}
+		if payload.protocolVersion ~= PROTOCOL_VERSION or type(payload.connectionId) ~= "string" then
+			warn("[Dominus 2] Server returned an invalid authentication response")
+			WsClient.disconnect()
+			return
 		end
-		local result = Properties.get(payload.path, compact)
-		WsClient.send(Protocol.encode({
-			id = msgId,
-			type = "studio:response",
-			payload = result,
-			ts = os.clock(),
-		}))
-	elseif msgType == "cli:get:descendants:properties" then
-		local compact = payload.compact
-		if compact == nil then
-			compact = true
-		end -- default to compact
-		local result = Properties.getDescendants(payload.path, compact)
-		WsClient.send(Protocol.encode({
-			id = msgId,
-			type = "studio:response",
-			payload = result,
-			ts = os.clock(),
-		}))
-	elseif msgType == "cli:set:properties" then
-		local result = Properties.set(payload.path, payload.properties)
-		WsClient.send(Protocol.encode({
-			id = msgId,
-			type = "studio:response",
-			payload = result,
-			ts = os.clock(),
-		}))
-	elseif msgType == "cli:insert" then
-		local result = Explorer.insertInstance(payload.className, payload.parentPath, payload.name, payload.properties)
-		WsClient.send(Protocol.encode({
-			id = msgId,
-			type = "studio:response",
-			payload = result,
-			ts = os.clock(),
-		}))
-	elseif msgType == "cli:delete" then
-		local result = Explorer.deleteInstance(payload.path)
-		WsClient.send(Protocol.encode({
-			id = msgId,
-			type = "studio:response",
-			payload = result,
-			ts = os.clock(),
-		}))
-	elseif msgType == "cli:get:selection" then
-		local Selection = game:GetService("Selection")
-		local selected = Selection:Get()
-		local paths = {}
-		local classNames = {}
-		for _, obj in selected do
-			table.insert(paths, Explorer.getPath(obj))
-			table.insert(classNames, obj.ClassName)
+		if type(payload.token) == "string" then
+			bridgeToken = payload.token
+			pcall(function()
+				plugin:SetSetting("Dominus2BridgeToken", bridgeToken)
+			end)
 		end
-		WsClient.send(Protocol.encode({
-			id = msgId,
-			type = "studio:response",
-			payload = { paths = paths, classNames = classNames },
-			ts = os.clock(),
-		}))
-	elseif msgType == "cli:search:scripts" then
-		local result = Explorer.searchScripts(payload.query, payload.maxResults or 20)
-		WsClient.send(Protocol.encode({
-			id = msgId,
-			type = "studio:response",
-			payload = result,
-			ts = os.clock(),
-		}))
-	elseif msgType == "cli:get:output" then
-		local result = Watcher.getRecentOutput(payload.limit or 50, payload.level)
-		WsClient.send(Protocol.encode({
-			id = msgId,
-			type = "studio:response",
-			payload = result,
-			ts = os.clock(),
-		}))
-	elseif msgType == "cli:run:tests" then
-		local result = TestRunner.run(payload.testName)
-		WsClient.send(Protocol.encode({
-			id = msgId,
-			type = "studio:response",
-			payload = result,
-			ts = os.clock(),
-		}))
-	elseif msgType == "cli:test:run" then
-		-- ExecuteRunModeAsync yields until EndTest is called, so run in a thread
-		task.spawn(function()
-			local result = TestRunner.executeRunMode(payload.args)
-			WsClient.send(Protocol.encode({
-				id = msgId,
-				type = "studio:response",
-				payload = result,
-				ts = os.clock(),
-			}))
-		end)
-	elseif msgType == "cli:test:play" then
-		-- ExecutePlayModeAsync yields until EndTest is called, so run in a thread
-		task.spawn(function()
-			local result = TestRunner.executePlayMode(payload.args)
-			WsClient.send(Protocol.encode({
-				id = msgId,
-				type = "studio:response",
-				payload = result,
-				ts = os.clock(),
-			}))
-		end)
-	elseif msgType == "cli:test:end" then
-		local result = TestRunner.endTest(payload.value)
-		WsClient.send(Protocol.encode({
-			id = msgId,
-			type = "studio:response",
-			payload = result,
-			ts = os.clock(),
-		}))
-	elseif msgType == "cli:build:ui" then
-		local result = UIBuilder.build(payload)
-		WsClient.send(Protocol.encode({
-			id = msgId,
-			type = "studio:response",
-			payload = result,
-			ts = os.clock(),
-		}))
-	elseif msgType == "cli:get:reflection" then
-		local result = Reflection.getClassInfo(payload.className)
-		WsClient.send(Protocol.encode({
-			id = msgId,
-			type = "studio:response",
-			payload = result,
-			ts = os.clock(),
-		}))
-	elseif msgType == "cli:serialize:ui" then
-		local result = UIBuilder.serialize(payload)
-		WsClient.send(Protocol.encode({
-			id = msgId,
-			type = "studio:response",
-			payload = result,
-			ts = os.clock(),
-		}))
-	elseif msgType == "cli:move:instance" then
-		local instance = Explorer.resolveInstance(payload.path)
-		if not instance then
-			WsClient.send(Protocol.encode({
-				id = msgId,
-				type = "studio:response",
-				payload = { success = false, error = "Instance not found: " .. payload.path },
-				ts = os.clock(),
-			}))
-		else
-			local newParent = Explorer.resolveInstance(payload.newParent)
-			if not newParent then
-				local svcOk, svc = pcall(function()
-					return game:GetService(payload.newParent)
-				end)
-				if svcOk then
-					newParent = svc
-				end
-			end
-			if not newParent then
-				WsClient.send(Protocol.encode({
-					id = msgId,
-					type = "studio:response",
-					payload = { success = false, error = "New parent not found: " .. payload.newParent },
-					ts = os.clock(),
-				}))
-			else
-				local ChangeHistoryService = game:GetService("ChangeHistoryService")
-				local recording = ChangeHistoryService:TryBeginRecording("Dominus: Move " .. instance.Name)
-				local ok, err = pcall(function()
-					instance.Parent = newParent
-				end)
-				if recording then
-					if ok then
-						ChangeHistoryService:FinishRecording(recording, Enum.FinishRecordingOperation.Commit)
-					else
-						ChangeHistoryService:FinishRecording(recording, Enum.FinishRecordingOperation.Cancel)
-					end
-				end
-				if ok then
-					WsClient.send(Protocol.encode({
-						id = msgId,
-						type = "studio:response",
-						payload = { success = true, newPath = Explorer.getPath(instance) },
-						ts = os.clock(),
-					}))
-				else
-					WsClient.send(Protocol.encode({
-						id = msgId,
-						type = "studio:response",
-						payload = { success = false, error = "Move failed: " .. tostring(err) },
-						ts = os.clock(),
-					}))
-				end
+		authenticated = true
+		connectButton:SetActive(true)
+		log("Authenticated as Studio session " .. payload.connectionId)
+		Watcher.start(function(entry)
+			if authenticated then
+				sendEnvelope({ id = Protocol.generateId(), type = "studio:output", payload = entry, ts = os.time() })
 			end
-		end
-	elseif msgType == "cli:undo" then
-		local ChangeHistoryService = game:GetService("ChangeHistoryService")
-		local ok, err = pcall(function()
-			ChangeHistoryService:Undo()
-		end)
-		WsClient.send(Protocol.encode({
-			id = msgId,
-			type = "studio:response",
-			payload = ok and { success = true } or { success = false, error = tostring(err) },
-			ts = os.clock(),
-		}))
-	elseif msgType == "cli:redo" then
-		local ChangeHistoryService = game:GetService("ChangeHistoryService")
-		local ok, err = pcall(function()
-			ChangeHistoryService:Redo()
 		end)
-		WsClient.send(Protocol.encode({
-			id = msgId,
-			type = "studio:response",
-			payload = ok and { success = true } or { success = false, error = tostring(err) },
-			ts = os.clock(),
-		}))
-	elseif msgType == "cli:bulk:set:properties" then
-		local result = Properties.bulkSet(payload)
-		WsClient.send(Protocol.encode({
-			id = msgId,
-			type = "studio:response",
-			payload = result,
-			ts = os.clock(),
-		}))
-	elseif msgType == "cli:find:replace:scripts" then
-		local result = Explorer.findReplaceScripts(payload)
-		WsClient.send(Protocol.encode({
-			id = msgId,
-			type = "studio:response",
-			payload = result,
-			ts = os.clock(),
-		}))
-	else
-		warn("[Dominus] Unknown message type:", msgType)
+		return
 	end
-end
 
-local function tryConnectToPort(port)
-	log("Trying ws://localhost:" .. port .. "...")
+	if message.type == "server:auth_rejected" then
+		warn("[Dominus 2] Authentication rejected: " .. tostring(message.payload and message.payload.error))
+		WsClient.disconnect()
+		return
+	end
 
-	local success = WsClient.connect("ws://localhost:" .. port, {
+	if not authenticated then
+		warn("[Dominus 2] Ignored a command before authentication")
+		return
+	end
+	CommandRouter.handle(message, sendResponse)
+end
+
+local function connect()
+	if shuttingDown or transportOpen then
+		return
+	end
+	reconnectScheduled = false
+	local success = WsClient.connect("ws://127.0.0.1:" .. port, {
 		onOpen = function()
-			connected = true
-			activePort = port
+			transportOpen = true
+			reconnectDelay = 1
 			connectButton:SetActive(true)
-			log("Connected to Dominus on port " .. port .. "!")
-
-			local studioVersion = game:GetService("RobloxPluginGuiService") and "Studio" or "Unknown"
-			local placeName = "Unknown"
-			pcall(function()
-				placeName = game:GetService("MarketplaceService"):GetProductInfo(game.PlaceId).Name
-			end)
-
-			WsClient.send(Protocol.encode({
+			local sent, sendErr = sendEnvelope({
 				id = Protocol.generateId(),
-				type = "studio:connected",
+				type = "studio:hello",
 				payload = {
-					studioVersion = studioVersion,
+					protocolVersion = PROTOCOL_VERSION,
+					clientId = clientId,
+					token = bridgeToken,
+					studioVersion = version(),
 					placeId = game.PlaceId,
-					placeName = placeName,
+					placeName = game.Name,
 				},
-				ts = os.clock(),
-			}))
-
-			Watcher.start(function(entry)
-				if connected then
-					WsClient.send(Protocol.encode({
-						id = Protocol.generateId(),
-						type = "studio:output",
-						payload = entry,
-						ts = os.clock(),
-					}))
-				end
-			end)
+				ts = os.time(),
+			})
+			if not sent then
+				warn("[Dominus 2] Could not send handshake: " .. tostring(sendErr))
+			end
 		end,
-
 		onMessage = function(message)
 			task.spawn(handleMessage, message)
 		end,
-
 		onClose = function()
-			connected = false
-			connectButton:SetActive(false)
-			Watcher.stop()
-			log("Disconnected from Dominus CLI")
+			local wasOpen = transportOpen
+			stopSession()
+			if wasOpen and not shuttingDown then
+				log("Studio bridge disconnected")
+			end
+			if shouldReconnect then
+				scheduleReconnect()
+			end
 		end,
-
-		onError = function(err)
-			warn("[Dominus] WebSocket error on port " .. port .. ":", err)
+		onError = function(statusCode, errorMessage)
+			if transportOpen then
+				warn("[Dominus 2] Bridge error " .. tostring(statusCode) .. ": " .. tostring(errorMessage))
+			end
 		end,
-	})
-
-	return success
+	}, 4)
+	if not success then
+		stopSession()
+		if shouldReconnect then
+			scheduleReconnect()
+		end
+	end
 end
 
-local function connect()
-	if connected then
-		log("Already connected, disconnecting first...")
-		WsClient.disconnect()
-		connected = false
+scheduleReconnect = function()
+	if reconnectScheduled or shuttingDown or not shouldReconnect then
+		return
 	end
-
-	log("Scanning for Dominus server (ports " .. BASE_PORT .. "-" .. (BASE_PORT + MAX_PORT_SCAN - 1) .. ")...")
-
-	for i = 0, MAX_PORT_SCAN - 1 do
-		local port = BASE_PORT + i
-		local success = tryConnectToPort(port)
-		if success then
-			return
+	reconnectScheduled = true
+	local delaySeconds = reconnectDelay
+	reconnectDelay = math.min(reconnectDelay * 2, 30)
+	task.delay(delaySeconds, function()
+		reconnectScheduled = false
+		if shouldReconnect and not shuttingDown then
+			task.spawn(connect)
 		end
-	end
-
-	warn("[Dominus] No Dominus server found. Is `dominus` or `dominus mcp` running?")
+	end)
 end
 
 connectButton.Click:Connect(function()
-	if connected then
+	if transportOpen or WsClient.isConnected() then
+		shouldReconnect = false
 		WsClient.disconnect()
-		connected = false
-		connectButton:SetActive(false)
+		stopSession()
+		log("Automatic connection disabled; click Dominus 2 to reconnect")
 	else
-		connect()
+		shouldReconnect = true
+		reconnectDelay = 1
+		task.spawn(connect)
 	end
 end)
 
--- Auto-connect on plugin load
 task.delay(1, function()
-	connect()
+	task.spawn(connect)
 end)
 
 plugin.Unloading:Connect(function()
-	Watcher.stop()
+	shuttingDown = true
+	shouldReconnect = false
+	stopSession()
 	WsClient.disconnect()
 end)