diff --git a/daemon/codex-usage-daemon.py b/daemon/codex-usage-daemon.py index 57cb090..a26cbd1 100755 --- a/daemon/codex-usage-daemon.py +++ b/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__":