From 3dfe062ee4202c35f6cfdd6392f2f345617d1f27 Mon Sep 17 00:00:00 2001
From: henrygd <hank@henrygd.me>
Date: Fri, 25 Sep 2026 13:24:11 -0400
Subject: [PATCH] fix(agent): promote existing entry when root key is already
 registered

When the root drive is also listed in EXTRA_FILESYSTEMS, its key is
taken before addPartitionRootFs runs. Returning false there sent the
agent to addLastResortRootFs, which picks the most active device and
could register a different drive as root, overwriting that drive's
extra entry. addPartitionRootFs has already resolved the root device,
so promote the existing entry to root instead.
---
 agent/disk.go      | 16 +++++++++++-----
 agent/disk_test.go | 20 ++++++++++----------
 2 files changed, 21 insertions(+), 15 deletions(-)

diff --git a/agent/disk.go b/agent/disk.go
index ec7006f9..732872e0 100644
--- a/agent/disk.go
+++ b/agent/disk.go
@@ -155,7 +155,7 @@ func registerFilesystemStats(existing map[string]*system.FsStats, device, mountp
 
 // addFsStat inserts a discovered filesystem if it resolves to a new tracking
 // key and reports whether it was added. The key selection itself lives in
-// buildFsStatRegistration so that logic can stay directly unit-tested.
+// registerFilesystemStats so that logic can stay directly unit-tested.
 func (d *diskDiscovery) addFsStat(device, mountpoint string, root bool, customName string) bool {
 	key, fsStats, ok := registerFilesystemStats(d.agent.fsStats, device, mountpoint, root, customName, d.ctx)
 	if !ok {
@@ -212,10 +212,16 @@ func (d *diskDiscovery) addPartitionRootFs(device, mountpoint string) bool {
 	if !match {
 		return false
 	}
-	// The resolved I/O device is already known here, so use it directly to avoid
-	// a second fallback search inside buildFsStatRegistration. Report failure if
-	// the key was already taken (e.g. root drive listed in EXTRA_FILESYSTEMS) so
-	// the caller can still fall back to addLastResortRootFs.
+	// The root device is already resolved, so if it was registered earlier as an
+	// extra filesystem (e.g. root drive listed in EXTRA_FILESYSTEMS), promote that
+	// entry rather than letting addLastResortRootFs guess a different device.
+	if stats, exists := d.agent.fsStats[fs]; exists {
+		stats.Root = true
+		stats.Mountpoint = mountpoint
+		return true
+	}
+	// Use the resolved I/O device directly to avoid a second fallback search
+	// inside registerFilesystemStats.
 	return d.addFsStat(fs, mountpoint, true, "")
 }
 
diff --git a/agent/disk_test.go b/agent/disk_test.go
index 7d5a3e4d..f37f64f6 100644
--- a/agent/disk_test.go
+++ b/agent/disk_test.go
@@ -1128,11 +1128,11 @@ func TestAddPartitionRootFsWindowsDrive(t *testing.T) {
 
 func TestAddPartitionRootFsKeyAlreadyRegistered(t *testing.T) {
 	// The root drive is also listed in EXTRA_FILESYSTEMS, so its key is taken
-	// before the root fallback runs. addPartitionRootFs must report failure so
-	// the caller still falls back to addLastResortRootFs instead of ending up
-	// with no root filesystem.
+	// before the root fallback runs. The existing entry must be promoted to root
+	// rather than falling back to the most active device, which here is D:.
 	agent := &Agent{fsStats: map[string]*system.FsStats{
-		"C:": {Mountpoint: `C:\`},
+		"C:": {Mountpoint: `C:\`, Name: "System"},
+		"D:": {Mountpoint: `D:\`},
 	}}
 	discovery := diskDiscovery{
 		agent:          agent,
@@ -1141,16 +1141,16 @@ func TestAddPartitionRootFsKeyAlreadyRegistered(t *testing.T) {
 			isWindows: true,
 			diskIoCounters: map[string]disk.IOCountersStat{
 				"C:": {Name: "C:", ReadBytes: 10},
-				"D:": {Name: "D:"},
+				"D:": {Name: "D:", ReadBytes: 100},
 			},
 		},
 	}
 
 	ok := discovery.addPartitionRootFs("C:", `C:\`)
-	assert.False(t, ok)
-	assert.False(t, agent.fsStats["C:"].Root)
-
-	discovery.addLastResortRootFs()
-	assert.Len(t, agent.fsStats, 1)
+	assert.True(t, ok)
+	assert.Len(t, agent.fsStats, 2)
 	assert.True(t, agent.fsStats["C:"].Root)
+	assert.Equal(t, `C:\`, agent.fsStats["C:"].Mountpoint)
+	assert.Equal(t, "System", agent.fsStats["C:"].Name)
+	assert.False(t, agent.fsStats["D:"].Root)
 }