feat: attach image drives to warm starts

crates/aenv/src/client/sandboxes.rs

@@ -22,15 +22,32 @@ fn volume_mounts_request(
     })
 }
 
+#[derive(Debug, PartialEq, Serialize)]
+pub struct AttachedDrive {
+    #[serde(rename = "driveID")]
+    pub drive_id: String,
+    #[serde(rename = "readOnly")]
+    pub read_only: bool,
+    #[serde(rename = "mountPath")]
+    pub mount_path: String,
+    pub source: AttachedDriveSource,
+}
+
+#[derive(Debug, PartialEq, Serialize)]
+pub struct AttachedDriveSource {
+    pub image: String,
+}
+
 #[derive(Debug, Serialize)]
-pub struct NewSandbox<'a> {
+pub struct NewSandboxV2<'a> {
     #[serde(rename = "templateID")]
     pub template_id: &'a str,
     #[serde(skip_serializing_if = "Option::is_none")]
     pub timeout: Option<u32>,
-    pub secure: bool,
     #[serde(skip_serializing_if = "Option::is_none", rename = "volumeMounts")]
     pub volume_mounts: Option<Vec<SandboxVolumeMount>>,
+    #[serde(skip_serializing_if = "Option::is_none", rename = "attachedDrives")]
+    pub attached_drives: Option<Vec<AttachedDrive>>,
 }
 
 #[derive(Debug, Serialize)]
@@ -47,6 +64,8 @@ pub struct NewColdSandbox<'a> {
     pub secure: bool,
     #[serde(skip_serializing_if = "Option::is_none", rename = "volumeMounts")]
     pub volume_mounts: Option<Vec<SandboxVolumeMount>>,
+    #[serde(skip_serializing_if = "Option::is_none", rename = "attachedDrives")]
+    pub attached_drives: Option<Vec<AttachedDrive>>,
 }
 
 #[derive(Deserialize)]
@@ -100,14 +119,15 @@ impl Client {
         template_id: &str,
         timeout: Option<u32>,
         volume_mounts: Option<HashMap<String, String>>,
+        attached_drives: Option<Vec<AttachedDrive>>,
     ) -> Result<Sandbox> {
-        let body = NewSandbox {
+        let body = NewSandboxV2 {
             template_id,
             timeout,
-            secure: true,
             volume_mounts: volume_mounts_request(volume_mounts),
+            attached_drives,
         };
-        let resp = handle_status(self.post("/sandboxes").send_json(&body))?;
+        let resp = handle_status(self.post("/v2/sandboxes").send_json(&body))?;
         let sandbox: Sandbox = resp.into_json()?;
         Ok(sandbox)
     }
@@ -121,6 +141,7 @@ impl Client {
         memory_mb: Option<u32>,
         disk_size_mb: Option<u32>,
         volume_mounts: Option<HashMap<String, String>>,
+        attached_drives: Option<Vec<AttachedDrive>>,
     ) -> Result<Sandbox> {
         let body = NewColdSandbox {
             image,
@@ -130,6 +151,7 @@ impl Client {
             disk_size_mb,
             secure: true,
             volume_mounts: volume_mounts_request(volume_mounts),
+            attached_drives,
         };
         let resp = handle_status(self.post("/sandboxes-cold").send_json(&body))?;
         let sandbox: Sandbox = resp.into_json()?;
@@ -203,23 +225,48 @@ impl Client {
 
 #[cfg(test)]
 mod tests {
-    use super::{NewColdSandbox, NewSandbox, RefreshSandbox};
+    use super::{AttachedDrive, AttachedDriveSource, NewColdSandbox, NewSandboxV2, RefreshSandbox};
 
     #[test]
     fn new_sandbox_serializes_template_start() {
-        let body = NewSandbox {
+        let body = NewSandboxV2 {
             template_id: "base-template",
             timeout: Some(300),
-            secure: true,
             volume_mounts: None,
+            attached_drives: None,
         };
 
         let value = serde_json::to_value(body).unwrap();
         assert_eq!(value["templateID"], "base-template");
         assert_eq!(value["timeout"], 300);
-        assert_eq!(value["secure"], true);
+        assert!(value.get("secure").is_none());
         assert!(value.get("cpuCount").is_none());
         assert!(value.get("memoryMB").is_none());
+        assert!(value.get("attachedDrives").is_none());
+    }
+
+    #[test]
+    fn new_sandbox_serializes_attached_drives() {
+        let body = NewSandboxV2 {
+            template_id: "base-template",
+            timeout: None,
+            volume_mounts: None,
+            attached_drives: Some(vec![AttachedDrive {
+                drive_id: "image_mnt_models".to_owned(),
+                read_only: false,
+                mount_path: "/mnt/models".to_owned(),
+                source: AttachedDriveSource {
+                    image: "aenv/models:v1".to_owned(),
+                },
+            }]),
+        };
+
+        let value = serde_json::to_value(body).unwrap();
+        let drive = &value["attachedDrives"][0];
+        assert_eq!(drive["driveID"], "image_mnt_models");
+        assert_eq!(drive["readOnly"], false);
+        assert_eq!(drive["mountPath"], "/mnt/models");
+        assert_eq!(drive["source"]["image"], "aenv/models:v1");
     }
 
     #[test]
@@ -232,6 +279,7 @@ mod tests {
             disk_size_mb: Some(8192),
             secure: true,
             volume_mounts: None,
+            attached_drives: None,
         };
 
         let value = serde_json::to_value(body).unwrap();

crates/aenv/src/commands/start.rs

@@ -1,7 +1,8 @@
+use crate::client::sandboxes::{AttachedDrive, AttachedDriveSource};
 use crate::client::{volumes::validate_volume_reference, Client};
 use anyhow::{Context, Result};
 use clap::Args as ClapArgs;
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 use std::time::{Duration, Instant};
 
 const ENVD_READY_TIMEOUT: Duration = Duration::from_secs(30);
@@ -14,6 +15,8 @@ const ENVD_READY_PROBE_INTERVAL: Duration = Duration::from_millis(100);
   aenv start --cold docker.io/library/ubuntu:24.04
   aenv start --cold docker.io/library/ubuntu:24.04 --cpu 2 --mem 2048 --disk-size-mb 8192
   aenv start --cold -d docker.io/library/ubuntu:24.04
+  aenv start my-template --image-mount /workspace=aenv/repo:main
+  aenv start my-template --image-mount-ro /mnt/models=aenv/models:v1
 
 Resource overrides are only supported with --cold.")]
 pub struct Args {
@@ -37,6 +40,12 @@ pub struct Args {
     /// Mount a persistent volume using MOUNT_PATH=VOLUME_ID_OR_NAME (repeatable)
     #[arg(long = "volume", value_name = "MOUNT_PATH=VOLUME", action = clap::ArgAction::Append)]
     volumes: Vec<String>,
+    /// Attach a writable drive from an image using MOUNT_PATH=IMAGE (repeatable)
+    #[arg(long = "image-mount", value_name = "MOUNT_PATH=IMAGE", action = clap::ArgAction::Append)]
+    image_mounts: Vec<String>,
+    /// Attach a read-only drive from an image using MOUNT_PATH=IMAGE (repeatable)
+    #[arg(long = "image-mount-ro", value_name = "MOUNT_PATH=IMAGE", action = clap::ArgAction::Append)]
+    read_only_image_mounts: Vec<String>,
 }
 
 fn parse_disk_size_mb(value: &str) -> std::result::Result<u32, String> {
@@ -52,6 +61,11 @@ fn parse_disk_size_mb(value: &str) -> std::result::Result<u32, String> {
 pub fn run(args: Args) -> Result<()> {
     let client = Client::from_env()?;
     let volume_mounts = parse_volume_mounts(&args.volumes)?;
+    let attached_drives = parse_image_mounts(
+        &args.image_mounts,
+        &args.read_only_image_mounts,
+        volume_mounts.as_ref(),
+    )?;
     let sandbox = if args.cold {
         client.create_cold_sandbox(
             &args.target,
@@ -60,12 +74,18 @@ pub fn run(args: Args) -> Result<()> {
             args.resources.memory_mb,
             args.disk_size_mb,
             volume_mounts,
+            attached_drives,
         )?
     } else {
         if args.resources.is_set() || args.disk_size_mb.is_some() {
             anyhow::bail!("--cpu-count, --memory-mb, and --disk-size-mb require --cold");
         }
-        client.create_sandbox(&args.target, Some(args.timeout), volume_mounts)?
+        client.create_sandbox(
+            &args.target,
+            Some(args.timeout),
+            volume_mounts,
+            attached_drives,
+        )?
     };
     let sandbox_id = sandbox.sandbox_id;
 
@@ -111,6 +131,69 @@ fn parse_volume_mounts(values: &[String]) -> Result<Option<HashMap<String, Strin
     Ok(Some(mounts))
 }
 
+/// Builds attached drives from writable and read-only MOUNT_PATH=IMAGE values.
+/// Their paths must not overlap each other or a volume mount.
+fn parse_image_mounts(
+    writable: &[String],
+    read_only: &[String],
+    volume_mounts: Option<&HashMap<String, String>>,
+) -> Result<Option<Vec<AttachedDrive>>> {
+    if writable.is_empty() && read_only.is_empty() {
+        return Ok(None);
+    }
+    let mut paths: Vec<String> = volume_mounts
+        .map(|mounts| mounts.keys().cloned().collect())
+        .unwrap_or_default();
+    let mut drive_ids = HashSet::new();
+    let mut drives = Vec::with_capacity(writable.len() + read_only.len());
+    let values = writable
+        .iter()
+        .map(|value| (value, false))
+        .chain(read_only.iter().map(|value| (value, true)));
+    for (value, read_only) in values {
+        let Some((mount_path, image)) = value
+            .split_once('=')
+            .filter(|(_, image)| !image.trim().is_empty())
+        else {
+            anyhow::bail!("image mount must use MOUNT_PATH=IMAGE syntax: {value}");
+        };
+        let mount_path = normalize_mount_path(mount_path)
+            .with_context(|| format!("invalid image mount: {value}"))?;
+        if paths
+            .iter()
+            .any(|existing| mount_paths_overlap(existing, &mount_path))
+        {
+            anyhow::bail!("image mount path overlaps another mount: {mount_path}");
+        }
+        paths.push(mount_path.clone());
+        drives.push(AttachedDrive {
+            drive_id: image_drive_id(&mount_path, &mut drive_ids),
+            read_only,
+            mount_path,
+            source: AttachedDriveSource {
+                image: image.trim().to_owned(),
+            },
+        });
+    }
+    Ok(Some(drives))
+}
+
+/// Derives a unique drive ID such as `image_mnt_models` from a normalized
+/// mount path. The prefix keeps it clear of reserved drive IDs.
+fn image_drive_id(mount_path: &str, used: &mut HashSet<String>) -> String {
+    let base: String = format!("image{mount_path}")
+        .chars()
+        .map(|c| if c.is_ascii_alphanumeric() { c } else { '_' })
+        .collect();
+    let mut drive_id = base.clone();
+    let mut suffix = 2;
+    while !used.insert(drive_id.clone()) {
+        drive_id = format!("{base}_{suffix}");
+        suffix += 1;
+    }
+    drive_id
+}
+
 fn normalize_mount_path(value: &str) -> Result<String> {
     if !value.starts_with('/') || value == "/" {
         anyhow::bail!("volume mount path must be an absolute guest path other than /: {value}");
@@ -186,7 +269,7 @@ async fn wait_for_envd(
 
 #[cfg(test)]
 mod tests {
-    use super::{parse_volume_mounts, Args};
+    use super::{parse_image_mounts, parse_volume_mounts, Args};
     use clap::Parser;
 
     #[derive(Parser)]
@@ -221,4 +304,86 @@ mod tests {
                 .is_err()
         );
     }
+
+    #[test]
+    fn parses_image_mounts_as_attached_drives() {
+        assert!(parse_image_mounts(&[], &[], None).unwrap().is_none());
+        let drives = parse_image_mounts(
+            &[
+                "/workspace=registry.example.com:5000/repo:main".to_owned(),
+                "/mnt/a-b=aenv/a".to_owned(),
+            ],
+            &[
+                "/mnt/models/=aenv/models:v1".to_owned(),
+                "/mnt/a.b=aenv/b".to_owned(),
+            ],
+            None,
+        )
+        .unwrap()
+        .unwrap();
+        let summary: Vec<_> = drives
+            .iter()
+            .map(|drive| {
+                (
+                    drive.drive_id.as_str(),
+                    drive.mount_path.as_str(),
+                    drive.source.image.as_str(),
+                    drive.read_only,
+                )
+            })
+            .collect();
+        assert_eq!(
+            summary,
+            [
+                (
+                    "image_workspace",
+                    "/workspace",
+                    "registry.example.com:5000/repo:main",
+                    false
+                ),
+                ("image_mnt_a_b", "/mnt/a-b", "aenv/a", false),
+                ("image_mnt_models", "/mnt/models", "aenv/models:v1", true),
+                ("image_mnt_a_b_2", "/mnt/a.b", "aenv/b", true),
+            ]
+        );
+    }
+
+    #[test]
+    fn image_mount_flags_parse_for_warm_and_cold_starts() {
+        for cold in [false, true] {
+            let mut values = vec!["test", "my-template"];
+            if cold {
+                values.push("--cold");
+            }
+            values.extend([
+                "--image-mount",
+                "/workspace=aenv/repo:main",
+                "--image-mount-ro",
+                "/mnt/models=aenv/models:v1",
+            ]);
+            let args = TestCli::try_parse_from(values).unwrap().args;
+            let drives = parse_image_mounts(&args.image_mounts, &args.read_only_image_mounts, None)
+                .unwrap()
+                .unwrap();
+            assert_eq!(drives.len(), 2);
+            assert!(!drives[0].read_only);
+            assert!(drives[1].read_only);
+        }
+    }
+
+    #[test]
+    fn rejects_invalid_or_overlapping_image_mounts() {
+        assert!(parse_image_mounts(&["/mnt/models".to_owned()], &[], None).is_err());
+        assert!(parse_image_mounts(&["/mnt/models=".to_owned()], &[], None).is_err());
+        assert!(parse_image_mounts(&["models=aenv/models".to_owned()], &[], None).is_err());
+        assert!(parse_image_mounts(
+            &["/mnt/models=aenv/a".to_owned()],
+            &["/mnt/models/sub=aenv/b".to_owned()],
+            None,
+        )
+        .is_err());
+        let volumes = parse_volume_mounts(&["/mnt/data=cache".to_owned()]).unwrap();
+        let overlapping = ["/mnt/data=aenv/models".to_owned()];
+        assert!(parse_image_mounts(&overlapping, &[], volumes.as_ref()).is_err());
+    }
 }