fix(agent): show correct NVMe capacity for Apple SSDs (#1873)

Co-authored-by: henrygd <hank@henrygd.me>

agent/smart.go

@@ -25,12 +25,15 @@ import (
 // SmartManager manages data collection for SMART devices
 type SmartManager struct {
 	sync.Mutex
-	SmartDataMap    map[string]*smart.SmartData
-	SmartDevices    []*DeviceInfo
-	refreshMutex    sync.Mutex
-	lastScanTime    time.Time
-	smartctlPath    string
-	excludedDevices map[string]struct{}
+	SmartDataMap       map[string]*smart.SmartData
+	SmartDevices       []*DeviceInfo
+	refreshMutex       sync.Mutex
+	lastScanTime       time.Time
+	smartctlPath       string
+	excludedDevices    map[string]struct{}
+	darwinNvmeOnce     sync.Once
+	darwinNvmeCapacity map[string]uint64      // serial → bytes cache, written once via darwinNvmeOnce
+	darwinNvmeProvider func() ([]byte, error) // overridable for testing
 }
 
 type scanOutput struct {
@@ -1033,6 +1036,52 @@ func parseScsiGigabytesProcessed(value string) int64 {
 	return parsed
 }
 
+// lookupDarwinNvmeCapacity returns the capacity in bytes for a given NVMe serial number on Darwin.
+// It uses system_profiler SPNVMeDataType to get capacity since Apple SSDs don't report user_capacity
+// via smartctl. Results are cached after the first call via sync.Once.
+func (sm *SmartManager) lookupDarwinNvmeCapacity(serial string) uint64 {
+	sm.darwinNvmeOnce.Do(func() {
+		sm.darwinNvmeCapacity = make(map[string]uint64)
+
+		provider := sm.darwinNvmeProvider
+		if provider == nil {
+			provider = func() ([]byte, error) {
+				ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+				defer cancel()
+				return exec.CommandContext(ctx, "system_profiler", "SPNVMeDataType", "-json").Output()
+			}
+		}
+
+		out, err := provider()
+		if err != nil {
+			slog.Debug("system_profiler NVMe lookup failed", "err", err)
+			return
+		}
+
+		var result struct {
+			SPNVMeDataType []struct {
+				Items []struct {
+					DeviceSerial string `json:"device_serial"`
+					SizeInBytes  uint64 `json:"size_in_bytes"`
+				} `json:"_items"`
+			} `json:"SPNVMeDataType"`
+		}
+		if err := json.Unmarshal(out, &result); err != nil {
+			slog.Debug("system_profiler NVMe parse failed", "err", err)
+			return
+		}
+
+		for _, controller := range result.SPNVMeDataType {
+			for _, item := range controller.Items {
+				if item.DeviceSerial != "" && item.SizeInBytes > 0 {
+					sm.darwinNvmeCapacity[item.DeviceSerial] = item.SizeInBytes
+				}
+			}
+		}
+	})
+	return sm.darwinNvmeCapacity[serial]
+}
+
 // parseSmartForNvme parses the output of smartctl --all -j /dev/nvmeX and updates the SmartDataMap
 // Returns hasValidData and exitStatus
 func (sm *SmartManager) parseSmartForNvme(output []byte) (bool, int) {
@@ -1069,6 +1118,9 @@ func (sm *SmartManager) parseSmartForNvme(output []byte) (bool, int) {
 	smartData.SerialNumber = data.SerialNumber
 	smartData.FirmwareVersion = data.FirmwareVersion
 	smartData.Capacity = data.UserCapacity.Bytes
+	if smartData.Capacity == 0 && (runtime.GOOS == "darwin" || sm.darwinNvmeProvider != nil) {
+		smartData.Capacity = sm.lookupDarwinNvmeCapacity(data.SerialNumber)
+	}
 	smartData.Temperature = data.NVMeSmartHealthInformationLog.Temperature
 	smartData.SmartStatus = getSmartStatus(smartData.Temperature, data.SmartStatus.Passed)
 	smartData.DiskName = data.Device.Name

agent/smart_test.go

@@ -1199,3 +1199,81 @@ func TestIsNvmeControllerPath(t *testing.T) {
 		})
 	}
 }
+
+func TestParseSmartForNvmeAppleSSD(t *testing.T) {
+	// Apple SSDs don't report user_capacity via smartctl; capacity should be fetched
+	// from system_profiler via the darwinNvmeProvider fallback.
+	fixturePath := filepath.Join("test-data", "smart", "apple_nvme.json")
+	data, err := os.ReadFile(fixturePath)
+	require.NoError(t, err)
+
+	providerCalls := 0
+	fakeProvider := func() ([]byte, error) {
+		providerCalls++
+		return []byte(`{
+			"SPNVMeDataType": [{
+				"_items": [{
+					"device_serial": "0ba0147940253c15",
+					"size_in_bytes": 251000193024
+				}]
+			}]
+		}`), nil
+	}
+
+	sm := &SmartManager{
+		SmartDataMap:       make(map[string]*smart.SmartData),
+		darwinNvmeProvider: fakeProvider,
+	}
+
+	hasData, _ := sm.parseSmartForNvme(data)
+	require.True(t, hasData)
+
+	deviceData, ok := sm.SmartDataMap["0ba0147940253c15"]
+	require.True(t, ok)
+	assert.Equal(t, "APPLE SSD AP0256Q", deviceData.ModelName)
+	assert.Equal(t, uint64(251000193024), deviceData.Capacity)
+	assert.Equal(t, uint8(42), deviceData.Temperature)
+	assert.Equal(t, "PASSED", deviceData.SmartStatus)
+	assert.Equal(t, 1, providerCalls, "system_profiler should be called once")
+
+	// Second parse: provider should NOT be called again (cache hit)
+	_, _ = sm.parseSmartForNvme(data)
+	assert.Equal(t, 1, providerCalls, "system_profiler should not be called again after caching")
+}
+
+func TestLookupDarwinNvmeCapacityMultipleDisks(t *testing.T) {
+	fakeProvider := func() ([]byte, error) {
+		return []byte(`{
+			"SPNVMeDataType": [
+				{
+					"_items": [
+						{"device_serial": "serial-disk0", "size_in_bytes": 251000193024},
+						{"device_serial": "serial-disk1", "size_in_bytes": 1000204886016}
+					]
+				},
+				{
+					"_items": [
+						{"device_serial": "serial-disk2", "size_in_bytes": 512110190592}
+					]
+				}
+			]
+		}`), nil
+	}
+
+	sm := &SmartManager{darwinNvmeProvider: fakeProvider}
+	assert.Equal(t, uint64(251000193024), sm.lookupDarwinNvmeCapacity("serial-disk0"))
+	assert.Equal(t, uint64(1000204886016), sm.lookupDarwinNvmeCapacity("serial-disk1"))
+	assert.Equal(t, uint64(512110190592), sm.lookupDarwinNvmeCapacity("serial-disk2"))
+	assert.Equal(t, uint64(0), sm.lookupDarwinNvmeCapacity("unknown-serial"))
+}
+
+func TestLookupDarwinNvmeCapacityProviderError(t *testing.T) {
+	fakeProvider := func() ([]byte, error) {
+		return nil, errors.New("system_profiler not found")
+	}
+
+	sm := &SmartManager{darwinNvmeProvider: fakeProvider}
+	assert.Equal(t, uint64(0), sm.lookupDarwinNvmeCapacity("any-serial"))
+	// Cache should be initialized even on error so we don't retry (Once already fired)
+	assert.NotNil(t, sm.darwinNvmeCapacity)
+}

agent/test-data/smart/apple_nvme.json

@@ -0,0 +1,51 @@
+{
+  "json_format_version": [1, 0],
+  "smartctl": {
+    "version": [7, 4],
+    "argv": ["smartctl", "-aix", "-j", "IOService:/AppleARMPE/arm-io@10F00000/AppleT810xIO/ans@77400000/AppleASCWrapV4/iop-ans-nub/RTBuddy(ANS2)/RTBuddyService/AppleANS3NVMeController/NS_01@1"],
+    "exit_status": 4
+  },
+  "device": {
+    "name": "IOService:/AppleARMPE/arm-io@10F00000/AppleT810xIO/ans@77400000/AppleASCWrapV4/iop-ans-nub/RTBuddy(ANS2)/RTBuddyService/AppleANS3NVMeController/NS_01@1",
+    "info_name": "IOService:/AppleARMPE/arm-io@10F00000/AppleT810xIO/ans@77400000/AppleASCWrapV4/iop-ans-nub/RTBuddy(ANS2)/RTBuddyService/AppleANS3NVMeController/NS_01@1",
+    "type": "nvme",
+    "protocol": "NVMe"
+  },
+  "model_name": "APPLE SSD AP0256Q",
+  "serial_number": "0ba0147940253c15",
+  "firmware_version": "555",
+  "smart_support": {
+    "available": true,
+    "enabled": true
+  },
+  "smart_status": {
+    "passed": true,
+    "nvme": {
+      "value": 0
+    }
+  },
+  "nvme_smart_health_information_log": {
+    "critical_warning": 0,
+    "temperature": 42,
+    "available_spare": 100,
+    "available_spare_threshold": 99,
+    "percentage_used": 1,
+    "data_units_read": 270189386,
+    "data_units_written": 166753862,
+    "host_reads": 7543766995,
+    "host_writes": 3761621926,
+    "controller_busy_time": 0,
+    "power_cycles": 366,
+    "power_on_hours": 2850,
+    "unsafe_shutdowns": 195,
+    "media_errors": 0,
+    "num_err_log_entries": 0
+  },
+  "temperature": {
+    "current": 42
+  },
+  "power_cycle_count": 366,
+  "power_on_time": {
+    "hours": 2850
+  }
+}