From 211145365308797e8e4424996ccb886141f3b2bf Mon Sep 17 00:00:00 2001
From: Graham McIntire <graham@mcintire.me>
Date: Fri, 6 Feb 2026 13:18:23 -0600
Subject: [PATCH] Fix TUI

---
 src/snmp/poller_registry.rs | 10 ++++++++--
 src/tui/events.rs           |  4 ++--
 src/tui/state.rs            | 14 +++++++-------
 src/websocket_client.rs     | 14 ++++++++------
 4 files changed, 25 insertions(+), 17 deletions(-)

diff --git a/src/snmp/poller_registry.rs b/src/snmp/poller_registry.rs
index f714e4e..6a77647 100644
--- a/src/snmp/poller_registry.rs
+++ b/src/snmp/poller_registry.rs
@@ -51,15 +51,20 @@ impl PollerRegistry {
 
     /// Remove a device poller (shutdown thread)
     /// Called when a device is deleted or no longer needs polling
-    pub fn remove(&self, device_id: &str) {
+    /// Returns the device IP if the poller was found
+    pub fn remove(&self, device_id: &str) -> Option<String> {
         let mut pollers = self.pollers.write().unwrap();
         if let Some(poller) = pollers.remove(device_id) {
+            let ip = poller.config().ip.clone();
             poller.shutdown();
             tracing::info!(
                 "Removed device poller for {} (remaining: {})",
                 device_id,
                 pollers.len()
             );
+            Some(ip)
+        } else {
+            None
         }
     }
 
@@ -121,8 +126,9 @@ mod tests {
         assert_eq!(poller.device_id(), "test-device");
 
         // Remove the poller
-        registry.remove("test-device");
+        let removed_ip = registry.remove("test-device");
         assert_eq!(registry.count(), 0);
+        assert_eq!(removed_ip, Some("127.0.0.1".to_string()));
     }
 
     #[test]
diff --git a/src/tui/events.rs b/src/tui/events.rs
index 9edb5ee..6870157 100644
--- a/src/tui/events.rs
+++ b/src/tui/events.rs
@@ -30,11 +30,11 @@ pub enum AgentEvent {
         sentence_count: usize,
     },
     PollerCreated {
-        device_id: String,
+        device_ip: String,
         total_count: usize,
     },
     PollerRemoved {
-        device_id: String,
+        device_ip: String,
         total_count: usize,
     },
     HeartbeatSent,
diff --git a/src/tui/state.rs b/src/tui/state.rs
index 4e35378..76aa743 100644
--- a/src/tui/state.rs
+++ b/src/tui/state.rs
@@ -106,27 +106,27 @@ impl AgentState {
                 )
             }
             AgentEvent::PollerCreated {
-                device_id,
+                device_ip,
                 total_count,
             } => {
                 self.active_pollers = *total_count;
-                if !self.active_devices.contains(device_id) {
-                    self.active_devices.push(device_id.clone());
+                if !self.active_devices.contains(device_ip) {
+                    self.active_devices.push(device_ip.clone());
                 }
                 format!(
                     "Poller created for {} ({} total pollers)",
-                    device_id, total_count
+                    device_ip, total_count
                 )
             }
             AgentEvent::PollerRemoved {
-                device_id,
+                device_ip,
                 total_count,
             } => {
                 self.active_pollers = *total_count;
-                self.active_devices.retain(|d| d != device_id);
+                self.active_devices.retain(|d| d != device_ip);
                 format!(
                     "Poller removed for {} ({} total pollers)",
-                    device_id, total_count
+                    device_ip, total_count
                 )
             }
             AgentEvent::HeartbeatSent => {
diff --git a/src/websocket_client.rs b/src/websocket_client.rs
index 7b81cc8..8477046 100644
--- a/src/websocket_client.rs
+++ b/src/websocket_client.rs
@@ -471,14 +471,16 @@ impl AgentClient {
                     "Removing poller for device no longer in job list: {}",
                     device_id
                 );
-                self.poller_registry.remove(&device_id);
+                let device_ip = self.poller_registry.remove(&device_id);
 
                 #[cfg(feature = "tui")]
                 if let Some(ref bus) = self.event_bus {
-                    let _ = bus.send(crate::tui::AgentEvent::PollerRemoved {
-                        device_id: device_id.clone(),
-                        total_count: self.poller_registry.count(),
-                    });
+                    if let Some(ip) = device_ip {
+                        let _ = bus.send(crate::tui::AgentEvent::PollerRemoved {
+                            device_ip: ip,
+                            total_count: self.poller_registry.count(),
+                        });
+                    }
                 }
             }
         }
@@ -873,7 +875,7 @@ async fn execute_snmp_job(
         if count_after > count_before {
             if let Some(ref bus) = event_bus {
                 let _ = bus.send(crate::tui::AgentEvent::PollerCreated {
-                    device_id: job.device_id.clone(),
+                    device_ip: snmp_device.ip.clone(),
                     total_count: count_after,
                 });
             }