Enhance session management and UI state handling

- Added checks to handle session disconnections during confirmation in ServeSession, preventing further actions on a disconnected session.
- Updated App:startSession to manage UI state in one-shot sync mode, ensuring a loading state is displayed instead of a connected page.
- Implemented safety measures in App:endSession to reset UI state if the session reference is cleared, addressing potential race conditions.

plugin/src/App/init.lua

@@ -663,15 +663,38 @@ function App:startSession()
 			self:setRunningConnectionInfo(baseUrl)
 
 			local address = ("%s:%s"):format(host, port)
-			self:setState({
+
-				appStatus = AppStatus.Connected,
+			if Settings:get("oneShotSync") then
-				projectName = details,
+				-- One-shot mode: Don't show Connected page since we're about to disconnect.
-				address = address,
+				-- Show a loading state while the patch applies and writes complete.
-				toolbarIcon = Assets.Images.PluginButtonConnected,
+				self:setState({
-			})
+					appStatus = AppStatus.Connecting,
-			self:addNotification({
+					connectingText = "Completing sync...",
-				text = string.format("Connected to session '%s' at %s.", details, address),
+					toolbarIcon = Assets.Images.PluginButton,
-			})
+				})
+				self:addNotification({
+					text = string.format("Synced with '%s'. Completing...", details),
+				})
+				-- Safety: if endSession hasn't fired within 30s, force disconnect.
+				-- This catches cases where the promise chain silently breaks
+				-- (e.g., write request hangs, promise never resolves).
+				task.delay(30, function()
+					if self.serveSession ~= nil and Settings:get("oneShotSync") then
+						Log.warn("One-shot sync safety timeout: forcing disconnect")
+						self:endSession()
+					end
+				end)
+			else
+				self:setState({
+					appStatus = AppStatus.Connected,
+					projectName = details,
+					address = address,
+					toolbarIcon = Assets.Images.PluginButtonConnected,
+				})
+				self:addNotification({
+					text = string.format("Connected to session '%s' at %s.", details, address),
+				})
+			end
 		elseif status == ServeSession.Status.Disconnected then
 			self.serveSession = nil
 			-- Only release lock if we claimed it (not in one-shot mode)
@@ -863,6 +886,19 @@ end
 
 function App:endSession()
 	if self.serveSession == nil then
+		-- Safety: if the session reference was cleared (e.g., by a Disconnected handler)
+		-- but the UI is still stuck on Connected/Connecting, force reset to NotConnected.
+		-- This handles race conditions where onStatusChanged(Connected) fires after
+		-- onStatusChanged(Disconnected) due to a coroutine resuming on a dead session.
+		if self.state.appStatus ~= AppStatus.NotConnected
+			and self.state.appStatus ~= AppStatus.Error
+			and self.state.appStatus ~= AppStatus.Settings
+		then
+			self:setState({
+				appStatus = AppStatus.NotConnected,
+				toolbarIcon = Assets.Images.PluginButton,
+			})
+		end
 		return
 	end
 

plugin/src/ServeSession.lua

@@ -278,6 +278,12 @@ function ServeSession:__onWebSocketMessage(messagesPacket)
 		self.__confirmingPatch = nil
 		self.__isConfirming = false
 
+		-- If the session was disconnected during the confirmation yield
+		-- (e.g., WebSocket error, server crash), bail out immediately.
+		if self.__status == Status.Disconnected then
+			return
+		end
+
 		-- Resume ChangeBatcher after confirmation
 		self.__changeBatcher:resume()
 
@@ -656,6 +662,13 @@ function ServeSession:__confirmAndApplyInitialPatch(catchUpPatch, serverInfo)
 		self.__confirmingPatch = nil
 		self.__isConfirming = false
 
+		-- If the session was disconnected during the confirmation yield
+		-- (e.g., WebSocket error, server crash), bail out immediately.
+		-- Don't resume batcher or apply patches on a dead session.
+		if self.__status == Status.Disconnected then
+			return Promise.reject("Session disconnected during confirmation")
+		end
+
 		-- Resume ChangeBatcher after confirmation
 		self.__changeBatcher:resume()
 	end