Add MachineRootfs type with SizeGB field (#206)

Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>

Author: dangra

machine_types.go

@@ -414,6 +414,11 @@ func (mpr *MachinePersistRootfs) UnmarshalJSON(raw []byte) error {
 	return nil
 }
 
+type MachineRootfs struct {
+	Persist MachinePersistRootfs `toml:"persist,omitempty" json:"persist,omitempty" enums:"never,always,restart"`
+	SizeGB  uint64               `toml:"size_gb,omitempty" json:"size_gb,omitempty"`
+}
+
 // @description The Machine restart policy defines whether and how flyd restarts a Machine after its main process exits. See https://fly.io/docs/machines/guides-examples/machine-restart-policy/.
 type MachineRestart struct {
 	// * no - Never try to restart a Machine automatically when its main process exits, whether that’s on purpose or on a crash.
@@ -439,13 +444,14 @@ type MachineMount struct {
 }
 
 type MachineGuest struct {
-	CPUKind          string               `toml:"cpu_kind,omitempty" json:"cpu_kind,omitempty"`
-	CPUs             int                  `toml:"cpus,omitempty" json:"cpus,omitempty"`
-	MemoryMB         int                  `toml:"memory_mb,omitempty" json:"memory_mb,omitempty"`
-	GPUs             int                  `toml:"gpus,omitempty" json:"gpus,omitempty"`
-	GPUKind          string               `toml:"gpu_kind,omitempty" json:"gpu_kind,omitempty"`
-	HostDedicationID string               `toml:"host_dedication_id,omitempty" json:"host_dedication_id,omitempty"`
-	PersistRootfs    MachinePersistRootfs `toml:"persist_rootfs,omitempty" json:"persist_rootfs,omitempty" enums:"never,always,restart"`
+	CPUKind          string `toml:"cpu_kind,omitempty" json:"cpu_kind,omitempty"`
+	CPUs             int    `toml:"cpus,omitempty" json:"cpus,omitempty"`
+	MemoryMB         int    `toml:"memory_mb,omitempty" json:"memory_mb,omitempty"`
+	GPUs             int    `toml:"gpus,omitempty" json:"gpus,omitempty"`
+	GPUKind          string `toml:"gpu_kind,omitempty" json:"gpu_kind,omitempty"`
+	HostDedicationID string `toml:"host_dedication_id,omitempty" json:"host_dedication_id,omitempty"`
+	// Deprecated: use MachineConfig.Rootfs instead
+	PersistRootfs MachinePersistRootfs `toml:"persist_rootfs,omitempty" json:"persist_rootfs,omitempty" enums:"never,always,restart"`
 
 	KernelArgs []string `toml:"kernel_args,omitempty" json:"kernel_args,omitempty"`
 }
@@ -804,6 +810,8 @@ type MachineConfig struct {
 	// with containers
 	Volumes []*VolumeConfig `toml:"volumes,omitempty" json:"volumes,omitempty"`
 
+	Rootfs *MachineRootfs `toml:"rootfs,omitempty" json:"rootfs,omitempty"`
+
 	// Deprecated: use Guest instead
 	VMSize string `toml:"size,omitempty" json:"size,omitempty"`
 	// Deprecated: use Service.Autostart instead

machine_types_test.go

@@ -303,6 +303,77 @@ func TestMachinePersistRootfsUnmarshalJSON(t *testing.T) {
 	}
 }
 
+func TestMachineRootfsJSON(t *testing.T) {
+	t.Run("marshal", func(t *testing.T) {
+		cases := []struct {
+			name   string
+			input  MachineConfig
+			output string
+		}{
+			{
+				name:   "rootfs with persist and size",
+				input:  MachineConfig{Rootfs: &MachineRootfs{Persist: MachinePersistRootfsAlways, SizeGB: 10}},
+				output: `{"init":{},"rootfs":{"persist":"always","size_gb":10}}`,
+			},
+			{
+				name:   "rootfs with persist only",
+				input:  MachineConfig{Rootfs: &MachineRootfs{Persist: MachinePersistRootfsRestart}},
+				output: `{"init":{},"rootfs":{"persist":"restart"}}`,
+			},
+			{
+				name:   "nil rootfs omitted",
+				input:  MachineConfig{},
+				output: `{"init":{}}`,
+			},
+		}
+		for _, tc := range cases {
+			b, err := json.Marshal(tc.input)
+			if err != nil {
+				t.Errorf("%s: unexpected error: %v", tc.name, err)
+			} else if string(b) != tc.output {
+				t.Errorf("%s: got %s, want %s", tc.name, string(b), tc.output)
+			}
+		}
+	})
+
+	t.Run("unmarshal", func(t *testing.T) {
+		cases := []struct {
+			name    string
+			input   string
+			persist MachinePersistRootfs
+			sizeGB  uint64
+		}{
+			{"persist and size", `{"rootfs":{"persist":"always","size_gb":10}}`, MachinePersistRootfsAlways, 10},
+			{"persist only", `{"rootfs":{"persist":"restart"}}`, MachinePersistRootfsRestart, 0},
+			{"size only", `{"rootfs":{"size_gb":5}}`, MachinePersistRootfsNone, 5},
+			{"no rootfs", `{}`, MachinePersistRootfsNone, 0},
+		}
+		for _, tc := range cases {
+			var mc MachineConfig
+			if err := json.Unmarshal([]byte(tc.input), &mc); err != nil {
+				t.Errorf("%s: unexpected error: %v", tc.name, err)
+				continue
+			}
+			if tc.persist == MachinePersistRootfsNone && tc.sizeGB == 0 {
+				if mc.Rootfs != nil {
+					t.Errorf("%s: expected nil rootfs, got %+v", tc.name, mc.Rootfs)
+				}
+				continue
+			}
+			if mc.Rootfs == nil {
+				t.Errorf("%s: expected non-nil rootfs", tc.name)
+				continue
+			}
+			if mc.Rootfs.Persist != tc.persist {
+				t.Errorf("%s: persist got %v, want %v", tc.name, mc.Rootfs.Persist, tc.persist)
+			}
+			if mc.Rootfs.SizeGB != tc.sizeGB {
+				t.Errorf("%s: size_gb got %d, want %d", tc.name, mc.Rootfs.SizeGB, tc.sizeGB)
+			}
+		}
+	})
+}
+
 func TestMachineAutostopUnmarshalJSON(t *testing.T) {
 	type testcase struct {
 		input  string