remote: support takeover of existing session and auto-pause of orphaned sessions

Platform/Shared/VMDetailsView.swift

@@ -162,7 +162,7 @@ struct Screenshot: View {
                 .blendMode(.hardLight)
             #if os(visionOS)
                 .overlay {
-                    if vm.isStopped {
+                    if vm.isStopped || vm.isTakeoverAllowed {
                         Image(systemName: "play.circle.fill")
                             .resizable()
                             .frame(width: 100, height: 100)
@@ -175,7 +175,7 @@ struct Screenshot: View {
             #endif
             if vm.isBusy {
                 Spinner(size: .large)
-            } else if vm.isStopped {
+            } else if vm.isStopped || vm.isTakeoverAllowed {
                 #if !os(visionOS)
                 Button(action: { data.run(vm: vm) }, label: {
                     Label("Run", systemImage: "play.circle.fill")

Platform/UTMData.swift

@@ -829,8 +829,9 @@ struct AlertMessage: Identifiable {
         })
         observers.insert(vm.$state.sink { state in
             Task {
+                let isTakeoverAllowed = self.vmWindows[vm] is VMRemoteSessionState && (state == .started || state == .paused)
                 await self.remoteServer.broadcast { remote in
-                    try await remote.virtualMachine(id: vm.id, didTransitionToState: state)
+                    try await remote.virtualMachine(id: vm.id, didTransitionToState: state, isTakeoverAllowed: isTakeoverAllowed)
                 }
             }
         })
@@ -1271,12 +1272,13 @@ class UTMRemoteData: UTMData {
         vm.updateMountedDrives(mountedDrives)
     }
 
-    func remoteVirtualMachineDidTransition(id: UUID, state: UTMVirtualMachineState) async {
+    func remoteVirtualMachineDidTransition(id: UUID, state: UTMVirtualMachineState, isTakeoverAllowed: Bool) async {
         guard let vm = virtualMachines.first(where: { $0.id == id }) else {
             return
         }
         let remoteVM = vm as! VMRemoteData
         let wrapped = remoteVM.wrapped as! UTMRemoteSpiceVirtualMachine
+        remoteVM.isTakeoverAllowed = isTakeoverAllowed
         await wrapped.updateRemoteState(state)
     }
 

Platform/VMData.swift

@@ -68,7 +68,10 @@ import SwiftUI
     
     /// Copy from wrapped VM
     @Published var screenshot: UTMVirtualMachineScreenshot?
-    
+
+    /// If true, it is possible to hijack the session.
+    @Published var isTakeoverAllowed: Bool = false
+
     /// Allows changes in the config, registry, and VM to be reflected
     private var observers: [AnyCancellable] = []
     
@@ -450,6 +453,7 @@ class VMRemoteData: VMData {
         self._isShortcut = item.isShortcut
         self.initialState = item.state
         super.init()
+        self.isTakeoverAllowed = item.isTakeoverAllowed
         self.registryEntryWrapped = UTMRegistry.shared.entry(uuid: item.id, name: item.name, path: item.path)
         self.registryEntryWrapped!.isSuspended = item.isSuspended
         self.registryEntryWrapped!.externalDrives = item.mountedDrives.mapValues({ UTMRegistryEntry.File(dummyFromPath: $0) })

Platform/iOS/UTMDataExtension.swift

@@ -30,9 +30,9 @@ extension UTMData {
         }
         if let session = VMSessionState.allActiveSessions.values.first(where: { $0.vm.id == wrapped.id }) {
             session.showWindow()
-        } else if vm.state == .stopped {
+        } else if vm.isStopped || vm.isTakeoverAllowed {
             let session = VMSessionState(for: wrapped as! (any UTMSpiceVirtualMachine))
-            session.start()
+            session.start(options: options)
         } else {
             showErrorAlert(message: NSLocalizedString("This virtual machine is already running. In order to run it from this device, you must stop it first.", comment: "UTMDataExtension"))
         }

Platform/iOS/VMSessionState.swift

@@ -434,7 +434,11 @@ extension VMSessionState {
         }
         Self.allActiveSessions[id] = self
         showWindow()
-        vm.requestVmStart(options: options)
+        if vm.state == .paused {
+            vm.requestVmResume()
+        } else {
+            vm.requestVmStart(options: options)
+        }
     }
 
     func showWindow() {

Platform/macOS/UTMDataExtension.swift

@@ -86,6 +86,13 @@ extension UTMData {
         guard let wrapped = vm.wrapped as? UTMQemuVirtualMachine, type(of: wrapped).capabilities.supportsRemoteSession else {
             throw UTMDataError.unsupportedBackend
         }
+        if let existingSession = vmWindows[vm] as? VMRemoteSessionState, let spiceServerInfo = wrapped.spiceServerInfo {
+            if wrapped.state == .paused {
+                try await wrapped.resume()
+            }
+            existingSession.client = client
+            return spiceServerInfo
+        }
         guard vmWindows[vm] == nil else {
             throw UTMDataError.virtualMachineUnavailable
         }

Platform/macOS/VMRemoteSessionState.swift

@@ -19,7 +19,7 @@ import IOKit.pwr_mgt
 
 /// Represents the UI state for a single headless VM session.
 class VMRemoteSessionState: VMHeadlessSessionState {
-    let client: UTMRemoteServer.Remote
+    public weak var client: UTMRemoteServer.Remote?
 
     init(for vm: any UTMVirtualMachine, client: UTMRemoteServer.Remote, onStop: (() -> Void)?) {
         self.client = client
@@ -28,7 +28,7 @@ class VMRemoteSessionState: VMHeadlessSessionState {
 
     override func virtualMachine(_ vm: any UTMVirtualMachine, didErrorWithMessage message: String) {
         Task {
-            try? await client.virtualMachine(id: vm.id, didErrorWithMessage: message)
+            try? await client?.virtualMachine(id: vm.id, didErrorWithMessage: message)
             super.virtualMachine(vm, didErrorWithMessage: message)
         }
     }

Remote/UTMRemoteClient.swift

@@ -313,7 +313,7 @@ extension UTMRemoteClient {
         }
 
         private func _virtualMachineDidTransition(parameters: M.VirtualMachineDidTransition.Request) async throws -> M.VirtualMachineDidTransition.Reply {
-            await data.remoteVirtualMachineDidTransition(id: parameters.id, state: parameters.state)
+            await data.remoteVirtualMachineDidTransition(id: parameters.id, state: parameters.state, isTakeoverAllowed: parameters.isTakeoverAllowed)
             return .init()
         }
 

Remote/UTMRemoteMessage.swift

@@ -74,6 +74,7 @@ extension UTMRemoteMessageServer {
         let path: String
         let isShortcut: Bool
         let isSuspended: Bool
+        let isTakeoverAllowed: Bool
         let backend: UTMBackend
         let state: UTMVirtualMachineState
         let mountedDrives: [String: String]
@@ -360,6 +361,7 @@ extension UTMRemoteMessageClient {
         struct Request: Serializable, Codable {
             let id: UUID
             let state: UTMVirtualMachineState
+            let isTakeoverAllowed: Bool
         }
 
         struct Reply: Serializable, Codable {}

Remote/UTMRemoteServer.swift

@@ -228,8 +228,33 @@ actor UTMRemoteServer {
             if !connectedClients.contains(client) {
                 if let remote = establishedConnections.removeValue(forKey: client) {
                     remote.close()
+                    Task { @MainActor in
+                        await suspendSessions(for: remote)
+                    }
+                }
+            }
+        }
+    }
+
+    @MainActor
+    private func suspendSessions(for remote: Remote) async {
+        let sessions = data.vmWindows.compactMap {
+            if let session = $0.value as? VMRemoteSessionState {
+                return ($0.key, session)
+            } else {
+                return nil
+            }
+        }
+        await withTaskGroup(of: Void.self) { group in
+            for (vm, session) in sessions {
+                if session.client?.id == remote.id {
+                    session.client = nil
+                }
+                group.addTask {
+                    try? await vm.wrapped?.pause()
                 }
             }
+            await group.waitForAll()
         }
     }
 
@@ -712,11 +737,13 @@ extension UTMRemoteServer {
                 try parameters.ids.map { id in
                     let vm = try findVM(withId: id)
                     let mountedDrives = vm.registryEntry?.externalDrives.mapValues({ $0.path }) ?? [:]
+                    let isTakeoverAllowed = data.vmWindows[vm] is VMRemoteSessionState && (vm.state == .started || vm.state == .paused)
                     return M.VirtualMachineInformation(id: vm.id,
                                                        name: vm.detailsTitleLabel,
                                                        path: vm.pathUrl.path,
                                                        isShortcut: vm.isShortcut,
                                                        isSuspended: vm.registryEntry?.isSuspended ?? false,
+                                                       isTakeoverAllowed: isTakeoverAllowed,
                                                        backend: vm.wrapped is UTMQemuVirtualMachine ? .qemu : .unknown,
                                                        state: vm.wrapped?.state ?? .stopped,
                                                        mountedDrives: mountedDrives)
@@ -850,9 +877,10 @@ extension UTMRemoteServer {
 }
 
 extension UTMRemoteServer {
-    class Remote {
+    class Remote: Identifiable {
         typealias M = UTMRemoteMessageClient
         fileprivate(set) var peer: Peer<UTMRemoteMessageServer>!
+        let id = UUID()
 
         func close() {
             peer.close()
@@ -876,8 +904,8 @@ extension UTMRemoteServer {
             try await _mountedDrivesHasChanged(parameters: .init(id: id, mountedDrives: mountedDrives))
         }
 
-        func virtualMachine(id: UUID, didTransitionToState state: UTMVirtualMachineState) async throws {
-            try await _virtualMachineDidTransition(parameters: .init(id: id, state: state))
+        func virtualMachine(id: UUID, didTransitionToState state: UTMVirtualMachineState, isTakeoverAllowed: Bool) async throws {
+            try await _virtualMachineDidTransition(parameters: .init(id: id, state: state, isTakeoverAllowed: isTakeoverAllowed))
         }
 
         func virtualMachine(id: UUID, didErrorWithMessage message: String) async throws {

Remote/UTMRemoteSpiceVirtualMachine.swift

@@ -197,7 +197,7 @@ extension UTMRemoteSpiceVirtualMachine {
     }
 
     func start(options: UTMVirtualMachineStartOptions) async throws {
-        try await _state.operation(before: .stopped, during: .starting, after: .started) {
+        try await _state.operation(before: [.stopped, .started, .paused], during: .starting, after: .started) {
             let spiceServer = try await server.startVirtualMachine(id: id, options: options)
             var options = UTMSpiceIOOptions()
             if await !config.sound.isEmpty {
@@ -250,8 +250,12 @@ extension UTMRemoteSpiceVirtualMachine {
     }
 
     func resume() async throws {
-        try await _state.operation(before: .paused, during: .resuming, after: .started) {
-            try await server.resumeVirtualMachine(id: id)
+        if ioService == nil {
+            return try await start(options: [])
+        } else {
+            try await _state.operation(before: .paused, during: .resuming, after: .started) {
+                try await server.resumeVirtualMachine(id: id)
+            }
         }
     }