Enhance MCP tool functionality and improve logging capabilities

- Introduced new MCP tools for executing Luau code, inserting models, and retrieving console output, enhancing the plugin's capabilities for user interaction.
- Implemented a logging system to capture console output and provide detailed feedback during code execution and play mode operations.
- Updated the MCP stream to handle passthrough commands, allowing for better integration of tool commands and improved error handling.
- Enhanced the TODO list to reflect new tasks related to documentation for the added MCP tools and logging features.

These changes aim to improve the overall functionality and user experience of the MCP integration, particularly in script execution and model management.

plugin/src/App/init.lua

@@ -38,6 +38,7 @@ local Dictionary = require(Plugin.Dictionary)
 local ServeSession = require(Plugin.ServeSession)
 local ApiContext = require(Plugin.ApiContext)
 local McpStream = require(Plugin.McpStream)
+local McpTools = require(Plugin.McpTools)
 local PatchSet = require(Plugin.PatchSet)
 local PatchTree = require(Plugin.PatchTree)
 local preloadAssets = require(Plugin.preloadAssets)
@@ -214,6 +215,24 @@ function App:init()
 			onGetScriptCommand = function(requestId, data)
 				return self:handleMcpGetScript(requestId, data)
 			end,
+			onToolCommand = function(requestId, toolType, args)
+				return Promise.new(function(resolve)
+					local ok, result = pcall(McpTools.dispatch, toolType, args)
+					if ok then
+						resolve({
+							requestId = requestId,
+							status = "success",
+							response = result or "",
+						})
+					else
+						resolve({
+							requestId = requestId,
+							status = "error",
+							response = tostring(result),
+						})
+					end
+				end)
+			end,
 			getPluginConfig = function(key)
 				return Settings:get(key)
 			end,

plugin/src/McpStream.lua

@@ -10,6 +10,15 @@ local strict = require(script.Parent.strict)
 
 local RECONNECT_INTERVAL = 5
 
+local PASSTHROUGH_TOOLS = {
+	run_code = true,
+	insert_model = true,
+	get_console_output = true,
+	get_studio_mode = true,
+	start_stop_play = true,
+	run_script_in_play_mode = true,
+}
+
 local McpStream = {}
 McpStream.__index = McpStream
 
@@ -18,6 +27,7 @@ function McpStream.new(options)
 
 	self._onSyncCommand = options.onSyncCommand
 	self._onGetScriptCommand = options.onGetScriptCommand
+	self._onToolCommand = options.onToolCommand
 	self._getPluginConfig = options.getPluginConfig
 	self._wsClient = nil
 	self._running = false
@@ -153,6 +163,33 @@ function McpStream:_tryConnect()
 						end)
 					end
 				end)
+		elseif data.type and data.requestId and PASSTHROUGH_TOOLS[data.type] then
+			Log.info("MCP stream: received {} command (requestId={})", data.type, data.requestId)
+
+			local resultPromise = self._onToolCommand(data.requestId, data.type, data.args or {})
+
+			resultPromise
+				:andThen(function(result)
+					if self._wsClient then
+						local json = HttpService:JSONEncode(result)
+						pcall(function()
+							self._wsClient:Send(json)
+						end)
+					end
+				end)
+				:catch(function(err)
+					Log.warn("MCP stream: {} command failed: {}", data.type, tostring(err))
+					if self._wsClient then
+						local errorResult = HttpService:JSONEncode({
+							requestId = data.requestId,
+							status = "error",
+							response = tostring(err),
+						})
+						pcall(function()
+							self._wsClient:Send(errorResult)
+						end)
+					end
+				end)
 		end
 	end)
 

plugin/src/McpTools/GetConsoleOutput.luau

@@ -0,0 +1,10 @@
+--!strict
+
+local ConsoleOutput = require(script.Parent.Utils.ConsoleOutput)
+
+local function handleGetConsoleOutput(): string?
+	local result = ConsoleOutput.outputMessage or ""
+	return result
+end
+
+return handleGetConsoleOutput

plugin/src/McpTools/GetStudioMode.luau

@@ -0,0 +1,9 @@
+--!strict
+
+local StudioModeState = require(script.Parent.Utils.StudioModeState)
+
+local function handleGetStudioMode(): string?
+	return StudioModeState.studioMode
+end
+
+return handleGetStudioMode

plugin/src/McpTools/InsertModel.luau

@@ -0,0 +1,135 @@
+--!strict
+--!nolint DeprecatedApi
+
+local InsertService = game:GetService("InsertService")
+
+type Args = {
+	query: string,
+}
+
+local INSERT_MAX_SEARCH_DEPTH = 2048
+local INSERT_MAX_DISTANCE_AWAY = 20
+
+local function getInsertPosition(): Vector3
+	local camera = workspace.CurrentCamera
+	local viewportPoint = camera.ViewportSize / 2
+	local unitRay = camera:ViewportPointToRay(viewportPoint.X, viewportPoint.Y, 0)
+
+	local ray = Ray.new(unitRay.Origin, unitRay.Direction * INSERT_MAX_SEARCH_DEPTH)
+	local params = RaycastParams.new()
+	params.BruteForceAllSlow = true
+
+	local result = workspace:Raycast(ray.Origin, ray.Direction, params)
+
+	if result then
+		return result.Position
+	else
+		return camera.CFrame.Position + unitRay.Direction * INSERT_MAX_DISTANCE_AWAY
+	end
+end
+
+type GetFreeModelsResponse = {
+	[number]: {
+		CurrentStartIndex: number,
+		TotalCount: number,
+		Results: {
+			[number]: {
+				Name: string,
+				AssetId: number,
+				AssetVersionId: number,
+				CreatorName: string,
+			},
+		},
+	},
+}
+
+local function toTitleCase(str: string): string
+	local function titleCase(first: string, rest: string): string
+		return first:upper() .. rest:lower()
+	end
+
+	local intermediate = string.gsub(str, "(%a)([%w_']*)", titleCase :: (string) -> string)
+	return intermediate:gsub("%s+", "")
+end
+
+local function collapseObjectsIntoContainer(objects: { Instance }): Instance?
+	local isPhysical = false
+	for _, object in objects do
+		if object:IsA("PVInstance") then
+			isPhysical = true
+			break
+		end
+	end
+
+	if isPhysical then
+		local model = Instance.new("Model")
+		for _, object in objects do
+			object.Parent = model
+		end
+		return model
+	end
+
+	if #objects > 1 then
+		local folder = Instance.new("Folder")
+		for _, object in objects do
+			object.Parent = folder
+		end
+		return folder
+	end
+
+	return objects[1]
+end
+
+local function loadAsset(assetId: number): Instance?
+	local objects = game:GetObjects(`rbxassetid://{assetId}`)
+	return collapseObjectsIntoContainer(objects)
+end
+
+local function getAssets(query: string): number?
+	local results: GetFreeModelsResponse = InsertService:GetFreeModels(query, 0)
+	local assets = {}
+	for i, result in results[1].Results do
+		if i > 6 then
+			break
+		end
+		table.insert(assets, result.AssetId)
+	end
+
+	return table.remove(assets, 1)
+end
+
+local function insertFromMarketplace(query: string): string
+	local primaryResult = getAssets(query)
+	if not primaryResult then
+		error("Failed to find asset")
+	end
+
+	local instance = loadAsset(primaryResult)
+	if not instance then
+		error("Failed to load asset")
+	end
+
+	local name = toTitleCase(query)
+	local i = 1
+	while workspace:FindFirstChild(name) do
+		name = query .. i
+		i += 1
+	end
+
+	instance.Name = name
+	instance.Parent = workspace
+
+	if instance:IsA("Model") then
+		instance:PivotTo(CFrame.new(getInsertPosition()))
+	end
+
+	return name
+end
+
+local function handleInsertModel(args: Args): string?
+	assert(type(args.query) == "string", "Missing query in insert_model")
+
+	return insertFromMarketplace(args.query)
+end
+
+return handleInsertModel

plugin/src/McpTools/RunCode.luau

@@ -0,0 +1,132 @@
+--!strict
+--!nolint DeprecatedApi
+
+local DeepCopy = require(script.Parent.Utils.DeepCopy)
+
+type Args = {
+	command: string,
+}
+
+local function valueToString(value: any): string
+	local t = type(value)
+	if t == "number" then
+		if value ~= value then
+			return "NaN"
+		elseif value == math.huge then
+			return "Infinity"
+		elseif value == -math.huge then
+			return "-Infinity"
+		end
+		return tostring(value)
+	elseif t == "nil" then
+		return "null"
+	elseif t == "string" then
+		return string.format("%q", value)
+	elseif t == "boolean" then
+		return tostring(value)
+	elseif t == "table" then
+		local parts: { string } = {}
+		local isArray = #value > 0
+		if isArray then
+			for _, v in value :: { any } do
+				table.insert(parts, valueToString(v))
+			end
+			return "[" .. table.concat(parts, ", ") .. "]"
+		else
+			for k, v in value :: { [any]: any } do
+				table.insert(parts, valueToString(k) .. ": " .. valueToString(v))
+			end
+			return "{" .. table.concat(parts, ", ") .. "}"
+		end
+	end
+	return tostring(value)
+end
+
+local function runCodeWithOutput(command: string): string
+	local output = ""
+
+	local function serializeTable(arg: any): any
+		if typeof(arg) == "table" then
+			local hasTable = false
+			for key, _ in arg do
+				if typeof(key) ~= "number" then
+					hasTable = true
+					break
+				end
+			end
+
+			local newArg = {}
+			for key, value in arg do
+				local newKey = serializeTable(key)
+				newArg[if hasTable then tostring(newKey) else newKey] = serializeTable(value)
+			end
+			return newArg
+		elseif type(arg) == "userdata" then
+			return tostring(arg) or "UNABLE_TO_SERIALIZE"
+		end
+		return arg
+	end
+
+	local function toStrTable(t: { any }): { string }
+		local clonedTable = DeepCopy(t)
+		local strTable = table.create(#clonedTable)
+		for i, arg in clonedTable do
+			local serializedArg = serializeTable(arg)
+			strTable[i] = valueToString(serializedArg)
+		end
+		return strTable
+	end
+
+	local function addToOutput(header: string, ...)
+		local strResults = toStrTable(table.pack(...))
+		output ..= header .. " " .. table.concat(strResults, "\t") .. "\n"
+	end
+
+	local function executeCode(): { any }
+		local chunk = loadstring(command) :: any
+		local chunkfenv = getfenv(chunk)
+
+		local oldPrint = print
+		chunkfenv.print = function(...)
+			oldPrint(...)
+			addToOutput("[OUTPUT]", ...)
+		end
+
+		local oldWarn = warn
+		chunkfenv.warn = function(...)
+			oldWarn(...)
+			addToOutput("[WARNING]", ...)
+		end
+
+		local oldError = error
+		chunkfenv.error = function(...)
+			oldError(...)
+			addToOutput("[ERROR]", ...)
+		end
+
+		local results = table.pack(chunk())
+		if #results > 0 then
+			addToOutput("[RETURNED RESULTS]", table.unpack(results))
+		end
+
+		return results
+	end
+
+	local ok, errorMessage = pcall(executeCode)
+	if not ok then
+		addToOutput("[UNEXPECTED ERROR]", errorMessage)
+	end
+
+	return output
+end
+
+local ConsoleOutput = require(script.Parent.Utils.ConsoleOutput)
+
+local function handleRunCode(args: Args): string?
+	assert(type(args.command) == "string", "Missing command in run_code")
+
+	ConsoleOutput.outputMessage = ""
+	return runCodeWithOutput(args.command)
+end
+
+return handleRunCode

plugin/src/McpTools/RunScriptInPlayMode.luau

@@ -0,0 +1,145 @@
+--!strict
+
+local ServerScriptService = game:GetService("ServerScriptService")
+local StudioTestService = game:GetService("StudioTestService")
+
+local ConsoleOutput = require(script.Parent.Utils.ConsoleOutput)
+local StartStopPlay = require(script.Parent.StartStopPlay)
+local StudioModeState = require(script.Parent.Utils.StudioModeState)
+
+type Args = {
+	code: string,
+	timeout: number?,
+	mode: StudioModeState.TestMode,
+}
+
+local TEST_SCRIPT_NAME = "AtlasTestScript"
+
+local function buildTestRunnerSource(userCode: string, timeout: number): string
+	local timeoutStr = tostring(timeout)
+	return [[
+local StudioTestService = game:GetService("StudioTestService")
+local LogService = game:GetService("LogService")
+local RunService = game:GetService("RunService")
+
+if not RunService:IsRunning() then
+	return
+end
+
+local capturedLogs = {}
+local logConnection = nil
+
+local MESSAGE_TYPE_MAP = {
+	[Enum.MessageType.MessageOutput] = "output",
+	[Enum.MessageType.MessageInfo] = "info",
+	[Enum.MessageType.MessageWarning] = "warning",
+	[Enum.MessageType.MessageError] = "error",
+}
+
+logConnection = LogService.MessageOut:Connect(function(message, messageType)
+	if string.sub(message, 1, 7) == "[Atlas]" then return end
+	table.insert(capturedLogs, {
+		level = MESSAGE_TYPE_MAP[messageType] or "output",
+		message = message,
+		ts = os.clock(),
+	})
+end)
+
+local startTime = os.clock()
+local isTimeout = false
+
+local function endTest()
+	local duration = os.clock() - startTime
+	task.wait(0.1)
+
+	if logConnection then logConnection:Disconnect() end
+
+	local errors = {}
+	for _, log in ipairs(capturedLogs) do
+		if log.level == "error" or log.level == "warning" then
+			table.insert(errors, log)
+		end
+	end
+
+	StudioTestService:EndTest({
+		success = ok,
+		value = if ok then tostring(result) else nil,
+		error = if not ok then tostring(result) else nil,
+		logs = capturedLogs,
+		errors = errors,
+		duration = duration,
+		isTimeout = isTimeout,
+	})
+end
+
+
+task.spawn(function()
+	task.wait(]] .. timeoutStr .. [[)
+	isTimeout = true
+	endTest()
+end)
+
+local ok, result = pcall(function()
+]] .. userCode .. [[
+
+end)
+
+endTest()
+]]
+end
+
+local function removeTestScript()
+	local testScript = ServerScriptService:FindFirstChild(TEST_SCRIPT_NAME)
+	if testScript then
+		testScript:Destroy()
+	end
+end
+
+local function injectTestScript(args: Args)
+	removeTestScript()
+
+	local source = buildTestRunnerSource(args.code, args.timeout or 1_000_000)
+	local testScript = Instance.new("Script")
+	testScript.Name = TEST_SCRIPT_NAME
+	testScript.Parent = ServerScriptService
+	testScript.Source = source
+end
+
+local function runScriptInPlayMode(args: Args): string?
+	injectTestScript(args)
+
+	ConsoleOutput.outputMessage = ""
+	local success, result = pcall(function()
+		if args.mode == "start_play" then
+			StudioModeState.studioMode = "start_play"
+			return StudioTestService:ExecutePlayModeAsync({})
+		else
+			StudioModeState.studioMode = "run_server"
+			return StudioTestService:ExecuteRunModeAsync({})
+		end
+	end)
+
+	StudioModeState.studioMode = "stop"
+	removeTestScript()
+
+	if not success then
+		error(`Failed to run script in play mode: {result}`)
+	end
+
+	return result
+end
+
+local function handleRunScriptInPlayMode(args: Args): string?
+	assert(type(args.code) == "string" and args.code ~= "", "Missing code in run_script_in_play_mode")
+	assert(
+		args.timeout == nil or (type(args.timeout) == "number" and args.timeout >= 0),
+		"Timeout must be a number >= 0"
+	)
+
+	StartStopPlay({ mode = "stop" })
+	task.wait(0.2)
+
+	return runScriptInPlayMode(args)
+end
+
+return handleRunScriptInPlayMode

plugin/src/McpTools/StartStopPlay.luau

@@ -0,0 +1,88 @@
+--!strict
+
+local StudioTestService = game:GetService("StudioTestService")
+
+local ConsoleOutput = require(script.Parent.Utils.ConsoleOutput)
+local GameStopUtil = require(script.Parent.Utils.GameStopUtil)
+local StudioModeState = require(script.Parent.Utils.StudioModeState)
+
+type Args = {
+	mode: StudioModeState.StartStopPlayMode,
+}
+
+local function callWithTimeout(callback: () -> any, timeout: number, timeoutResult: string): (string?, boolean)
+	local BindableEvent = Instance.new("BindableEvent")
+	local isDone = false
+	local _success = true
+	local result = timeoutResult
+	local isTimeout = false
+
+	task.spawn(function()
+		_success, result = pcall(function()
+			return callback()
+		end)
+		if not isDone then
+			isDone = true
+			BindableEvent:Fire()
+		end
+	end)
+	task.spawn(function()
+		task.wait(timeout)
+		if not isDone then
+			isDone = true
+			isTimeout = true
+			BindableEvent:Fire()
+		end
+	end)
+
+	if not isDone then
+		BindableEvent.Event:Wait()
+	end
+	return result, isTimeout
+end
+
+local function startPlay(): string?
+	ConsoleOutput.outputMessage = ""
+	StudioModeState.studioMode = "start_play"
+	local result, isTimeout = callWithTimeout(function()
+		return StudioTestService:ExecutePlayModeAsync({})
+	end, 0.1, "Started play")
+	if not isTimeout then
+		return "Already in play mode"
+	end
+	return result
+end
+
+local function stop(): string?
+	GameStopUtil.stopPlay()
+	task.wait(1)
+	StudioModeState.studioMode = "stop"
+	return "Stopped"
+end
+
+local function runServer(): string?
+	StudioModeState.studioMode = "run_server"
+	local result, isTimeout = callWithTimeout(function()
+		return StudioTestService:ExecuteRunModeAsync({})
+	end, 0.1, "Ran server")
+	if not isTimeout then
+		return "Already in run_server mode"
+	end
+	return result
+end
+
+local function handleStartStopPlay(args: Args): string?
+	assert(type(args.mode) == "string", "Missing mode in start_stop_play")
+
+	if args.mode == "start_play" then
+		return startPlay()
+	elseif args.mode == "stop" then
+		return stop()
+	elseif args.mode == "run_server" then
+		return runServer()
+	end
+
+	error("Invalid mode in start_stop_play, must be start_play, stop, or run_server")
+end
+
+return handleStartStopPlay

plugin/src/McpTools/Utils/ConsoleOutput.luau

@@ -0,0 +1,24 @@
+--!strict
+
+local LogService = game:GetService("LogService")
+
+local MAX_OUTPUT = 10000
+
+local ConsoleOutput = {
+	outputMessage = "",
+}
+
+function ConsoleOutput.startListener(): RBXScriptConnection
+	local conn = LogService.MessageOut:Connect(function(message: string, _: Enum.MessageType)
+		if #ConsoleOutput.outputMessage + #message <= MAX_OUTPUT then
+			if #ConsoleOutput.outputMessage == 0 then
+				ConsoleOutput.outputMessage = message
+			else
+				ConsoleOutput.outputMessage = ConsoleOutput.outputMessage .. "\n" .. message
+			end
+		end
+	end)
+	return conn
+end
+
+return ConsoleOutput

plugin/src/McpTools/Utils/DeepCopy.luau

@@ -0,0 +1,31 @@
+--!strict
+--!optimize 2
+
+function DeepCopy<T>(value: T, cache_: { [any]: any }?): T
+	if type(value) ~= "table" then
+		return value
+	end
+
+	local cache: { [any]: any } = cache_ or {}
+
+	local cachedResult = cache[value]
+	if cachedResult then
+		return cachedResult
+	end
+
+	local result = {}
+	cache[value] = result
+
+	for k, v in next, value do
+		result[DeepCopy(k, cache)] = DeepCopy(v, cache)
+	end
+
+	local valueMT = getmetatable(value)
+	if valueMT then
+		setmetatable(result, valueMT :: any)
+	end
+
+	return result :: any
+end
+
+return DeepCopy

plugin/src/McpTools/Utils/GameStopUtil.luau

@@ -0,0 +1,32 @@
+--!strict
+
+local StudioTestService = game:GetService("StudioTestService")
+
+local MCP_STOP_PLAY_KEY = "Atlas_MCP_STOP_PLAY"
+
+local GameStopUtil = {}
+
+local _plugin: Plugin? = nil
+
+function GameStopUtil.setPlugin(p: Plugin)
+	_plugin = p
+end
+
+function GameStopUtil.stopPlay()
+	assert(_plugin, "GameStopUtil: plugin not set")
+	_plugin:SetSetting(MCP_STOP_PLAY_KEY, true)
+end
+
+function GameStopUtil.monitorForStopPlay()
+	assert(_plugin, "GameStopUtil: plugin not set")
+	_plugin:SetSetting(MCP_STOP_PLAY_KEY, false)
+	while true do
+		if _plugin:GetSetting(MCP_STOP_PLAY_KEY) then
+			_plugin:SetSetting(MCP_STOP_PLAY_KEY, false)
+			StudioTestService:EndTest({})
+		end
+		task.wait(1)
+	end
+end
+
+return GameStopUtil

plugin/src/McpTools/Utils/StudioModeState.luau

@@ -0,0 +1,8 @@
+--!strict
+
+export type TestMode = "start_play" | "run_server"
+export type StartStopPlayMode = TestMode | "stop"
+
+return {
+	studioMode = "stop" :: StartStopPlayMode,
+}

plugin/src/McpTools/init.luau

@@ -0,0 +1,26 @@
+--!strict
+
+local strict = require(script.Parent.strict)
+
+local tools: { [string]: (any) -> string? } = {
+	run_code = require(script.RunCode),
+	insert_model = require(script.InsertModel),
+	get_console_output = require(script.GetConsoleOutput),
+	get_studio_mode = require(script.GetStudioMode),
+	start_stop_play = require(script.StartStopPlay),
+	run_script_in_play_mode = require(script.RunScriptInPlayMode),
+}
+
+local McpTools = {}
+
+function McpTools.dispatch(toolType: string, args: any): string?
+	local handler = tools[toolType]
+	if not handler then
+		error("Unknown MCP tool: " .. tostring(toolType))
+	end
+	return handler(args)
+end
+
+return strict("McpTools", {
+	dispatch = McpTools.dispatch,
+})

plugin/src/init.server.lua

@@ -2,6 +2,8 @@ if not plugin then
 	return
 end
 
+local RunService = game:GetService("RunService")
+
 local Rojo = script:FindFirstAncestor("Rojo")
 local Packages = Rojo.Packages
 
@@ -12,10 +14,21 @@ local Settings = require(script.Settings)
 local Config = require(script.Config)
 local App = require(script.App)
 
+local ConsoleOutput = require(script.McpTools.Utils.ConsoleOutput)
+local GameStopUtil = require(script.McpTools.Utils.GameStopUtil)
+
+GameStopUtil.setPlugin(plugin)
+
 Log.setLogLevelThunk(function()
 	return Log.Level[Settings:get("logLevel")] or Log.Level.Info
 end)
 
+if RunService:IsRunning() and RunService:IsServer() then
+	task.spawn(GameStopUtil.monitorForStopPlay)
+end
+
+local consoleOutputConn = ConsoleOutput.startListener()
+
 local app = Roact.createElement(App, {
 	plugin = plugin,
 })
@@ -23,6 +36,7 @@ local tree = Roact.mount(app, game:GetService("CoreGui"), "Atlas UI")
 
 plugin.Unloading:Connect(function()
 	Roact.unmount(tree)
+	consoleOutputConn:Disconnect()
 end)
 
 if Config.isDevBuild then