From a9bbb21f54bb909c58879b23cb5447430b0a1ae5 Mon Sep 17 00:00:00 2001
From: zach <zach@brohn.se>
Date: Mon, 13 Jul 2026 19:31:40 +0200
Subject: [PATCH] improve device detection when plugging/unplugging

---
 cli/src/zshell/subcommands/battery.py | 144 +++++++++++++++++++++++++-
 1 file changed, 139 insertions(+), 5 deletions(-)

diff --git a/cli/src/zshell/subcommands/battery.py b/cli/src/zshell/subcommands/battery.py
index dd189b7..e5f0568 100644
--- a/cli/src/zshell/subcommands/battery.py
+++ b/cli/src/zshell/subcommands/battery.py
@@ -292,6 +292,7 @@ class _ListenerRegistry:
 		self._rescan_lock = threading.Lock()
 		self._min_rescan_interval = 1.0
 		self._last_rescan_time = 0.0
+
 		self._hidpp_io_lock = threading.RLock()
 
 	def _touch_and_notify(self, dev_or_receiver, *args, **kwargs):
@@ -309,6 +310,7 @@ class _ListenerRegistry:
 			return
 		if not self._rescan_lock.acquire(blocking=False):
 			return
+
 		try:
 			self._last_rescan_time = now
 			with self._hidpp_io_lock:
@@ -321,9 +323,38 @@ class _ListenerRegistry:
 						if not listener._active
 					]
 					for p in dead:
+						dying_obj, _dying_listener = self._entries[p]
+
+						if dying_obj.isDevice:
+							ident = getattr(
+								dying_obj, "unitId", None
+							) or getattr(dying_obj, "serial", None)
+							info = getattr(dying_obj, "battery_info", None)
+							if ident and info is not None:
+								for (
+									other_path,
+									(other_obj, _other_listener),
+								) in self._entries.items():
+									if other_path == p or other_obj.isDevice:
+										continue
+									for child in other_obj:
+										child_ident = getattr(
+											child, "unitId", None
+										) or getattr(child, "serial", None)
+										if child_ident == ident:
+											child.set_battery_info(info)
+											break
 						del self._entries[p]
 					known_paths = set(self._entries.keys())
 
+				with self._lock:
+					id_to_entry = {}
+					for path, (obj, listener) in self._entries.items():
+						if obj.isDevice:
+							ident = obj.unitId or obj.serial
+							if ident:
+								id_to_entry[ident] = (path, obj, listener)
+
 				for dev_info in base.receivers_and_devices():
 					if dev_info.path in known_paths:
 						continue
@@ -350,17 +381,108 @@ class _ListenerRegistry:
 						continue
 
 					listener = ListenerClass(obj, self._touch_and_notify)
-					listener.start()
+
+					if obj.isDevice:
+						ident = None
+						try:
+							if obj.protocol >= 2.0:
+								obj.get_ids()
+							ident = obj.unitId or obj.serial
+						except Exception:
+							pass
+
+						if ident:
+							with self._lock:
+								existing = id_to_entry.get(ident)
+								if existing:
+									old_path, old_obj, old_listener = existing
+
+									if old_path != dev_info.path:
+										logger.info(
+											"replacing device %s (old path %s) with new path %s",
+											ident,
+											old_path,
+											dev_info.path,
+										)
+
+										if (
+											getattr(obj, "battery_info", None)
+											is None
+											and getattr(
+												old_obj, "battery_info", None
+											)
+											is not None
+										):
+											obj.set_battery_info(
+												old_obj.battery_info
+											)
+
+										old_listener.stop()
+
+										del self._entries[old_path]
+
+										self._entries[dev_info.path] = (
+											obj,
+											listener,
+										)
+										self._last_change[id(obj)] = time.time()
+
+										to_join = old_listener
+										break
+									else:
+										pass
+								else:
+									if (
+										getattr(obj, "battery_info", None)
+										is None
+									):
+										for (
+											other_path,
+											(other_obj, _other_listener),
+										) in self._entries.items():
+											if other_obj.isDevice:
+												continue
+											for child in other_obj:
+												child_ident = getattr(
+													child, "unitId", None
+												) or getattr(
+													child, "serial", None
+												)
+												if child_ident == ident:
+													obj.set_battery_info(
+														child.battery_info
+													)
+													break
+											if (
+												getattr(
+													obj,
+													"battery_info",
+													None,
+												)
+												is not None
+											):
+												break
 
 					with self._lock:
-						self._entries[dev_info.path] = (obj, listener)
-						self._last_change[id(obj)] = time.time()
+						if dev_info.path not in self._entries:
+							self._entries[dev_info.path] = (obj, listener)
+							self._last_change[id(obj)] = time.time()
+
+					listener.start()
 
 					logger.info(
 						"listening on %s (%s)",
 						dev_info.path,
 						"device" if dev_info.isDevice else "receiver",
 					)
+
+					if "to_join" in locals():
+						try:
+							to_join.join(timeout=1.0)
+						except Exception:
+							logger.exception("error joining replaced listener")
+						del to_join
+
 		finally:
 			self._rescan_lock.release()
 
@@ -610,10 +732,22 @@ def daemon(
 	pending_write_lock = threading.Lock()
 	WRITE_SETTLE_SECONDS = 0.3
 
+	write_seq_lock = threading.Lock()
+	write_seq = 0
+	last_written_seq = 0
+
 	def write_if_changed():
-		nonlocal last_payload
+		nonlocal last_payload, write_seq, last_written_seq
+		with write_seq_lock:
+			write_seq += 1
+			my_seq = write_seq
+
+		results = _snapshot(registry)
 		with write_lock:
-			results = _snapshot(registry)
+			with write_seq_lock:
+				if my_seq <= last_written_seq:
+					return
+				last_written_seq = my_seq
 			payload = [asdict(r) for r in results]
 			if payload != last_payload:
 				_write_json_atomic(out, results)