diff --git a/README.md b/README.md index bef7a89b..9dcb7dd7 100644 --- a/README.md +++ b/README.md @@ -191,6 +191,67 @@ Atlas includes a sync lock to prevent conflicts when multiple people work on the **Recommended team workflow:** One designated person syncs in persistent mode (one-shot disabled). Everyone else uses one-shot mode for quick, conflict-free previews. +## MCP Integration (AI Agents) + +`atlas serve` automatically hosts an MCP (Model Context Protocol) server on the same port as the API. This lets AI agents like Cursor and Claude Code trigger syncs to Roblox Studio programmatically. + +### Available Tools + +| Tool | Description | +|------|-------------| +| `atlas_sync` | Sync filesystem changes to Studio. Auto-accepts if all changes are pre-selected by git. Blocks until the user accepts or rejects if there are unresolved changes. | + +### Adding to Cursor + +Create or edit `.cursor/mcp.json` in your project root: + +```json +{ + "mcpServers": { + "atlas": { + "url": "http://localhost:34873/mcp" + } + } +} +``` + +If your project uses a custom port (via `servePort` in the project file), update the URL accordingly. + +### Adding to Claude Code + +```bash +claude mcp add atlas --transport http http://localhost:34873/mcp +``` + +Or add it to your project-level config at `.claude/settings.json`: + +```json +{ + "mcpServers": { + "atlas": { + "type": "url", + "url": "http://localhost:34873/mcp" + } + } +} +``` + +### How It Works + +1. Start `atlas serve` — the MCP server is available immediately +2. The Studio plugin connects to the MCP stream in the background (retries every 5s) +3. When an agent calls `atlas_sync`: + - If all changes are pre-selected by the git system, they are **auto-accepted** (fast-forward) + - If there are unresolved changes, the **Confirming page** appears in Studio for the user to review + - The agent blocks until the sync completes and receives a list of synced file paths with directions (`push`/`pull`) +4. If the plugin is already connected in live sync mode, the tool returns immediately with a notice that sync is automatic + +### Requirements + +- `atlas serve` must be running +- Roblox Studio must be open with the Atlas plugin installed +- The plugin does **not** need to be manually connected — the MCP stream connects automatically + ## File Formats ### Scripts diff --git a/TODO.txt b/TODO.txt index 0a7bac63..739c5aa2 100644 --- a/TODO.txt +++ b/TODO.txt @@ -1,10 +1,6 @@ [vscode extension] get the schemas for everything into cursor somehow -[versioning] -setup nightly build -suffix version for beta - [merging] we should make the property diffs have arrows next to eachother so you can like "undo" certain things need proper merging support in this UI. @@ -12,13 +8,7 @@ make sure for fully deleted/added files we can use the merge UI as expected. [syncback] maybe make syncback automatically create a branch and merge it in, rather than overwriting your project files? - -[automated review of all changes done, similar to cursor agent review for the game] -send in discord -require people to green check it when done -have the cursor style issue context copyable -tag who last modified the file, too -https://cursor.com/dashboard?tab=bugbot +in order to do this, we must make dedup work flawlessly. [Move tests to dedicated location] Right now we have a severe amount of tests in the actual implementation files. I'd like to move those to a dedicated location because it's bloating the implementation files. @@ -327,6 +317,12 @@ because have already loaded it, we can detect this. this is an optimization to avoid writing 100,000 files or something on larger projects when file bytes match perfectly, or directories verifiably already exist. make sure this is applied to plugin-based syncback, too. +[fix dedup] +dedup is broken due to relying on sort order +recreate broken state by forcing random order in all matching dedup systems of the incoming children list +if it cannot match correctly with randomization, then it is inherently broken +in all 3 matchers, simply randomize the order of the children - this will break the tests. + [mcp] * atlas_sync - Basically equiv to git fast-forward merge @@ -348,3 +344,35 @@ make sure this is applied to plugin-based syncback, too. - Does nothing if we are running in connected-mode (already syncing) - Agent should await the completion of this sync before proceeding, this mcp command needs to yield and give the response as described above. +[mcp] +* syncback + - invokes plugin based syncback + - ensure this doesn't interrupt the connection of the mcp, since i think we shutdown the api? we should just not shutdown the api or something + +[automated review of all changes done, similar to cursor agent review for the game] +it's a discord bot +we send discord channels +some way for users to check off the item in the channel as "fixed" +have the cursor style issue context copyable +this would work by requesting the rbxl from roblox like syncback already does +tag who last modified the file, too if we can - we need to somehow request the version history for a given script file - this is kinda hard, but we should investigate how hard it would be to do via OpenCloud +ideally for the .rbxl downloading, we use OpenCloud so we can use an API key rather than having my personal cookie +inspo: Cursor's "Agent Review" feature +additional inspo: https://cursor.com/dashboard?tab=bugbot +- create an atlast clone of the project +- every ~5 minutes, download and syncback. + - if there are changes, perform an audit on the diff + - copy the Agent Review prompt that cursor uses, somehow + - if there are any issues, send the exact same format that Cursor's Agent Review uses into the discord in a code block + - agent should also determine a severity (Low/Medium/High/Critical), and if it's Critical - it will @here +- should be really easy to spin these up for other games + - possibly configurable automatically throgh discord commands? + - could allow for nice administration like resetting the project from scratch (rm -rf ing it and starting over) + - reports errors + - add games via the discord by PlaceId + - automatically remembers which discord channels to which placeid + - support multiple discords using the same placeid (runs under like DiscordServerId/PlaceId) or something + - only allow specific user ids to administer these commands: 176643682689089545 (me), 1010301288946339920 (Tony) + - alternatively, and this would be better: allow the user to specify their own OpenCloud api key (or cookie if we have to revert to that) to use (make sure providing this is very secure) and then we can let them have at it. + - need to ensure that only certain perm users can use the bot + - we need to make sure not just any public server with a bot link can just run this, because it will be using my system resources and could be insecure. this is a proprietary tool. diff --git a/plugin/src/App/init.lua b/plugin/src/App/init.lua index fb3f693f..1c581b26 100644 --- a/plugin/src/App/init.lua +++ b/plugin/src/App/init.lua @@ -37,6 +37,7 @@ local strict = require(Plugin.strict) local Dictionary = require(Plugin.Dictionary) local ServeSession = require(Plugin.ServeSession) local ApiContext = require(Plugin.ApiContext) +local McpStream = require(Plugin.McpStream) local PatchSet = require(Plugin.PatchSet) local PatchTree = require(Plugin.PatchTree) local preloadAssets = require(Plugin.preloadAssets) @@ -198,6 +199,18 @@ function App:init() self:clearRunningConnectionInfo() end end) + + if RunService:IsEdit() then + self.mcpStream = McpStream.new({ + onSyncCommand = function(requestId) + return self:startMcpSync(requestId) + end, + getPluginConfig = function(key) + return Settings:get(key) + end, + }) + self.mcpStream:start() + end end function App:willUnmount() @@ -216,6 +229,11 @@ function App:willUnmount() self.autoConnectPlaytestServerListener() self:clearRunningConnectionInfo() + + if self.mcpStream then + self.mcpStream:stop() + self.mcpStream = nil + end end function App:addNotification(notif: { @@ -1099,6 +1117,253 @@ function App:endSession() Log.trace("Session terminated by user") end +function App:startMcpSync(requestId) + return Promise.new(function(resolve, _reject) + if self.state.appStatus == AppStatus.Connected then + resolve({ + requestId = requestId, + status = "already_connected", + changes = {}, + message = "Atlas is already connected in live sync mode. All changes sync automatically.", + }) + return + end + + if self.state.appStatus == AppStatus.Connecting or self.state.appStatus == AppStatus.Confirming then + resolve({ + requestId = requestId, + status = "sync_in_progress", + changes = {}, + message = "A sync operation is already in progress.", + }) + return + end + + if self.serveSession ~= nil then + Log.trace("MCP sync: ending existing session before starting") + self:endSession() + end + + local host = Config.defaultHost + local port = Config.defaultPort + local baseUrl = ("http://%s:%s"):format(host, port) + local apiContext = ApiContext.new(baseUrl) + + local serveSession = ServeSession.new({ + apiContext = apiContext, + twoWaySync = false, + }) + + local mcpPatchTree = nil + local mcpSyncResolved = false + + local function resolveOnce(result) + if mcpSyncResolved then + return + end + mcpSyncResolved = true + result.requestId = requestId + resolve(result) + end + + serveSession:setUpdateLoadingTextCallback(function(text: string) + self:setState({ + connectingText = text, + }) + end) + + local cachedServerInfo = nil + + serveSession:setConfirmCallback(function(instanceMap, patch, serverInfo) + cachedServerInfo = serverInfo + + PatchSet.removeDataModelName(patch, instanceMap) + + if PatchSet.isEmpty(patch) then + Log.info("MCP sync: no changes to sync") + resolveOnce({ + status = "empty", + changes = {}, + }) + return "Accept" + end + + local gitMetadata = serverInfo.gitMetadata + local patchTree = PatchTree.build(patch, instanceMap, { "Property", "Current", "Incoming" }, gitMetadata) + mcpPatchTree = patchTree + + local allPreSelected = true + local hasSelectableNodes = false + patchTree:forEach(function(node) + if node.patchType then + hasSelectableNodes = true + if node.defaultSelection == nil then + allPreSelected = false + end + end + end) + + if not hasSelectableNodes then + Log.info("MCP sync: no selectable changes, auto-accepting") + resolveOnce({ + status = "empty", + changes = {}, + }) + return "Accept" + end + + if allPreSelected then + Log.info("MCP sync: all changes pre-selected, fast-forwarding") + local selections = PatchTree.buildInitialSelections(patchTree) + local changes = self:_buildMcpChangeList(patchTree, selections) + + local autoSelectedIds = {} + patchTree:forEach(function(node) + if node.patchType and node.defaultSelection and node.defaultSelection ~= "ignore" then + autoSelectedIds[node.id] = true + end + end) + + resolveOnce({ + status = "success", + changes = changes, + }) + + return { + type = "Confirm", + selections = selections, + autoSelectedIds = autoSelectedIds, + } + end + + Log.info("MCP sync: changes require user review, showing confirmation UI") + self:setState({ + appStatus = AppStatus.Confirming, + patchTree = patchTree, + confirmData = { + serverInfo = serverInfo, + }, + toolbarIcon = Assets.Images.PluginButton, + }) + + self:addNotification({ + text = "An AI agent requested a sync. Please review and confirm the changes.", + timeout = 10, + }) + + local result = self.confirmationEvent:Wait() + + if result == "Abort" then + local changes = self:_buildMcpChangeList(patchTree, nil) + resolveOnce({ + status = "rejected", + changes = changes, + message = "User rejected the sync changes.", + }) + return result + end + + if type(result) == "table" and result.type == "Confirm" then + local changes = self:_buildMcpChangeList(patchTree, result.selections) + resolveOnce({ + status = "success", + changes = changes, + }) + return result + end + + resolveOnce({ + status = "success", + changes = {}, + }) + return result + end) + + serveSession:onStatusChanged(function(status, details) + if status == ServeSession.Status.Connecting then + self:setState({ + appStatus = AppStatus.Connecting, + toolbarIcon = Assets.Images.PluginButton, + }) + elseif status == ServeSession.Status.Connected then + self:setState({ + appStatus = AppStatus.Connecting, + connectingText = "Completing sync...", + toolbarIcon = Assets.Images.PluginButton, + }) + elseif status == ServeSession.Status.Disconnected then + self.serveSession = nil + self:setState({ + appStatus = AppStatus.NotConnected, + toolbarIcon = Assets.Images.PluginButton, + patchTree = nil, + }) + if details ~= nil then + Log.warn("MCP sync disconnected with error: {}", details) + resolveOnce({ + status = "error", + changes = {}, + message = tostring(details), + }) + end + end + end) + + serveSession:setInitialSyncCompleteCallback(function() + Log.info("MCP sync: initial sync complete, disconnecting") + if self.serveSession and self.serveSession == serveSession then + self:endSession() + end + end) + + serveSession:start() + self.serveSession = serveSession + end) +end + +function App:_buildMcpChangeList(patchTree, selections) + local changes = {} + + patchTree:forEach(function(node) + if not node.patchType then + return + end + + local direction = nil + if selections then + local sel = selections[node.id] + if sel == "push" then + direction = "push" + elseif sel == "pull" then + direction = "pull" + else + return + end + else + direction = node.defaultSelection or "push" + end + + local segments = {} + local current = node + while current and current.id ~= "ROOT" do + table.insert(segments, 1, current.name or current.id) + if current.parentId then + current = patchTree:getNode(current.parentId) + else + break + end + end + + local path = table.concat(segments, "/") + table.insert(changes, { + path = path, + direction = direction, + }) + end) + + return changes +end + function App:render() local pluginName = "Atlas " .. Version.display(Config.version) diff --git a/plugin/src/McpStream.lua b/plugin/src/McpStream.lua new file mode 100644 index 00000000..da3e2805 --- /dev/null +++ b/plugin/src/McpStream.lua @@ -0,0 +1,150 @@ +local HttpService = game:GetService("HttpService") + +local Rojo = script:FindFirstAncestor("Rojo") +local Packages = Rojo.Packages + +local Log = require(Packages.Log) + +local Config = require(script.Parent.Config) +local strict = require(script.Parent.strict) + +local RECONNECT_INTERVAL = 5 + +local McpStream = {} +McpStream.__index = McpStream + +function McpStream.new(options) + local self = setmetatable({}, McpStream) + + self._onSyncCommand = options.onSyncCommand + self._getPluginConfig = options.getPluginConfig + self._wsClient = nil + self._running = false + self._loopThread = nil + + return self +end + +function McpStream:start() + if self._running then + return + end + self._running = true + + self._loopThread = task.spawn(function() + while self._running do + self:_tryConnect() + if self._running then + task.wait(RECONNECT_INTERVAL) + end + end + end) +end + +function McpStream:stop() + self._running = false + if self._loopThread then + task.cancel(self._loopThread) + self._loopThread = nil + end + self:_disconnect() +end + +function McpStream:_disconnect() + if self._wsClient then + pcall(function() + self._wsClient:Destroy() + end) + self._wsClient = nil + end +end + +function McpStream:_tryConnect() + local host = Config.defaultHost + local port = Config.defaultPort + local url = ("ws://%s:%s/api/mcp/stream"):format(host, port) + + local success, wsClient = + pcall(HttpService.CreateWebStreamClient, HttpService, Enum.WebStreamClientType.WebSocket, { Url = url }) + if not success then + return + end + + self._wsClient = wsClient + Log.trace("MCP stream connected to {}", url) + + local greeting = HttpService:JSONEncode({ + type = "hello", + twoWaySync = self._getPluginConfig("twoWaySync"), + oneShotSync = self._getPluginConfig("oneShotSync"), + confirmationBehavior = self._getPluginConfig("confirmationBehavior"), + }) + pcall(function() + wsClient:Send(greeting) + end) + + local done = false + local closed, errored, received + + received = wsClient.MessageReceived:Connect(function(msg) + local ok, data = pcall(HttpService.JSONDecode, HttpService, msg) + if not ok then + return + end + + if data.type == "sync" and data.requestId then + Log.info("MCP stream: received sync command (requestId={})", data.requestId) + + local resultPromise = self._onSyncCommand(data.requestId) + + 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: sync command failed: {}", tostring(err)) + if self._wsClient then + local errorResult = HttpService:JSONEncode({ + requestId = data.requestId, + status = "error", + changes = {}, + message = tostring(err), + }) + pcall(function() + self._wsClient:Send(errorResult) + end) + end + end) + end + end) + + closed = wsClient.Closed:Connect(function() + done = true + closed:Disconnect() + errored:Disconnect() + received:Disconnect() + Log.trace("MCP stream disconnected") + end) + + errored = wsClient.Error:Connect(function(_code, _msg) + done = true + closed:Disconnect() + errored:Disconnect() + received:Disconnect() + end) + + while not done and self._running do + task.wait(1) + end + + self:_disconnect() +end + +return strict("McpStream", { + new = McpStream.new, +}) diff --git a/plugin/src/PatchTree.spec.lua b/plugin/src/PatchTree.spec.lua index 88f066b0..0fdeaddc 100644 --- a/plugin/src/PatchTree.spec.lua +++ b/plugin/src/PatchTree.spec.lua @@ -734,4 +734,206 @@ return function() instanceMap:stop() end) end) + + describe("MCP fast-forward detection", function() + it("should detect all-pre-selected tree as fast-forwardable", function() + local patch = PatchSet.newEmpty() + local instanceMap = InstanceMap.new() + + local workspace = game:GetService("Workspace") + local testContainer = Instance.new("Folder") + testContainer.Name = "McpFFTest" + testContainer.Parent = workspace + + local script1 = Instance.new("ModuleScript") + script1.Name = "ModuleA" + script1.Parent = testContainer + + local id1 = HttpService:GenerateGUID(false) + instanceMap:insert(id1, script1) + + table.insert(patch.updated, { + id = id1, + changedName = nil, + changedProperties = { Source = "-- changed" }, + }) + + local gitMetadata = { + changedIds = { id1 }, + scriptCommittedHashes = { [id1] = {} }, + newFileIds = {}, + } + + local tree = PatchTree.build(patch, instanceMap, { "Property", "Current", "Incoming" }, gitMetadata) + + local allPreSelected = true + local hasSelectable = false + tree:forEach(function(node) + if node.patchType then + hasSelectable = true + if node.defaultSelection == nil then + allPreSelected = false + end + end + end) + + -- With git metadata for a changed script, it should have a defaultSelection. + -- The exact value depends on hash matching, but we verify the structure works. + expect(hasSelectable).to.equal(true) + + local selections = PatchTree.buildInitialSelections(tree) + local count = 0 + for _ in selections do + count += 1 + end + -- If allPreSelected, count should match selectable count + if allPreSelected then + local selectableCount = 0 + tree:forEach(function(node) + if node.patchType then + selectableCount += 1 + end + end) + expect(count).to.equal(selectableCount) + end + + testContainer:Destroy() + instanceMap:stop() + end) + + it("should detect tree with nil defaultSelection as NOT fast-forwardable", function() + local patch = PatchSet.newEmpty() + local instanceMap = InstanceMap.new() + + local workspace = game:GetService("Workspace") + local testContainer = Instance.new("Folder") + testContainer.Name = "McpNoFF" + testContainer.Parent = workspace + + local folder = Instance.new("Folder") + folder.Name = "SomeFolder" + folder.Parent = testContainer + + local folderId = HttpService:GenerateGUID(false) + instanceMap:insert(folderId, folder) + + -- An update without git metadata => defaultSelection stays nil + table.insert(patch.updated, { + id = folderId, + changedName = "Renamed", + changedProperties = {}, + }) + + local tree = PatchTree.build(patch, instanceMap, { "Property", "Current", "Incoming" }) + + local allPreSelected = true + tree:forEach(function(node) + if node.patchType and node.defaultSelection == nil then + allPreSelected = false + end + end) + + expect(allPreSelected).to.equal(false) + + testContainer:Destroy() + instanceMap:stop() + end) + end) + + describe("MCP change list path construction", function() + it("should build paths by walking ancestry to ROOT", function() + local patch = PatchSet.newEmpty() + local instanceMap = InstanceMap.new() + + local workspace = game:GetService("Workspace") + local testContainer = Instance.new("Folder") + testContainer.Name = "PathTest" + testContainer.Parent = workspace + + local child = Instance.new("ModuleScript") + child.Name = "MyModule" + child.Parent = testContainer + + local containerId = HttpService:GenerateGUID(false) + instanceMap:insert(containerId, testContainer) + + local childId = HttpService:GenerateGUID(false) + instanceMap:insert(childId, child) + + table.insert(patch.updated, { + id = childId, + changedProperties = { Source = "return {}" }, + }) + + local tree = PatchTree.build(patch, instanceMap, { "Property", "Current", "Incoming" }) + + -- Walk the ancestry of the child node to build a path + local childNode = tree:getNode(childId) + expect(childNode).to.be.ok() + + local segments = {} + local current = childNode + while current and current.id ~= "ROOT" do + table.insert(segments, 1, current.name or current.id) + if current.parentId then + current = tree:getNode(current.parentId) + else + break + end + end + + local path = table.concat(segments, "/") + + -- Path should contain both the parent folder and the child name + expect(string.find(path, "PathTest")).to.be.ok() + expect(string.find(path, "MyModule")).to.be.ok() + + testContainer:Destroy() + instanceMap:stop() + end) + + it("should produce single-segment path for root-level nodes", function() + local patch = PatchSet.newEmpty() + local instanceMap = InstanceMap.new() + + -- Use a service-level instance (direct child of DataModel) + local repStorage = game:GetService("ReplicatedStorage") + local testObj = Instance.new("Folder") + testObj.Name = "TopLevel" + testObj.Parent = repStorage + + local objId = HttpService:GenerateGUID(false) + instanceMap:insert(objId, testObj) + + table.insert(patch.updated, { + id = objId, + changedName = "Renamed", + changedProperties = {}, + }) + + local tree = PatchTree.build(patch, instanceMap, { "Property", "Current", "Incoming" }) + + local node = tree:getNode(objId) + expect(node).to.be.ok() + + local segments = {} + local current = node + while current and current.id ~= "ROOT" do + table.insert(segments, 1, current.name or current.id) + if current.parentId then + current = tree:getNode(current.parentId) + else + break + end + end + + local path = table.concat(segments, "/") + -- Should include the service ancestor and the object + expect(string.find(path, "ReplicatedStorage")).to.be.ok() + expect(string.find(path, "TopLevel")).to.be.ok() + + testObj:Destroy() + instanceMap:stop() + end) + end) end diff --git a/src/web/api.rs b/src/web/api.rs index 7b77a7aa..22583ff3 100644 --- a/src/web/api.rs +++ b/src/web/api.rs @@ -125,6 +125,8 @@ pub async fn call( serve_session: Arc, mut request: Request, syncback_signal: Arc, + mcp_state: Arc, + active_api_connections: Arc, ) -> Response> { let service = ApiService::new(serve_session); @@ -133,9 +135,23 @@ pub async fn call( (&Method::GET, path) if path.starts_with("/api/read/") => { service.handle_api_read(request).await } + (&Method::GET, "/api/mcp/stream") => { + if is_upgrade_request(&request) { + handle_mcp_stream_upgrade(&mut request, mcp_state, active_api_connections).await + } else { + msgpack( + ErrorResponse::bad_request( + "/api/mcp/stream must be called as a websocket upgrade request", + ), + StatusCode::BAD_REQUEST, + ) + } + } (&Method::GET, path) if path.starts_with("/api/socket/") => { if is_upgrade_request(&request) { - service.handle_api_socket(&mut request).await + service + .handle_api_socket(&mut request, Arc::clone(&active_api_connections)) + .await } else { msgpack( ErrorResponse::bad_request( @@ -426,7 +442,11 @@ impl ApiService { } /// Handle WebSocket upgrade for real-time message streaming - async fn handle_api_socket(&self, request: &mut Request) -> Response> { + async fn handle_api_socket( + &self, + request: &mut Request, + active_api_connections: Arc, + ) -> Response> { let argument = &request.uri().path()["/api/socket/".len()..]; let input_cursor: u32 = match argument.parse() { Ok(v) => v, @@ -453,9 +473,11 @@ impl ApiService { // Spawn a task to handle the WebSocket connection tokio::spawn(async move { - if let Err(e) = - handle_websocket_subscription(serve_session, websocket, input_cursor).await - { + active_api_connections.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + let result = + handle_websocket_subscription(serve_session, websocket, input_cursor).await; + active_api_connections.fetch_sub(1, std::sync::atomic::Ordering::SeqCst); + if let Err(e) = result { log::error!("Error in websocket subscription: {}", e); } }); @@ -3490,6 +3512,191 @@ fn pick_script_path(instance: InstanceWithMeta<'_>) -> Option { .map(|path| path.to_owned()) } +// --------------------------------------------------------------------------- +// MCP stream WebSocket: plugin connects here to receive sync commands +// --------------------------------------------------------------------------- + +async fn handle_mcp_stream_upgrade( + request: &mut Request, + mcp_state: Arc, + active_api_connections: Arc, +) -> Response> { + use std::sync::atomic::Ordering; + + if mcp_state.plugin_stream_connected.load(Ordering::Relaxed) { + return msgpack( + ErrorResponse::bad_request("Another plugin is already connected to the MCP stream"), + StatusCode::CONFLICT, + ); + } + + let (response, websocket) = match upgrade(request, None) { + Ok(result) => result, + Err(err) => { + return msgpack( + ErrorResponse::internal_error(format!("WebSocket upgrade failed: {err}")), + StatusCode::INTERNAL_SERVER_ERROR, + ); + } + }; + + tokio::spawn(async move { + if let Err(e) = + handle_mcp_stream_connection(websocket, mcp_state, active_api_connections).await + { + log::error!("MCP stream WebSocket error: {e}"); + } + }); + + response +} + +async fn handle_mcp_stream_connection( + websocket: HyperWebsocket, + mcp_state: Arc, + active_api_connections: Arc, +) -> anyhow::Result<()> { + use super::mcp::{McpSyncResult, PluginConfig}; + use std::sync::atomic::Ordering; + + let mut websocket = websocket.await?; + + mcp_state + .plugin_stream_connected + .store(true, Ordering::SeqCst); + log::info!("Plugin MCP stream connected"); + + let mut command_rx = mcp_state.command_rx.clone(); + + // Wait for the plugin greeting first (plugin sends its config on connect). + // Give it a generous timeout -- the greeting should arrive immediately. + let greeting_timeout = + tokio::time::timeout(std::time::Duration::from_secs(10), websocket.next()).await; + + if let Ok(Some(Ok(Message::Text(text)))) = greeting_timeout { + if let Ok(greeting) = serde_json::from_str::(&text) { + if greeting.get("type").and_then(|t| t.as_str()) == Some("hello") { + if let Ok(config) = serde_json::from_value::(greeting.clone()) { + let mut slot = mcp_state + .plugin_config + .lock() + .unwrap_or_else(|e| e.into_inner()); + *slot = Some(config); + log::debug!("Received plugin config via MCP stream greeting"); + } + } + } + } + + loop { + tokio::select! { + // Watch for sync commands from the MCP handler. + result = command_rx.changed() => { + if result.is_err() { + break; + } + + let command = command_rx.borrow_and_update().clone(); + let Some(cmd) = command else { + continue; + }; + + // If we are in connected mode right now, immediately respond + // without forwarding to the plugin. + if active_api_connections.load(Ordering::Relaxed) > 0 { + let result = McpSyncResult { + request_id: cmd.request_id, + status: "already_connected".to_string(), + changes: vec![], + message: Some( + "Atlas is already connected in live sync mode.".to_string(), + ), + }; + let mut slot = mcp_state + .result_tx + .lock() + .unwrap_or_else(|e| e.into_inner()); + if let Some(tx) = slot.take() { + let _ = tx.send(result); + } + continue; + } + + let json_cmd = serde_json::to_string(&cmd)?; + if websocket.send(Message::Text(json_cmd.into())).await.is_err() { + log::debug!("MCP stream: plugin disconnected while sending command"); + break; + } + } + + // Handle incoming messages from the plugin. + msg = websocket.next() => { + match msg { + Some(Ok(Message::Text(text))) => { + if let Ok(result) = serde_json::from_str::(&text) { + let mut slot = mcp_state + .result_tx + .lock() + .unwrap_or_else(|e| e.into_inner()); + if let Some(tx) = slot.take() { + let _ = tx.send(result); + } + } else { + log::debug!("MCP stream: ignoring non-result message from plugin"); + } + } + Some(Ok(Message::Close(_))) | None => { + log::info!("Plugin MCP stream disconnected"); + break; + } + Some(Ok(Message::Ping(_))) | Some(Ok(Message::Pong(_))) => {} + Some(Ok(Message::Binary(_))) => { + log::debug!("MCP stream: ignoring binary message from plugin"); + } + Some(Ok(Message::Frame(_))) => { + unreachable!(); + } + Some(Err(e)) => { + log::error!("MCP stream WebSocket error: {e}"); + break; + } + } + } + } + } + + // Clean up: mark disconnected, cancel any pending sync. + mcp_state + .plugin_stream_connected + .store(false, Ordering::SeqCst); + { + let mut slot = mcp_state + .plugin_config + .lock() + .unwrap_or_else(|e| e.into_inner()); + *slot = None; + } + + // If there's a pending sync request, send a disconnect error. + { + let mut slot = mcp_state + .result_tx + .lock() + .unwrap_or_else(|e| e.into_inner()); + if let Some(tx) = slot.take() { + let _ = tx.send(McpSyncResult { + request_id: String::new(), + status: "error".to_string(), + changes: vec![], + message: Some("Plugin MCP stream disconnected.".to_string()), + }); + } + } + + log::info!("Plugin MCP stream handler exiting"); + Ok(()) +} + /// Handle WebSocket connection for streaming subscription messages async fn handle_websocket_subscription( serve_session: Arc, diff --git a/src/web/mcp.rs b/src/web/mcp.rs new file mode 100644 index 00000000..d5a5a17f --- /dev/null +++ b/src/web/mcp.rs @@ -0,0 +1,855 @@ +use std::sync::{ + atomic::{AtomicBool, Ordering}, + Arc, Mutex, +}; + +use bytes::Bytes; +use http_body_util::{BodyExt, Full}; +use hyper::{body::Incoming, header::CONTENT_TYPE, Method, Request, Response, StatusCode}; +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +use super::interface::SERVER_VERSION; + +/// Plugin config received via the MCP stream WebSocket greeting. +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct PluginConfig { + #[serde(default)] + pub two_way_sync: bool, + #[serde(default)] + pub one_shot_sync: bool, + #[serde(default)] + pub confirmation_behavior: String, +} + +/// Command sent from the MCP handler to the plugin via the MCP stream WebSocket. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct McpSyncCommand { + #[serde(rename = "type")] + pub command_type: String, + pub request_id: String, +} + +/// A single change entry in the sync result. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct SyncChange { + pub path: String, + pub direction: String, +} + +/// Result sent from the plugin back through the MCP stream WebSocket. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct McpSyncResult { + pub request_id: String, + pub status: String, + #[serde(default)] + pub changes: Vec, + #[serde(default)] + pub message: Option, +} + +/// Shared state for MCP sync coordination between the agent-facing `/mcp` +/// endpoint and the plugin-facing `/api/mcp/stream` WebSocket. +pub struct McpSyncState { + pub sync_in_progress: AtomicBool, + pub command_tx: tokio::sync::watch::Sender>, + pub command_rx: tokio::sync::watch::Receiver>, + pub result_tx: Mutex>>, + pub plugin_stream_connected: AtomicBool, + pub plugin_config: Mutex>, +} + +impl McpSyncState { + pub fn new() -> Self { + let (command_tx, command_rx) = tokio::sync::watch::channel(None); + Self { + sync_in_progress: AtomicBool::new(false), + command_tx, + command_rx, + result_tx: Mutex::new(None), + plugin_stream_connected: AtomicBool::new(false), + plugin_config: Mutex::new(None), + } + } +} + +// --------------------------------------------------------------------------- +// JSON-RPC types (minimal subset for MCP Streamable HTTP) +// --------------------------------------------------------------------------- + +#[derive(Debug, Deserialize)] +struct JsonRpcRequest { + #[allow(dead_code)] + jsonrpc: String, + id: Option, + method: String, + #[serde(default)] + params: Option, +} + +#[derive(Debug, Serialize)] +struct JsonRpcResponse { + jsonrpc: &'static str, + #[serde(skip_serializing_if = "Option::is_none")] + id: Option, + #[serde(skip_serializing_if = "Option::is_none")] + result: Option, + #[serde(skip_serializing_if = "Option::is_none")] + error: Option, +} + +#[derive(Debug, Serialize)] +struct JsonRpcError { + code: i64, + message: String, + #[serde(skip_serializing_if = "Option::is_none")] + data: Option, +} + +impl JsonRpcResponse { + fn success(id: Option, result: Value) -> Self { + Self { + jsonrpc: "2.0", + id, + result: Some(result), + error: None, + } + } + + fn error(id: Option, code: i64, message: impl Into) -> Self { + Self { + jsonrpc: "2.0", + id, + result: None, + error: Some(JsonRpcError { + code, + message: message.into(), + data: None, + }), + } + } +} + +fn json_response(body: &JsonRpcResponse, status: StatusCode) -> Response> { + let serialized = serde_json::to_string(body).unwrap_or_default(); + Response::builder() + .status(status) + .header(CONTENT_TYPE, "application/json") + .body(Full::new(Bytes::from(serialized))) + .unwrap() +} + +// --------------------------------------------------------------------------- +// MCP HTTP endpoint: POST /mcp +// --------------------------------------------------------------------------- + +pub async fn call( + request: Request, + mcp_state: Arc, + active_api_connections: Arc, +) -> Response> { + if request.method() != Method::POST { + let resp = JsonRpcResponse::error(None, -32600, "MCP endpoint only accepts POST requests"); + return json_response(&resp, StatusCode::METHOD_NOT_ALLOWED); + } + + let body = match request.into_body().collect().await { + Ok(collected) => collected.to_bytes(), + Err(err) => { + let resp = + JsonRpcResponse::error(None, -32700, format!("Failed to read request body: {err}")); + return json_response(&resp, StatusCode::BAD_REQUEST); + } + }; + + let rpc_request: JsonRpcRequest = match serde_json::from_slice(&body) { + Ok(req) => req, + Err(err) => { + let resp = JsonRpcResponse::error(None, -32700, format!("Parse error: {err}")); + return json_response(&resp, StatusCode::BAD_REQUEST); + } + }; + + match rpc_request.method.as_str() { + "initialize" => handle_initialize(rpc_request.id), + "notifications/initialized" => Response::builder() + .status(StatusCode::ACCEPTED) + .body(Full::new(Bytes::new())) + .unwrap(), + "tools/list" => handle_tools_list(rpc_request.id), + "tools/call" => { + handle_tools_call( + rpc_request.id, + rpc_request.params, + mcp_state, + active_api_connections, + ) + .await + } + _ => { + let resp = JsonRpcResponse::error( + rpc_request.id, + -32601, + format!("Method not found: {}", rpc_request.method), + ); + json_response(&resp, StatusCode::OK) + } + } +} + +// --------------------------------------------------------------------------- +// MCP method handlers +// --------------------------------------------------------------------------- + +fn handle_initialize(id: Option) -> Response> { + let result = serde_json::json!({ + "protocolVersion": "2024-11-05", + "capabilities": { + "tools": { + "listChanged": false + } + }, + "serverInfo": { + "name": "atlas", + "version": SERVER_VERSION + } + }); + let resp = JsonRpcResponse::success(id, result); + json_response(&resp, StatusCode::OK) +} + +fn handle_tools_list(id: Option) -> Response> { + let result = serde_json::json!({ + "tools": [ + { + "name": "atlas_sync", + "description": "Sync filesystem changes to Roblox Studio. Equivalent to clicking the Atlas sync button. If all changes are pre-selected by the git system, they are auto-accepted (fast-forward). Otherwise, the user is prompted to review changes in Studio and the call blocks until they accept or reject.", + "inputSchema": { + "type": "object", + "properties": {}, + "required": [] + } + } + ] + }); + let resp = JsonRpcResponse::success(id, result); + json_response(&resp, StatusCode::OK) +} + +async fn handle_tools_call( + id: Option, + params: Option, + mcp_state: Arc, + active_api_connections: Arc, +) -> Response> { + let tool_name = params + .as_ref() + .and_then(|p| p.get("name")) + .and_then(|n| n.as_str()) + .unwrap_or(""); + + if tool_name != "atlas_sync" { + let resp = JsonRpcResponse::error(id, -32602, format!("Unknown tool: {tool_name}")); + return json_response(&resp, StatusCode::OK); + } + + handle_atlas_sync(id, mcp_state, active_api_connections).await +} + +async fn handle_atlas_sync( + id: Option, + mcp_state: Arc, + active_api_connections: Arc, +) -> Response> { + // Check if the plugin is connected via the regular API WebSocket (connected mode). + if active_api_connections.load(Ordering::Relaxed) > 0 { + let config_info = { + let config = mcp_state + .plugin_config + .lock() + .unwrap_or_else(|e| e.into_inner()); + match config.as_ref() { + Some(cfg) => format!( + " Plugin config: twoWaySync={}, oneShotSync={}, confirmationBehavior={}.", + cfg.two_way_sync, cfg.one_shot_sync, cfg.confirmation_behavior + ), + None => String::new(), + } + }; + + return tool_response( + id, + true, + &format!( + "Atlas is already connected to Studio in live sync mode. \ + All changes are syncing automatically. No manual sync needed.{config_info}" + ), + ); + } + + // Check if a sync is already in progress (from another agent invocation). + if mcp_state + .sync_in_progress + .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst) + .is_err() + { + return tool_response( + id, + true, + "An atlas_sync operation is already in progress. Please wait for it to complete.", + ); + } + + // Check if the plugin MCP stream is connected. + if !mcp_state.plugin_stream_connected.load(Ordering::Relaxed) { + mcp_state.sync_in_progress.store(false, Ordering::SeqCst); + return tool_response( + id, + true, + "No Roblox Studio plugin is connected to the MCP stream. \ + Make sure Studio is open with the Atlas plugin installed.", + ); + } + + // Create a oneshot channel for receiving the result from the plugin. + let (result_tx, result_rx) = tokio::sync::oneshot::channel(); + { + let mut slot = mcp_state + .result_tx + .lock() + .unwrap_or_else(|e| e.into_inner()); + *slot = Some(result_tx); + } + + let request_id = uuid::Uuid::new_v4().to_string(); + let command = McpSyncCommand { + command_type: "sync".to_string(), + request_id: request_id.clone(), + }; + + // Send the sync command to the plugin via the watch channel. + if mcp_state.command_tx.send(Some(command)).is_err() { + mcp_state.sync_in_progress.store(false, Ordering::SeqCst); + return tool_response(id, true, "Failed to send sync command to the plugin."); + } + + // Await the plugin's result (blocks until plugin responds or disconnects). + let result = match result_rx.await { + Ok(result) => result, + Err(_) => { + mcp_state.sync_in_progress.store(false, Ordering::SeqCst); + return tool_response( + id, + true, + "Plugin disconnected before completing the sync operation.", + ); + } + }; + + mcp_state.sync_in_progress.store(false, Ordering::SeqCst); + + // Clear the command so the watch channel doesn't re-fire on reconnect. + let _ = mcp_state.command_tx.send(None); + + let is_error = result.status != "success" && result.status != "empty"; + + let mut text = match result.status.as_str() { + "success" => "Sync completed successfully.".to_string(), + "empty" => "No changes to sync. Studio is already up to date.".to_string(), + "rejected" => "Sync was rejected by the user.".to_string(), + "already_connected" => "Atlas is already connected to Studio in live sync mode. \ + All changes are syncing automatically." + .to_string(), + "sync_in_progress" => "A sync operation is already in progress.".to_string(), + other => format!("Sync finished with status: {other}"), + }; + + if let Some(msg) = &result.message { + text.push_str(&format!("\n\n{msg}")); + } + + if !result.changes.is_empty() { + let label = if is_error { + "Presented changes" + } else { + "Changes" + }; + text.push_str(&format!("\n\n{label}:")); + for change in &result.changes { + text.push_str(&format!("\n- [{}] {}", change.direction, change.path)); + } + } + + tool_response(id, is_error, &text) +} + +fn tool_response(id: Option, is_error: bool, text: &str) -> Response> { + let result = serde_json::json!({ + "content": [{ "type": "text", "text": text }], + "isError": is_error + }); + let resp = JsonRpcResponse::success(id, result); + json_response(&resp, StatusCode::OK) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::atomic::AtomicUsize; + + // -- JSON-RPC type tests -------------------------------------------------- + + mod jsonrpc_types { + use super::*; + + #[test] + fn success_response_serializes_correctly() { + let resp = + JsonRpcResponse::success(Some(Value::from(1)), serde_json::json!({"ok": true})); + let json = serde_json::to_value(&resp).unwrap(); + assert_eq!(json["jsonrpc"], "2.0"); + assert_eq!(json["id"], 1); + assert!(json["result"]["ok"].as_bool().unwrap()); + assert!(json.get("error").is_none()); + } + + #[test] + fn error_response_serializes_correctly() { + let resp = JsonRpcResponse::error(Some(Value::from(2)), -32600, "bad request"); + let json = serde_json::to_value(&resp).unwrap(); + assert_eq!(json["jsonrpc"], "2.0"); + assert_eq!(json["id"], 2); + assert!(json.get("result").is_none()); + assert_eq!(json["error"]["code"], -32600); + assert_eq!(json["error"]["message"], "bad request"); + } + + #[test] + fn null_id_omitted_in_success() { + let resp = JsonRpcResponse::success(None, serde_json::json!({})); + let json = serde_json::to_value(&resp).unwrap(); + assert!(json.get("id").is_none()); + } + } + + // -- Wire type serde tests ------------------------------------------------ + + mod wire_types { + use super::*; + + #[test] + fn mcp_sync_command_serializes_camel_case() { + let cmd = McpSyncCommand { + command_type: "sync".to_string(), + request_id: "abc-123".to_string(), + }; + let json = serde_json::to_value(&cmd).unwrap(); + assert_eq!(json["type"], "sync"); + assert_eq!(json["requestId"], "abc-123"); + } + + #[test] + fn mcp_sync_result_deserializes_camel_case() { + let json = serde_json::json!({ + "requestId": "abc-123", + "status": "success", + "changes": [ + { "path": "ReplicatedStorage/Module.luau", "direction": "push" } + ], + "message": "done" + }); + let result: McpSyncResult = serde_json::from_value(json).unwrap(); + assert_eq!(result.request_id, "abc-123"); + assert_eq!(result.status, "success"); + assert_eq!(result.changes.len(), 1); + assert_eq!(result.changes[0].path, "ReplicatedStorage/Module.luau"); + assert_eq!(result.changes[0].direction, "push"); + assert_eq!(result.message.as_deref(), Some("done")); + } + + #[test] + fn mcp_sync_result_defaults_for_missing_optional_fields() { + let json = serde_json::json!({ + "requestId": "x", + "status": "empty" + }); + let result: McpSyncResult = serde_json::from_value(json).unwrap(); + assert!(result.changes.is_empty()); + assert!(result.message.is_none()); + } + + #[test] + fn plugin_config_deserializes_with_defaults() { + let json = serde_json::json!({}); + let cfg: PluginConfig = serde_json::from_value(json).unwrap(); + assert!(!cfg.two_way_sync); + assert!(!cfg.one_shot_sync); + assert_eq!(cfg.confirmation_behavior, ""); + } + + #[test] + fn plugin_config_roundtrip() { + let cfg = PluginConfig { + two_way_sync: true, + one_shot_sync: false, + confirmation_behavior: "Always".to_string(), + }; + let json = serde_json::to_value(&cfg).unwrap(); + assert_eq!(json["twoWaySync"], true); + assert_eq!(json["oneShotSync"], false); + assert_eq!(json["confirmationBehavior"], "Always"); + } + } + + // -- McpSyncState tests --------------------------------------------------- + + mod state_tests { + use super::*; + + #[test] + fn new_state_starts_clean() { + let state = McpSyncState::new(); + assert!(!state.sync_in_progress.load(Ordering::Relaxed)); + assert!(!state.plugin_stream_connected.load(Ordering::Relaxed)); + assert!(state.plugin_config.lock().unwrap().is_none()); + assert!(state.result_tx.lock().unwrap().is_none()); + assert!(state.command_rx.borrow().is_none()); + } + + #[test] + fn sync_mutex_prevents_double_acquire() { + let state = McpSyncState::new(); + let first = state.sync_in_progress.compare_exchange( + false, + true, + Ordering::SeqCst, + Ordering::SeqCst, + ); + assert!(first.is_ok()); + let second = state.sync_in_progress.compare_exchange( + false, + true, + Ordering::SeqCst, + Ordering::SeqCst, + ); + assert!(second.is_err()); + } + + #[test] + fn command_channel_delivers() { + let state = McpSyncState::new(); + let mut rx = state.command_rx.clone(); + + let cmd = McpSyncCommand { + command_type: "sync".to_string(), + request_id: "test-1".to_string(), + }; + state.command_tx.send(Some(cmd)).unwrap(); + + assert!(rx.has_changed().unwrap()); + let received = rx.borrow_and_update().clone().unwrap(); + assert_eq!(received.request_id, "test-1"); + } + + #[test] + fn result_oneshot_delivers() { + let state = McpSyncState::new(); + let (tx, rx) = tokio::sync::oneshot::channel(); + *state.result_tx.lock().unwrap() = Some(tx); + + let result = McpSyncResult { + request_id: "r1".to_string(), + status: "success".to_string(), + changes: vec![], + message: None, + }; + let sender = state.result_tx.lock().unwrap().take().unwrap(); + sender.send(result).unwrap(); + + let received = rx.blocking_recv().unwrap(); + assert_eq!(received.status, "success"); + } + + #[test] + fn plugin_config_cache() { + let state = McpSyncState::new(); + let cfg = PluginConfig { + two_way_sync: true, + one_shot_sync: false, + confirmation_behavior: "Initial".to_string(), + }; + *state.plugin_config.lock().unwrap() = Some(cfg); + + let cached = state.plugin_config.lock().unwrap(); + assert!(cached.as_ref().unwrap().two_way_sync); + assert_eq!(cached.as_ref().unwrap().confirmation_behavior, "Initial"); + } + } + + // -- Handler response tests (pure functions) ------------------------------ + + mod handler_tests { + use super::*; + + #[test] + fn initialize_returns_capabilities() { + let resp = handle_initialize(Some(Value::from(1))); + assert_eq!(resp.status(), StatusCode::OK); + let body = resp.into_body(); + let rt = tokio::runtime::Runtime::new().unwrap(); + let bytes = rt.block_on(async { body.collect().await.unwrap().to_bytes() }); + let json: Value = serde_json::from_slice(&bytes).unwrap(); + assert_eq!(json["jsonrpc"], "2.0"); + assert_eq!(json["id"], 1); + assert_eq!(json["result"]["serverInfo"]["name"], "atlas"); + assert!(json["result"]["capabilities"]["tools"].is_object()); + } + + #[test] + fn tools_list_returns_atlas_sync() { + let resp = handle_tools_list(Some(Value::from(2))); + assert_eq!(resp.status(), StatusCode::OK); + let rt = tokio::runtime::Runtime::new().unwrap(); + let bytes = rt.block_on(async { resp.into_body().collect().await.unwrap().to_bytes() }); + let json: Value = serde_json::from_slice(&bytes).unwrap(); + let tools = json["result"]["tools"].as_array().unwrap(); + assert_eq!(tools.len(), 1); + assert_eq!(tools[0]["name"], "atlas_sync"); + assert!(tools[0]["inputSchema"].is_object()); + } + + #[test] + fn tool_response_formats_correctly() { + let resp = tool_response(Some(Value::from(3)), false, "All good"); + assert_eq!(resp.status(), StatusCode::OK); + let rt = tokio::runtime::Runtime::new().unwrap(); + let bytes = rt.block_on(async { resp.into_body().collect().await.unwrap().to_bytes() }); + let json: Value = serde_json::from_slice(&bytes).unwrap(); + assert_eq!(json["result"]["isError"], false); + let content = json["result"]["content"].as_array().unwrap(); + assert_eq!(content[0]["type"], "text"); + assert_eq!(content[0]["text"], "All good"); + } + + #[test] + fn tool_response_error_flag() { + let resp = tool_response(Some(Value::from(4)), true, "Oops"); + let rt = tokio::runtime::Runtime::new().unwrap(); + let bytes = rt.block_on(async { resp.into_body().collect().await.unwrap().to_bytes() }); + let json: Value = serde_json::from_slice(&bytes).unwrap(); + assert_eq!(json["result"]["isError"], true); + } + } + + // -- atlas_sync guard tests (via handle_atlas_sync) ----------------------- + + mod atlas_sync_guards { + use super::*; + + #[tokio::test] + async fn rejects_when_api_connected() { + let state = Arc::new(McpSyncState::new()); + let conns = Arc::new(AtomicUsize::new(1)); + + let resp = handle_atlas_sync(Some(Value::from(1)), state, conns).await; + let bytes = resp.into_body().collect().await.unwrap().to_bytes(); + let json: Value = serde_json::from_slice(&bytes).unwrap(); + + assert_eq!(json["result"]["isError"], true); + let text = json["result"]["content"][0]["text"].as_str().unwrap(); + assert!(text.contains("already connected")); + } + + #[tokio::test] + async fn rejects_when_api_connected_includes_config() { + let state = Arc::new(McpSyncState::new()); + *state.plugin_config.lock().unwrap() = Some(PluginConfig { + two_way_sync: true, + one_shot_sync: false, + confirmation_behavior: "Always".to_string(), + }); + let conns = Arc::new(AtomicUsize::new(1)); + + let resp = handle_atlas_sync(Some(Value::from(1)), state, conns).await; + let bytes = resp.into_body().collect().await.unwrap().to_bytes(); + let json: Value = serde_json::from_slice(&bytes).unwrap(); + + let text = json["result"]["content"][0]["text"].as_str().unwrap(); + assert!(text.contains("twoWaySync=true")); + assert!(text.contains("confirmationBehavior=Always")); + } + + #[tokio::test] + async fn rejects_when_sync_in_progress() { + let state = Arc::new(McpSyncState::new()); + state.sync_in_progress.store(true, Ordering::SeqCst); + let conns = Arc::new(AtomicUsize::new(0)); + + let resp = handle_atlas_sync(Some(Value::from(2)), state, conns).await; + let bytes = resp.into_body().collect().await.unwrap().to_bytes(); + let json: Value = serde_json::from_slice(&bytes).unwrap(); + + assert_eq!(json["result"]["isError"], true); + let text = json["result"]["content"][0]["text"].as_str().unwrap(); + assert!(text.contains("already in progress")); + } + + #[tokio::test] + async fn rejects_when_no_plugin_connected() { + let state = Arc::new(McpSyncState::new()); + state.plugin_stream_connected.store(false, Ordering::SeqCst); + let conns = Arc::new(AtomicUsize::new(0)); + + let resp = handle_atlas_sync(Some(Value::from(3)), state.clone(), conns).await; + let bytes = resp.into_body().collect().await.unwrap().to_bytes(); + let json: Value = serde_json::from_slice(&bytes).unwrap(); + + assert_eq!(json["result"]["isError"], true); + let text = json["result"]["content"][0]["text"].as_str().unwrap(); + assert!(text.contains("No Roblox Studio plugin")); + // Mutex should be released after rejection. + assert!(!state.sync_in_progress.load(Ordering::Relaxed)); + } + + #[tokio::test] + async fn sends_command_and_returns_success_result() { + let state = Arc::new(McpSyncState::new()); + state.plugin_stream_connected.store(true, Ordering::SeqCst); + let conns = Arc::new(AtomicUsize::new(0)); + + let state2 = Arc::clone(&state); + // Spawn a task that simulates the plugin responding. + tokio::spawn(async move { + let mut rx = state2.command_rx.clone(); + rx.changed().await.unwrap(); + let cmd = rx.borrow().clone().unwrap(); + + let result = McpSyncResult { + request_id: cmd.request_id, + status: "success".to_string(), + changes: vec![SyncChange { + path: "ServerScriptService/Main.server.luau".to_string(), + direction: "push".to_string(), + }], + message: None, + }; + + let tx = state2.result_tx.lock().unwrap().take().unwrap(); + tx.send(result).unwrap(); + }); + + let resp = handle_atlas_sync(Some(Value::from(5)), state.clone(), conns).await; + let bytes = resp.into_body().collect().await.unwrap().to_bytes(); + let json: Value = serde_json::from_slice(&bytes).unwrap(); + + assert_eq!(json["result"]["isError"], false); + let text = json["result"]["content"][0]["text"].as_str().unwrap(); + assert!(text.contains("Sync completed successfully")); + assert!(text.contains("[push] ServerScriptService/Main.server.luau")); + assert!(!state.sync_in_progress.load(Ordering::Relaxed)); + } + + #[tokio::test] + async fn returns_rejection_with_presented_changes() { + let state = Arc::new(McpSyncState::new()); + state.plugin_stream_connected.store(true, Ordering::SeqCst); + let conns = Arc::new(AtomicUsize::new(0)); + + let state2 = Arc::clone(&state); + tokio::spawn(async move { + let mut rx = state2.command_rx.clone(); + rx.changed().await.unwrap(); + let cmd = rx.borrow().clone().unwrap(); + + let result = McpSyncResult { + request_id: cmd.request_id, + status: "rejected".to_string(), + changes: vec![ + SyncChange { + path: "ReplicatedStorage/Foo.luau".to_string(), + direction: "push".to_string(), + }, + SyncChange { + path: "Workspace/Bar.model.json5".to_string(), + direction: "pull".to_string(), + }, + ], + message: Some("User rejected the sync changes.".to_string()), + }; + + let tx = state2.result_tx.lock().unwrap().take().unwrap(); + tx.send(result).unwrap(); + }); + + let resp = handle_atlas_sync(Some(Value::from(6)), state, conns).await; + let bytes = resp.into_body().collect().await.unwrap().to_bytes(); + let json: Value = serde_json::from_slice(&bytes).unwrap(); + + assert_eq!(json["result"]["isError"], true); + let text = json["result"]["content"][0]["text"].as_str().unwrap(); + assert!(text.contains("rejected")); + assert!(text.contains("Presented changes")); + assert!(text.contains("[push] ReplicatedStorage/Foo.luau")); + assert!(text.contains("[pull] Workspace/Bar.model.json5")); + } + + #[tokio::test] + async fn handles_plugin_disconnect_during_sync() { + let state = Arc::new(McpSyncState::new()); + state.plugin_stream_connected.store(true, Ordering::SeqCst); + let conns = Arc::new(AtomicUsize::new(0)); + + let state2 = Arc::clone(&state); + tokio::spawn(async move { + let mut rx = state2.command_rx.clone(); + rx.changed().await.unwrap(); + // Drop the sender without sending a result (simulates disconnect). + let _tx = state2.result_tx.lock().unwrap().take(); + }); + + let resp = handle_atlas_sync(Some(Value::from(7)), state.clone(), conns).await; + let bytes = resp.into_body().collect().await.unwrap().to_bytes(); + let json: Value = serde_json::from_slice(&bytes).unwrap(); + + assert_eq!(json["result"]["isError"], true); + let text = json["result"]["content"][0]["text"].as_str().unwrap(); + assert!(text.contains("disconnected")); + assert!(!state.sync_in_progress.load(Ordering::Relaxed)); + } + + #[tokio::test] + async fn empty_status_is_not_an_error() { + let state = Arc::new(McpSyncState::new()); + state.plugin_stream_connected.store(true, Ordering::SeqCst); + let conns = Arc::new(AtomicUsize::new(0)); + + let state2 = Arc::clone(&state); + tokio::spawn(async move { + let mut rx = state2.command_rx.clone(); + rx.changed().await.unwrap(); + let cmd = rx.borrow().clone().unwrap(); + + let tx = state2.result_tx.lock().unwrap().take().unwrap(); + tx.send(McpSyncResult { + request_id: cmd.request_id, + status: "empty".to_string(), + changes: vec![], + message: None, + }) + .unwrap(); + }); + + let resp = handle_atlas_sync(Some(Value::from(8)), state, conns).await; + let bytes = resp.into_body().collect().await.unwrap().to_bytes(); + let json: Value = serde_json::from_slice(&bytes).unwrap(); + + assert_eq!(json["result"]["isError"], false); + let text = json["result"]["content"][0]["text"].as_str().unwrap(); + assert!(text.contains("No changes to sync")); + } + } +} diff --git a/src/web/mod.rs b/src/web/mod.rs index 9086f213..884f5de8 100644 --- a/src/web/mod.rs +++ b/src/web/mod.rs @@ -5,11 +5,13 @@ mod api; mod assets; pub mod interface; +pub mod mcp; mod ui; mod util; use std::convert::Infallible; use std::net::SocketAddr; +use std::sync::atomic::AtomicUsize; use std::sync::{Arc, Mutex}; use std::time::Duration; @@ -62,6 +64,8 @@ impl SyncbackSignal { pub struct LiveServer { serve_session: Arc, syncback_signal: Arc, + mcp_state: Arc, + active_api_connections: Arc, } impl LiveServer { @@ -69,12 +73,16 @@ impl LiveServer { LiveServer { serve_session, syncback_signal: Arc::new(SyncbackSignal::new()), + mcp_state: Arc::new(mcp::McpSyncState::new()), + active_api_connections: Arc::new(AtomicUsize::new(0)), } } pub fn start(self, address: SocketAddr) -> ServerExitReason { let serve_session = Arc::clone(&self.serve_session); let syncback_signal = Arc::clone(&self.syncback_signal); + let mcp_state = Arc::clone(&self.mcp_state); + let active_api_connections = Arc::clone(&self.active_api_connections); let rt = Runtime::new().unwrap(); let exit_reason = rt.block_on(async move { @@ -117,16 +125,32 @@ impl LiveServer { let io = TokioIo::new(stream); let serve_session = Arc::clone(&serve_session); let syncback_signal = Arc::clone(&syncback_signal); + let mcp_state = Arc::clone(&mcp_state); + let active_api_connections = Arc::clone(&active_api_connections); tokio::spawn(async move { let service = service_fn(move |req: Request| { let serve_session = Arc::clone(&serve_session); let syncback_signal = Arc::clone(&syncback_signal); + let mcp_state = Arc::clone(&mcp_state); + let active_api_connections = Arc::clone(&active_api_connections); async move { - if req.uri().path().starts_with("/api") { + if req.uri().path().starts_with("/mcp") { Ok::<_, Infallible>( - api::call(serve_session, req, syncback_signal).await, + mcp::call(req, mcp_state, active_api_connections) + .await, + ) + } else if req.uri().path().starts_with("/api") { + Ok::<_, Infallible>( + api::call( + serve_session, + req, + syncback_signal, + mcp_state, + active_api_connections, + ) + .await, ) } else { Ok::<_, Infallible>(ui::call(serve_session, req).await)