fix(lume): handle guest-initiated VM shutdown via VZVirtualMachineDelegate

Author: NikkeTryHard

libs/lume/src/LumeController.swift

@@ -1122,7 +1122,8 @@ final class LumeController {
                 usbMassStoragePaths: usbMassStoragePaths,
                 networkMode: networkMode,
                 clipboard: clipboard)
-            Logger.info("VM started successfully", metadata: ["name": normalizedName])
+            SharedVM.shared.removeVM(name: normalizedName)
+            Logger.info("VM exited", metadata: ["name": normalizedName])
         } catch {
             SharedVM.shared.removeVM(name: normalizedName)
             Logger.error("Failed to run VM", metadata: ["error": error.localizedDescription])

libs/lume/src/VM/VM.swift

@@ -325,9 +325,26 @@ class VM {
                 await clipboardWatcher?.start()
             }
 
-            while true {
-                try await Task.sleep(nanoseconds: UInt64(1e9))
+            // Block until the guest OS shuts down or crashes
+            if let guestError = await service.waitForGuestStop() {
+                Logger.error("VM stopped unexpectedly", metadata: [
+                    "name": vmDirContext.name,
+                    "error": guestError.localizedDescription,
+                ])
+            } else {
+                Logger.info("Guest initiated shutdown", metadata: ["name": vmDirContext.name])
             }
+
+            // Clean up after guest stop
+            await clipboardWatcher?.stop()
+            clipboardWatcher = nil
+            virtualizationService = nil
+            vncService.stop()
+
+            Logger.info("Releasing file lock after guest stop", metadata: ["name": vmDirContext.name])
+            flock(fileHandle.fileDescriptor, LOCK_UN)
+            try? fileHandle.close()
+            unlockConfigFile()
         } catch {
             Logger.error(
                 "Failed in VM.run",
@@ -1051,9 +1068,19 @@ class VM {
                 }
             }
 
-            while true {
-                try await Task.sleep(nanoseconds: UInt64(1e9))
+            if let guestError = await service.waitForGuestStop() {
+                Logger.error("VM stopped unexpectedly", metadata: [
+                    "name": vmDirContext.name,
+                    "error": guestError.localizedDescription,
+                ])
+            } else {
+                Logger.info("Guest initiated shutdown", metadata: ["name": vmDirContext.name])
             }
+
+            virtualizationService = nil
+            vncService.stop()
+            flock(fileHandle.fileDescriptor, LOCK_UN)
+            try? fileHandle.close()
         } catch {
             Logger.error(
                 "Failed to create/start VM with USB storage",

libs/lume/src/Virtualization/VMVirtualizationService.swift

@@ -28,21 +28,62 @@ protocol VMVirtualizationService {
     func pause() async throws
     func resume() async throws
     func getVirtualMachine() -> Any
+    func waitForGuestStop() async -> Error?
 }
 
 /// Base implementation of VMVirtualizationService using VZVirtualMachine
 @MainActor
-class BaseVirtualizationService: VMVirtualizationService {
+class BaseVirtualizationService: NSObject, VMVirtualizationService, VZVirtualMachineDelegate {
     let virtualMachine: VZVirtualMachine
     let recoveryMode: Bool  // Store whether we should start in recovery mode
 
+    private var guestStopContinuation: CheckedContinuation<Error?, Never>?
+    private var pendingGuestStop: (fired: Bool, error: Error?) = (false, nil)
+
     var state: VZVirtualMachine.State {
         virtualMachine.state
     }
 
     init(virtualMachine: VZVirtualMachine, recoveryMode: Bool = false) {
         self.virtualMachine = virtualMachine
         self.recoveryMode = recoveryMode
+        super.init()
+        self.virtualMachine.delegate = self
+    }
+
+    func waitForGuestStop() async -> Error? {
+        if pendingGuestStop.fired {
+            return pendingGuestStop.error
+        }
+        return await withCheckedContinuation { continuation in
+            guestStopContinuation = continuation
+        }
+    }
+
+    // MARK: - VZVirtualMachineDelegate
+
+    nonisolated func guestDidStop(_ virtualMachine: VZVirtualMachine) {
+        Task { @MainActor in
+            Logger.info("Guest initiated shutdown")
+            if let continuation = guestStopContinuation {
+                guestStopContinuation = nil
+                continuation.resume(returning: nil)
+            } else {
+                pendingGuestStop = (true, nil)
+            }
+        }
+    }
+
+    nonisolated func virtualMachine(_ virtualMachine: VZVirtualMachine, didStopWithError error: Error) {
+        Task { @MainActor in
+            Logger.error("Guest stopped with error", metadata: ["error": error.localizedDescription])
+            if let continuation = guestStopContinuation {
+                guestStopContinuation = nil
+                continuation.resume(returning: error)
+            } else {
+                pendingGuestStop = (true, error)
+            }
+        }
     }
 
     func start() async throws {

libs/lume/tests/Mocks/MockVMVirtualizationService.swift

@@ -62,4 +62,8 @@ final class MockVMVirtualizationService: VMVirtualizationService {
     func getVirtualMachine() -> Any {
         return "mock_vm"
     }
+
+    func waitForGuestStop() async -> Error? {
+        return nil
+    }
 }