From d1719208826538fd6b7d8abd683399d57cf696ca Mon Sep 17 00:00:00 2001
From: f-trycua <f@trycua.com>
Date: Tue, 22 Apr 2025 15:19:45 -0700
Subject: [PATCH] Handle computer sparse image

---
 libs/computer/computer/computer.py | 20 ++++-------
 libs/lume/src/LumeController.swift | 57 +++++++++++++++++++++++-------
 2 files changed, 52 insertions(+), 25 deletions(-)

diff --git a/libs/computer/computer/computer.py b/libs/computer/computer/computer.py
index 6c1119ac..f4d9d9bf 100644
--- a/libs/computer/computer/computer.py
+++ b/libs/computer/computer/computer.py
@@ -227,24 +227,18 @@ class Computer:
                         self.logger.error(f"Failed to initialize PyLume context: {e}")
                         raise RuntimeError(f"Failed to initialize PyLume: {e}")
 
-                # Try to get the VM, if it doesn't exist, create it and pull the image
+                # Try to get the VM, if it doesn't exist, return an error
                 try:
                     vm = await self.config.pylume.get_vm(self.config.name)  # type: ignore[attr-defined]
                     self.logger.verbose(f"Found existing VM: {self.config.name}")
                 except Exception as e:
-                    self.logger.verbose(f"VM not found, pulling image: {e}")
-                    image_ref = ImageRef(
-                        image=self.config.image,
-                        tag=self.config.tag,
-                        registry="ghcr.io",
-                        organization="trycua",
+                    self.logger.error(f"VM not found: {self.config.name}")
+                    self.logger.error(
+                        f"Please pull the VM first with lume pull macos-sequoia-cua-sparse:latest: {e}"
+                    )
+                    raise RuntimeError(
+                        f"VM not found: {self.config.name}. Please pull the VM first."
                     )
-                    self.logger.info(f"Pulling image {self.config.image}:{self.config.tag}...")
-                    try:
-                        await self.config.pylume.pull_image(image_ref, name=self.config.name)  # type: ignore[attr-defined]
-                    except Exception as pull_error:
-                        self.logger.error(f"Failed to pull image: {pull_error}")
-                        raise RuntimeError(f"Failed to pull VM image: {pull_error}")
 
                 # Convert paths to SharedDirectory objects
                 shared_directories = []
diff --git a/libs/lume/src/LumeController.swift b/libs/lume/src/LumeController.swift
index 1329f8c5..c50edd90 100644
--- a/libs/lume/src/LumeController.swift
+++ b/libs/lume/src/LumeController.swift
@@ -399,20 +399,53 @@ final class LumeController {
         storage: String? = nil
     ) async throws {
         do {
-            let vmName: String = name ?? normalizeVMName(name: image)
+            // Convert non-sparse image to sparse version if needed
+            var actualImage = image
+            var actualName = name
+
+            // Check if image is a non-sparse version (doesn't contain -sparse)
+            if !image.contains("-sparse") {
+                // Split the image to get name and tag
+                let components = image.split(separator: ":")
+                guard components.count == 2 else {
+                    throw ValidationError("Invalid image format. Expected format: name:tag")
+                }
+
+                let originalName = String(components[0])
+                let tag = String(components[1])
+
+                // Create sparse version of the image name
+                actualImage = "\(originalName)-sparse:\(tag)"
+
+                // If name wasn't explicitly provided, use the original image name (without -sparse)
+                if actualName == nil {
+                    actualName = originalName
+                }
+
+                Logger.info(
+                    "Converting to sparse image",
+                    metadata: [
+                        "original": image,
+                        "sparse": actualImage,
+                        "vm_name": actualName ?? "default",
+                    ]
+                )
+            }
+
+            let vmName: String = actualName ?? normalizeVMName(name: actualImage)
 
             Logger.info(
                 "Pulling image",
                 metadata: [
-                    "image": image,
-                    "name": name ?? "default",
+                    "image": actualImage,
+                    "name": actualName ?? "default",
                     "registry": registry,
                     "organization": organization,
                     "location": storage ?? "default",
                 ])
 
             try self.validatePullParameters(
-                image: image,
+                image: actualImage,
                 name: vmName,
                 registry: registry,
                 organization: organization,
@@ -422,7 +455,7 @@ final class LumeController {
             let imageContainerRegistry = ImageContainerRegistry(
                 registry: registry, organization: organization)
             try await imageContainerRegistry.pull(
-                image: image,
+                image: actualImage,
                 name: vmName,
                 locationName: storage)
 
@@ -440,7 +473,7 @@ final class LumeController {
             Logger.info(
                 "Image pulled successfully",
                 metadata: [
-                    "image": image,
+                    "image": actualImage,
                     "name": vmName,
                     "registry": registry,
                     "organization": organization,
@@ -477,7 +510,7 @@ final class LumeController {
                     "location": storage ?? "default",
                     "chunk_size": "\(chunkSizeMb)MB",
                     "dry_run": "\(dryRun)",
-                    "reassemble": "\(reassemble)"
+                    "reassemble": "\(reassemble)",
                 ])
 
             try validatePushParameters(
@@ -490,14 +523,14 @@ final class LumeController {
 
             // Find the actual location of the VM
             let actualLocation = try self.validateVMExists(name, storage: storage)
-            
+
             // Get the VM directory
             let vmDir = try home.getVMDirectory(name, storage: actualLocation)
-            
-            // Use ImageContainerRegistry to push the VM 
+
+            // Use ImageContainerRegistry to push the VM
             let imageContainerRegistry = ImageContainerRegistry(
                 registry: registry, organization: organization)
-            
+
             try await imageContainerRegistry.push(
                 vmDirPath: vmDir.dir.path,
                 imageName: imageName,
@@ -849,7 +882,7 @@ final class LumeController {
         guard !organization.isEmpty else {
             throw ValidationError("Organization cannot be empty")
         }
-        
+
         // Verify VM exists (this will throw if not found)
         _ = try self.validateVMExists(name)
     }