From 249e26367b89d8c609c330e9afa29ad2308253fb Mon Sep 17 00:00:00 2001
From: Joe <joe.melsha@live.com>
Date: Fri, 30 Jan 2026 05:11:44 -0600
Subject: [PATCH] Add initial sync complete callback to ServeSession for
 one-shot mode handling

- Introduced `__initialSyncCompleteCallback` in `ServeSession` to manage actions after the initial sync is complete.
- Updated `App:startSession` to set the callback, ensuring proper disconnection after initial sync in one-shot mode.
- Refactored comments in `App:startSession` for clarity on sync behavior and disconnection logic.
---
 .../Components/PatchVisualizer/DomLabel.lua   |  1 -
 plugin/src/App/init.lua                       | 28 ++++++++-----------
 plugin/src/ServeSession.lua                   | 16 ++++++++---
 3 files changed, 24 insertions(+), 21 deletions(-)

diff --git a/plugin/src/App/Components/PatchVisualizer/DomLabel.lua b/plugin/src/App/Components/PatchVisualizer/DomLabel.lua
index 876df759..6c0f3134 100644
--- a/plugin/src/App/Components/PatchVisualizer/DomLabel.lua
+++ b/plugin/src/App/Components/PatchVisualizer/DomLabel.lua
@@ -14,7 +14,6 @@ local bindingUtil = require(Plugin.App.bindingUtil)
 local e = Roact.createElement
 
 local ChangeList = require(script.Parent.ChangeList)
-local Tooltip = require(Plugin.App.Components.Tooltip)
 local ClassIcon = require(Plugin.App.Components.ClassIcon)
 
 local function ChangeTag(props)
diff --git a/plugin/src/App/init.lua b/plugin/src/App/init.lua
index 1b37cfd6..508c330c 100644
--- a/plugin/src/App/init.lua
+++ b/plugin/src/App/init.lua
@@ -800,23 +800,10 @@ function App:startSession()
 
 		local result = self.confirmationEvent:Wait()
 
-		-- Track if this was the initial sync for this project
-		local isInitialSyncForProject = not self.knownProjects[serverInfo.projectName]
-
-		-- One-shot sync: disconnect immediately after confirmation
-		-- Skip the "Connected" UI state entirely to avoid flash
+		-- One-shot sync: Don't transition to Connected UI
+		-- The actual disconnect is handled by setInitialSyncCompleteCallback
+		-- to ensure any pending writes complete first
 		if Settings:get("oneShotSync") then
-			local resultType = if type(result) == "table" then result.type else result
-			if resultType == "Confirm" or resultType == "Accept" then
-				-- Use task.defer to disconnect after patch is applied (runs after current thread yields)
-				task.defer(function()
-					if self.serveSession then
-						Log.info("One-shot sync: disconnecting after sync")
-						self:endSession()
-					end
-				end)
-			end
-			-- Don't transition to Connected UI, just return result and let disconnect happen
 			return result
 		end
 
@@ -858,6 +845,15 @@ function App:startSession()
 		})
 	end)
 
+	-- One-shot sync: disconnect after initial sync is fully complete (including any writes)
+	-- This ensures that "pull" changes are sent to the server before disconnecting
+	serveSession:setInitialSyncCompleteCallback(function()
+		if Settings:get("oneShotSync") and self.serveSession then
+			Log.info("One-shot sync: disconnecting after initial sync complete")
+			self:endSession()
+		end
+	end)
+
 	serveSession:start()
 
 	self.serveSession = serveSession
diff --git a/plugin/src/ServeSession.lua b/plugin/src/ServeSession.lua
index 754b3afe..34c10a1c 100644
--- a/plugin/src/ServeSession.lua
+++ b/plugin/src/ServeSession.lua
@@ -116,6 +116,7 @@ function ServeSession.new(options)
 		__statusChangedCallback = nil,
 		__userConfirmCallback = nil,
 		__patchUpdateCallback = nil,
+		__initialSyncCompleteCallback = nil,
 		__serverInfo = nil,
 		__confirmingPatch = nil,
 		__connections = connections,
@@ -156,6 +157,10 @@ function ServeSession:setPatchUpdateCallback(callback)
 	self.__patchUpdateCallback = callback
 end
 
+function ServeSession:setInitialSyncCompleteCallback(callback)
+	self.__initialSyncCompleteCallback = callback
+end
+
 function ServeSession:setUpdateLoadingTextCallback(callback)
 	self.__updateLoadingText = callback
 end
@@ -306,6 +311,12 @@ function ServeSession:start()
 				return self:__confirmAndApplyInitialPatch(catchUpPatch, serverInfo)
 			end)
 		end)
+		:andThen(function()
+			-- Initial sync is fully complete (including any writes to the server)
+			if self.__initialSyncCompleteCallback ~= nil then
+				self.__initialSyncCompleteCallback()
+			end
+		end)
 		:catch(function(err)
 			if self.__status ~= Status.Disconnected then
 				self:__stopInternal(err)
@@ -658,11 +669,8 @@ function ServeSession:__confirmAndApplyInitialPatch(catchUpPatch, serverInfo)
 			if selection == "push" then
 				-- Apply Rojo removal to Studio
 				table.insert(pushPatch.removed, idOrInstance)
-			elseif selection == "pull" and self.__twoWaySync then
-				-- Don't remove in Studio, but we can't easily "add back" to Rojo
-				-- For now, just skip (similar to ignore)
 			end
-			-- "ignore" items are skipped
+			-- "pull" skipped: can't easily "add back" to Rojo, same as "ignore"
 		end
 
 		-- Process added items