fix(hub): remove stale smart_devices records when a drive is no longer reported (#2178)

* Fix duplicate /dev/sdg-style entries

* only prune devices after complete refreshes

---------

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

agent/handlers.go

@@ -166,14 +166,16 @@ type GetSmartDataHandler struct{}
 
 func (h *GetSmartDataHandler) Handle(hctx *HandlerContext) error {
 	if hctx.Agent.smartManager == nil {
-		// return empty map to indicate no data
-		return hctx.SendResponse(map[string]smart.SmartData{}, hctx.RequestID)
+		return hctx.SendResponse(smart.SmartDataResponse{Data: map[string]smart.SmartData{}}, hctx.RequestID)
 	}
-	if err := hctx.Agent.smartManager.Refresh(false); err != nil {
+	complete, err := hctx.Agent.smartManager.Refresh(false)
+	if err != nil {
 		slog.Debug("smart refresh failed", "err", err)
 	}
-	data := hctx.Agent.smartManager.GetCurrentData()
-	return hctx.SendResponse(data, hctx.RequestID)
+	return hctx.SendResponse(smart.SmartDataResponse{
+		Data:     hctx.Agent.smartManager.GetCurrentData(),
+		Complete: complete,
+	}, hctx.RequestID)
 }
 
 ////////////////////////////////////////////////////////////////////////////

agent/handlers_test.go

@@ -7,6 +7,7 @@ import (
 
 	"github.com/fxamacker/cbor/v2"
 	"github.com/henrygd/beszel/internal/common"
+	"github.com/henrygd/beszel/internal/entities/smart"
 	"github.com/stretchr/testify/assert"
 )
 
@@ -17,6 +18,18 @@ type MockHandler struct {
 	handleFunc           func(ctx *HandlerContext) error
 }
 
+func TestNewAgentResponseSmartData(t *testing.T) {
+	response := newAgentResponse(smart.SmartDataResponse{
+		Data: map[string]smart.SmartData{
+			"AAA": {SerialNumber: "AAA"},
+		},
+		Complete: true,
+	}, nil)
+
+	assert.Equal(t, "AAA", response.SmartData["AAA"].SerialNumber)
+	assert.True(t, response.SmartComplete)
+}
+
 func (m *MockHandler) Handle(ctx *HandlerContext) error {
 	if m.handleFunc != nil {
 		return m.handleFunc(ctx)

agent/response.go

@@ -21,6 +21,9 @@ func newAgentResponse(data any, requestID *uint32) common.AgentResponse {
 		response.String = &v
 	case map[string]smart.SmartData:
 		response.SmartData = v
+	case smart.SmartDataResponse:
+		response.SmartData = v.Data
+		response.SmartComplete = v.Complete
 	case systemd.ServiceDetails:
 		response.ServiceInfo = v
 	default:

agent/smart.go

@@ -70,8 +70,9 @@ type deviceKey struct {
 
 var errNoValidSmartData = fmt.Errorf("no valid SMART data found") // Error for missing data
 
-// Refresh updates SMART data for all known devices
-func (sm *SmartManager) Refresh(forceScan bool) error {
+// Refresh updates SMART data for all known devices and reports whether every
+// discovered device was collected successfully.
+func (sm *SmartManager) Refresh(forceScan bool) (bool, error) {
 	sm.refreshMutex.Lock()
 	defer sm.refreshMutex.Unlock()
 
@@ -92,7 +93,7 @@ func (sm *SmartManager) Refresh(forceScan bool) error {
 		}
 	}
 
-	return sm.resolveRefreshError(scanErr, collectErr)
+	return scanErr == nil && collectErr == nil, sm.resolveRefreshError(scanErr, collectErr)
 }
 
 // devicesSnapshot returns a copy of the current device slice to avoid iterating

internal/common/common-ws.go

@@ -42,7 +42,8 @@ type AgentResponse struct {
 	SmartData   map[string]smart.SmartData `cbor:"5,keyasint,omitempty,omitzero"` // Legacy (<= 0.17)
 	ServiceInfo systemd.ServiceDetails     `cbor:"6,keyasint,omitempty,omitzero"` // Legacy (<= 0.17)
 	// Data is the generic response payload for new endpoints (0.18+)
-	Data cbor.RawMessage `cbor:"7,keyasint,omitempty,omitzero"`
+	Data          cbor.RawMessage `cbor:"7,keyasint,omitempty,omitzero"`
+	SmartComplete bool            `cbor:"8,keyasint,omitempty,omitzero"`
 }
 
 type FingerprintRequest struct {

internal/entities/smart/smart.go

@@ -494,7 +494,7 @@ type SmartInfoForNvme struct {
 	FirmwareVersion string           `json:"firmware_version"`
 	// NVMePCIVendor                 NVMePCIVendor                 `json:"nvme_pci_vendor"`
 	// NVMeIEEEOUIIdentifier         uint32                        `json:"nvme_ieee_oui_identifier"`
-	NVMeTotalCapacity             uint64                        `json:"nvme_total_capacity"`
+	NVMeTotalCapacity uint64 `json:"nvme_total_capacity"`
 	// NVMeUnallocatedCapacity       uint64                        `json:"nvme_unallocated_capacity"`
 	// NVMeControllerID              uint16                        `json:"nvme_controller_id"`
 	// NVMeVersion                   VersionStringInfo             `json:"nvme_version"`
@@ -531,6 +531,13 @@ type SmartData struct {
 	Attributes      []*SmartAttribute `json:"a,omitempty" cbor:"9,keyasint,omitempty"`
 }
 
+// SmartDataResponse contains the collected data and whether every discovered
+// device was collected. Older agents omit Complete, so hubs must not prune from it.
+type SmartDataResponse struct {
+	Data     map[string]SmartData `json:"data" cbor:"0,keyasint"`
+	Complete bool                 `json:"complete" cbor:"1,keyasint,omitempty"` // Whether every discovered device was collected
+}
+
 type SmartAttribute struct {
 	ID         uint16 `json:"id,omitempty" cbor:"0,keyasint,omitempty"`
 	Name       string `json:"n" cbor:"1,keyasint"`

internal/hub/systems/system.go

@@ -532,11 +532,16 @@ func (sys *System) FetchSystemdInfoFromAgent(serviceName string) (systemd.Servic
 	return result, err
 }
 
-// FetchSmartDataFromAgent fetches SMART data from the agent
-func (sys *System) FetchSmartDataFromAgent() (map[string]smart.SmartData, error) {
+// FetchSmartDataFromAgent fetches SMART data from the agent.
+func (sys *System) FetchSmartDataFromAgent() (smart.SmartDataResponse, error) {
 	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
 	defer cancel()
-	var result map[string]smart.SmartData
+	if sys.agentVersion.LT(beszel.MinVersionAgentResponse) {
+		var data map[string]smart.SmartData
+		err := sys.request(ctx, common.GetSmartData, nil, &data)
+		return smart.SmartDataResponse{Data: data}, err
+	}
+	var result smart.SmartDataResponse
 	err := sys.request(ctx, common.GetSmartData, nil, &result)
 	return result, err
 }

internal/hub/systems/system_smart.go

@@ -7,6 +7,7 @@ import (
 	"time"
 
 	"github.com/henrygd/beszel/internal/entities/smart"
+	"github.com/pocketbase/dbx"
 	"github.com/pocketbase/pocketbase/core"
 )
 
@@ -17,13 +18,13 @@ type smartFetchState struct {
 
 // FetchAndSaveSmartDevices fetches SMART data from the agent and saves it to the database
 func (sys *System) FetchAndSaveSmartDevices() error {
-	smartData, err := sys.FetchSmartDataFromAgent()
+	response, err := sys.FetchSmartDataFromAgent()
 	if err != nil {
 		sys.recordSmartFetchResult(err, 0)
 		return err
 	}
-	err = sys.saveSmartDevices(smartData)
-	sys.recordSmartFetchResult(err, len(smartData))
+	err = sys.saveSmartDevices(response.Data, response.Complete)
+	sys.recordSmartFetchResult(err, len(response.Data))
 	return err
 }
 
@@ -61,8 +62,9 @@ func (sys *System) smartFetchInterval() time.Duration {
 	return time.Hour
 }
 
-// saveSmartDevices saves SMART device data to the smart_devices collection
-func (sys *System) saveSmartDevices(smartData map[string]smart.SmartData) error {
+// saveSmartDevices saves SMART device data and, after a complete refresh,
+// removes rows for devices that are no longer reported.
+func (sys *System) saveSmartDevices(smartData map[string]smart.SmartData, complete bool) error {
 	if len(smartData) == 0 {
 		return nil
 	}
@@ -73,20 +75,49 @@ func (sys *System) saveSmartDevices(smartData map[string]smart.SmartData) error
 		return err
 	}
 
-	for deviceKey, device := range smartData {
-		if err := sys.upsertSmartDeviceRecord(collection, deviceKey, device); err != nil {
-			return err
-		}
+	currentIDs := make(map[string]struct{}, len(smartData))
+	for deviceKey := range smartData {
+		currentIDs[makeStableHashId(sys.Id, deviceKey)] = struct{}{}
 	}
 
-	return nil
+	err = hub.RunInTransaction(func(txApp core.App) error {
+		if complete {
+			existing, err := txApp.FindRecordsByFilter(
+				collection,
+				"system = {:system}",
+				"",
+				0,
+				0,
+				dbx.Params{"system": sys.Id},
+			)
+			if err != nil {
+				return err
+			}
+			for _, record := range existing {
+				if _, ok := currentIDs[record.Id]; ok {
+					continue
+				}
+				if err := txApp.Delete(record); err != nil {
+					return err
+				}
+			}
+		}
+
+		for deviceKey, device := range smartData {
+			if err := sys.upsertSmartDeviceRecord(txApp, collection, deviceKey, device); err != nil {
+				return err
+			}
+		}
+
+		return nil
+	})
+	return err
 }
 
-func (sys *System) upsertSmartDeviceRecord(collection *core.Collection, deviceKey string, device smart.SmartData) error {
-	hub := sys.manager.hub
+func (sys *System) upsertSmartDeviceRecord(app core.App, collection *core.Collection, deviceKey string, device smart.SmartData) error {
 	recordID := makeStableHashId(sys.Id, deviceKey)
 
-	record, err := hub.FindRecordById(collection, recordID)
+	record, err := app.FindRecordById(collection, recordID)
 	if err != nil {
 		if !errors.Is(err, sql.ErrNoRows) {
 			return err
@@ -114,7 +145,7 @@ func (sys *System) upsertSmartDeviceRecord(collection *core.Collection, deviceKe
 	record.Set("cycles", powerCycles)
 	record.Set("attributes", device.Attributes)
 
-	return hub.SaveNoValidate(record)
+	return app.SaveNoValidate(record)
 }
 
 // extractPowerMetrics extracts power on hours and power cycles from SMART attributes

internal/hub/systems/system_smart_test.go

@@ -7,10 +7,53 @@ import (
 	"testing"
 	"time"
 
+	"github.com/henrygd/beszel/internal/entities/smart"
+	esystem "github.com/henrygd/beszel/internal/entities/system"
 	"github.com/henrygd/beszel/internal/hub/expirymap"
+	_ "github.com/henrygd/beszel/internal/migrations"
+
+	"github.com/pocketbase/pocketbase/core"
+	pbtests "github.com/pocketbase/pocketbase/tests"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"golang.org/x/crypto/ssh"
 )
 
+// stubHub implements hubLike using a plain pocketbase test app, so
+// smart-device DB tests can run in-package without an import cycle to
+// internal/hub (which imports this package).
+type stubHub struct{ core.App }
+
+func (stubHub) GetSSHKey(dataDir string) (ssh.Signer, error) { return nil, nil }
+func (stubHub) HandleSystemAlerts(systemRecord *core.Record, data *esystem.CombinedData) error {
+	return nil
+}
+func (stubHub) HandleStatusAlerts(status string, systemRecord *core.Record) error { return nil }
+func (stubHub) CancelPendingStatusAlerts(systemID string)                         {}
+
+// newTestSystemWithHub creates a System backed by a real (temp) database, along
+// with a matching "systems" record, for tests that need to exercise DB reads/writes.
+func newTestSystemWithHub(t *testing.T) (*System, *pbtests.TestApp) {
+	t.Helper()
+	testApp, err := pbtests.NewTestApp(t.TempDir())
+	require.NoError(t, err)
+	t.Cleanup(testApp.Cleanup)
+
+	sm := &SystemManager{hub: stubHub{testApp}, smartFetchMap: expirymap.New[smartFetchState](time.Hour)}
+	t.Cleanup(sm.smartFetchMap.StopCleaner)
+
+	col, err := testApp.FindCachedCollectionByNameOrId("systems")
+	require.NoError(t, err)
+	systemRecord := core.NewRecord(col)
+	systemRecord.Set("name", "test-system")
+	systemRecord.Set("host", "127.0.0.1")
+	systemRecord.Set("port", "45876")
+	require.NoError(t, testApp.SaveNoValidate(systemRecord))
+
+	sys := &System{Id: systemRecord.Id, manager: sm}
+	return sys, testApp
+}
+
 func TestRecordSmartFetchResult(t *testing.T) {
 	sm := &SystemManager{smartFetchMap: expirymap.New[smartFetchState](time.Hour)}
 	t.Cleanup(sm.smartFetchMap.StopCleaner)
@@ -92,3 +135,95 @@ func TestResetFailedSmartFetchState(t *testing.T) {
 	_, ok = sm.smartFetchMap.GetOk("system-1")
 	assert.True(t, ok, "expected successful smart fetch state to be preserved")
 }
+
+// countSmartDeviceRecords returns the number of smart_devices rows for the given system.
+func countSmartDeviceRecords(t *testing.T, app core.App, systemID string) []*core.Record {
+	t.Helper()
+	records, err := app.FindAllRecords("smart_devices", nil)
+	require.NoError(t, err)
+	var forSystem []*core.Record
+	for _, r := range records {
+		if r.GetString("system") == systemID {
+			forSystem = append(forSystem, r)
+		}
+	}
+	return forSystem
+}
+
+func TestSaveSmartDevices_RemovesStaleDevices(t *testing.T) {
+	sys, testApp := newTestSystemWithHub(t)
+
+	// first fetch reports two devices: sda (serial AAA) and sdb (serial BBB)
+	err := sys.saveSmartDevices(map[string]smart.SmartData{
+		"AAA": {SerialNumber: "AAA", DiskName: "sda", ModelName: "Disk A"},
+		"BBB": {SerialNumber: "BBB", DiskName: "sdb", ModelName: "Disk B"},
+	}, true)
+	require.NoError(t, err)
+
+	records := countSmartDeviceRecords(t, testApp, sys.Id)
+	require.Len(t, records, 2, "expected both devices to be saved")
+
+	var recordA *core.Record
+	for _, r := range records {
+		if r.GetString("serial") == "AAA" {
+			recordA = r
+		}
+	}
+	require.NotNil(t, recordA, "expected to find device AAA")
+	originalID := recordA.Id
+
+	deleteEvents := 0
+	testApp.OnRecordAfterDeleteSuccess("smart_devices").BindFunc(func(e *core.RecordEvent) error {
+		deleteEvents++
+		return e.Next()
+	})
+
+	// A complete refresh confirms that BBB is gone, so remove it through
+	// PocketBase and notify realtime subscribers.
+	err = sys.saveSmartDevices(map[string]smart.SmartData{
+		"AAA": {SerialNumber: "AAA", DiskName: "sda", ModelName: "Disk A", Temperature: 42},
+	}, true)
+	require.NoError(t, err)
+
+	records = countSmartDeviceRecords(t, testApp, sys.Id)
+	require.Len(t, records, 1, "expected stale device BBB to be removed")
+	assert.Equal(t, "AAA", records[0].GetString("serial"))
+	assert.Equal(t, originalID, records[0].Id, "expected existing device to be updated in place, not recreated")
+	assert.EqualValues(t, 42, records[0].GetInt("temp"))
+	assert.Equal(t, 1, deleteEvents, "expected PocketBase delete hooks to run")
+}
+
+func TestSaveSmartDevices_IncompleteDataDoesNotRemoveDevices(t *testing.T) {
+	sys, testApp := newTestSystemWithHub(t)
+
+	require.NoError(t, sys.saveSmartDevices(map[string]smart.SmartData{
+		"AAA": {SerialNumber: "AAA", DiskName: "sda"},
+		"BBB": {SerialNumber: "BBB", DiskName: "sdb"},
+	}, true))
+
+	// AAA was collected but BBB failed. The response is useful for updating AAA,
+	// but it is not authoritative enough to remove BBB.
+	require.NoError(t, sys.saveSmartDevices(map[string]smart.SmartData{
+		"AAA": {SerialNumber: "AAA", DiskName: "sda", Temperature: 42},
+	}, false))
+
+	assert.Len(t, countSmartDeviceRecords(t, testApp, sys.Id), 2)
+	recordA, err := testApp.FindRecordById("smart_devices", makeStableHashId(sys.Id, "AAA"))
+	require.NoError(t, err)
+	assert.EqualValues(t, 42, recordA.GetInt("temp"))
+}
+
+func TestSaveSmartDevices_EmptyDataIsNoop(t *testing.T) {
+	sys, testApp := newTestSystemWithHub(t)
+
+	err := sys.saveSmartDevices(map[string]smart.SmartData{
+		"AAA": {SerialNumber: "AAA", DiskName: "sda"},
+	}, true)
+	require.NoError(t, err)
+
+	err = sys.saveSmartDevices(map[string]smart.SmartData{}, true)
+	require.NoError(t, err)
+
+	records := countSmartDeviceRecords(t, testApp, sys.Id)
+	assert.Len(t, records, 1, "empty fetch result should not delete existing devices")
+}

internal/hub/transport/transport.go

@@ -88,15 +88,23 @@ func unmarshalLegacyResponse(resp common.AgentResponse, action common.WebSocketA
 		*d = *resp.String
 		return nil
 	case common.GetSmartData:
-		d, ok := dest.(*map[string]smart.SmartData)
-		if !ok {
+		switch d := dest.(type) {
+		case *map[string]smart.SmartData:
+			if resp.SmartData == nil {
+				return errors.New("no SMART data in response")
+			}
+			*d = resp.SmartData
+			return nil
+		case *smart.SmartDataResponse:
+			if resp.SmartData == nil {
+				return errors.New("no SMART data in response")
+			}
+			d.Data = resp.SmartData
+			d.Complete = resp.SmartComplete
+			return nil
+		default:
 			return fmt.Errorf("unexpected dest type for GetSmartData: %T", dest)
 		}
-		if resp.SmartData == nil {
-			return errors.New("no SMART data in response")
-		}
-		*d = resp.SmartData
-		return nil
 	case common.GetSystemdInfo:
 		d, ok := dest.(*systemd.ServiceDetails)
 		if !ok {

internal/hub/transport/transport_test.go

@@ -0,0 +1,33 @@
+package transport
+
+import (
+	"testing"
+
+	"github.com/henrygd/beszel/internal/common"
+	"github.com/henrygd/beszel/internal/entities/smart"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestUnmarshalSmartDataResponse(t *testing.T) {
+	for _, test := range []struct {
+		name     string
+		complete bool
+	}{
+		{name: "complete response", complete: true},
+		{name: "older agent defaults to incomplete", complete: false},
+	} {
+		t.Run(test.name, func(t *testing.T) {
+			response := common.AgentResponse{
+				SmartData: map[string]smart.SmartData{
+					"AAA": {SerialNumber: "AAA"},
+				},
+				SmartComplete: test.complete,
+			}
+			var result smart.SmartDataResponse
+			require.NoError(t, UnmarshalResponse(response, common.GetSmartData, &result))
+			assert.Equal(t, test.complete, result.Complete)
+			assert.Equal(t, "AAA", result.Data["AAA"].SerialNumber)
+		})
+	}
+}