Fix mDNS registration for multi-provider daemons

daemon/codex-usage-daemon.py

@@ -12,13 +12,13 @@
 
 import argparse
 import asyncio
-import atexit
 import glob
 import http.server
 import inspect
 import json
 import os
 import re
+import signal
 import socket
 import ssl
 import subprocess
@@ -33,40 +33,45 @@
 from typing import Any
 
 try:
-    from zeroconf import ServiceInfo, Zeroconf  # type: ignore[import-not-found]
+    from zeroconf import ServiceInfo  # type: ignore[import-not-found]
+    from zeroconf.asyncio import AsyncZeroconf  # type: ignore[import-not-found]
 except ImportError:  # pragma: no cover - optional
     ServiceInfo = None
-    Zeroconf = None
+    AsyncZeroconf = None
 
-# Global zeroconf instance and service info so we can unregister on exit
-_ZEROCONF: Any | None = None
+# Global AsyncZeroconf instance and service info so we can unregister on exit.
+# Cleanup is done via _async_cleanup_zeroconf() from main's finally block while
+# the asyncio loop is still alive — calling the sync wrappers after
+# asyncio.run() returns raises RuntimeError: Event loop is closed.
+_ASYNC_ZEROCONF: Any | None = None
 _SERVICE_INFO: Any | None = None
 
 
-def _cleanup_zeroconf() -> None:
-    global _ZEROCONF, _SERVICE_INFO
+async def _async_cleanup_zeroconf() -> None:
+    global _ASYNC_ZEROCONF, _SERVICE_INFO
+    if _ASYNC_ZEROCONF is None or _SERVICE_INFO is None:
+        return
+    azc, info = _ASYNC_ZEROCONF, _SERVICE_INFO
+    _ASYNC_ZEROCONF = None
+    _SERVICE_INFO = None
     try:
-        if _ZEROCONF is not None and _SERVICE_INFO is not None:
-            try:
-                _ZEROCONF.unregister_service(_SERVICE_INFO)
-            except Exception:
-                pass
-            try:
-                _ZEROCONF.close()
-            except Exception:
-                pass
-            log("Unregistered mDNS service and closed Zeroconf")
+        await azc.async_unregister_service(info)
     except Exception as exc:
-        # log may not be defined this early if imported differently; guard
         try:
-            log(f"zeroconf cleanup error: {exc}")
+            log(f"async_unregister_service failed: {type(exc).__name__}: {exc!r}")
         except Exception:
             pass
-    _ZEROCONF = None
-    _SERVICE_INFO = None
-
-
-atexit.register(_cleanup_zeroconf)
+    try:
+        await azc.async_close()
+    except Exception as exc:
+        try:
+            log(f"async_close failed: {type(exc).__name__}: {exc!r}")
+        except Exception:
+            pass
+    try:
+        log("Unregistered mDNS service and closed Zeroconf")
+    except Exception:
+        pass
 
 try:
     from bleak import BleakClient, BleakScanner  # type: ignore[import-not-found]
@@ -1514,7 +1519,7 @@ def log_message(self, format: str, *args) -> None:
         pass  # suppress HTTP access logs
 
 
-def start_http_server(port: int, shared: SharedSnapshot) -> threading.Thread:
+async def start_http_server(port: int, shared: SharedSnapshot) -> threading.Thread:
     CodexMeterHTTPHandler.shared = shared
     try:
         server = http.server.HTTPServer(("0.0.0.0", port), CodexMeterHTTPHandler)
@@ -1524,8 +1529,11 @@ def start_http_server(port: int, shared: SharedSnapshot) -> threading.Thread:
     t.start()
     log(f"HTTP server listening on http://0.0.0.0:{port} (endpoints: /usage, /status)")
 
-    # Advertise via mDNS if zeroconf is available
-    if Zeroconf is not None and ServiceInfo is not None:
+    # Advertise via mDNS if zeroconf is available. We use AsyncZeroconf and
+    # async_register_service because calling the sync register_service from
+    # within asyncio.run() triggers EventLoopBlocked on Python 3.14 +
+    # zeroconf 0.148.
+    if AsyncZeroconf is not None and ServiceInfo is not None:
         try:
             local_ip = detect_local_ip()
             if not local_ip:
@@ -1538,26 +1546,26 @@ def start_http_server(port: int, shared: SharedSnapshot) -> threading.Thread:
             except OSError:
                 log(f"mDNS discovery disabled: non-IPv4 LAN IP detected ({local_ip})")
                 return t
+            instance_name = f"codexmeter-{PROVIDER}._http._tcp.local."
             info = ServiceInfo(
                 "_http._tcp.local.",
-                "codexmeter._http._tcp.local.",
+                instance_name,
                 addresses=[mdns_addr],
                 port=port,
-                properties={"path": "/usage"},
+                properties={"path": "/usage", "provider": PROVIDER},
             )
-            # store global Zeroconf so we can unregister on exit
             try:
-                global _ZEROCONF, _SERVICE_INFO
-                _ZEROCONF = Zeroconf()
+                global _ASYNC_ZEROCONF, _SERVICE_INFO
+                _ASYNC_ZEROCONF = AsyncZeroconf()
                 _SERVICE_INFO = info
-                _ZEROCONF.register_service(info)
+                await _ASYNC_ZEROCONF.async_register_service(info)
                 log(
-                    f"mDNS advertisement active: codexmeter._http._tcp.local at http://{local_ip}:{port}"
+                    f"mDNS advertisement active: {instance_name} at http://{local_ip}:{port}"
                 )
             except Exception as exc:
-                log(f"mDNS advertise failed during register: {exc}")
+                log(f"mDNS advertise failed during register: {type(exc).__name__}: {exc!r}")
         except Exception as exc:
-            log(f"mDNS advertise failed: {exc}")
+            log(f"mDNS advertise failed: {type(exc).__name__}: {exc!r}")
     else:
         log(
             "mDNS discovery disabled: install zeroconf to advertise the daemon on the local network"
@@ -1672,32 +1680,53 @@ async def main() -> int:
         return 0
 
     shared = SharedSnapshot()
-    start_http_server(args.http_port, shared)
-
-    # Start transport in background (non-fatal — HTTP stays up regardless)
-    if args.transport == "serial":
-        asyncio.create_task(
-            transport_serial_task(find_serial_port(args.serial_port), shared)
-        )
-    elif args.transport == "ble":
-        asyncio.create_task(transport_ble_task(shared))
-
-    # One-shot mode: fetch once, write, exit
-    if args.once:
-        snapshot, source = usage_snapshot_with_source()
-        snapshot = with_provider_activity(snapshot)
-        payload = snapshot.payload()
-        shared.set_payload(payload, source)
-        log(f"Payload: {payload}")
-        await asyncio.sleep(1)  # give HTTP server a moment
-        return 0
+    await start_http_server(args.http_port, shared)
+
+    # Install a SIGTERM handler so launchctl-initiated shutdowns reach the
+    # finally block below (which sends the mDNS goodbye/unregister). SIGINT
+    # already raises KeyboardInterrupt via asyncio.run.
+    shutdown_event = asyncio.Event()
+    loop = asyncio.get_running_loop()
+    for sig in (signal.SIGTERM,):
+        try:
+            loop.add_signal_handler(sig, shutdown_event.set)
+        except (NotImplementedError, RuntimeError):
+            pass
 
-    # Main loop: keep refreshing usage for HTTP consumers
     try:
-        await run_http_loop(shared)
-    except KeyboardInterrupt:
+        # Start transport in background (non-fatal — HTTP stays up regardless)
+        if args.transport == "serial":
+            asyncio.create_task(
+                transport_serial_task(find_serial_port(args.serial_port), shared)
+            )
+        elif args.transport == "ble":
+            asyncio.create_task(transport_ble_task(shared))
+
+        # One-shot mode: fetch once, write, exit
+        if args.once:
+            snapshot, source = usage_snapshot_with_source()
+            snapshot = with_provider_activity(snapshot)
+            payload = snapshot.payload()
+            shared.set_payload(payload, source)
+            log(f"Payload: {payload}")
+            await asyncio.sleep(1)  # give HTTP server a moment
+            return 0
+
+        # Main loop: keep refreshing usage for HTTP consumers. Exit cleanly
+        # on SIGTERM or KeyboardInterrupt so async cleanup runs.
+        try:
+            main_task = asyncio.create_task(run_http_loop(shared))
+            shutdown_wait = asyncio.create_task(shutdown_event.wait())
+            done, pending = await asyncio.wait(
+                {main_task, shutdown_wait}, return_when=asyncio.FIRST_COMPLETED
+            )
+            for t in pending:
+                t.cancel()
+        except KeyboardInterrupt:
+            pass
         return 0
-    return 0
+    finally:
+        await _async_cleanup_zeroconf()
 
 
 if __name__ == "__main__":