fix: MQTT threading model and event emitter correctness (#93)

* fix: MQTT threading model and event emitter correctness

- Marshal MQTT message parse/dispatch from the AWS CRT network thread
  onto the asyncio event loop; keeps the SDK callback trivial, prevents
  blocking user callbacks from stalling MQTT processing, and eliminates
  the handler-registry iteration race during reconnect/resubscribe
- Iterate snapshots of handler registries so handlers can
  subscribe/unsubscribe during dispatch
- Isolate individual handler failures: catch Exception per handler with
  logging instead of only (TypeError, AttributeError, KeyError)
- Make the unit system preference process-wide instead of a ContextVar;
  tasks scheduled from CRT threads never inherited the caller's context,
  so --unit-system was silently ignored for all MQTT-delivered data
- Remove once-listeners before invoking them: raising callbacks no
  longer stay registered forever and concurrent emits no longer
  double-fire
- Clean up wait_for() listeners in a finally block so cancellation does
  not leak them; fix wait_for docstring examples that broke user code

* test: reset unit system before and after each threading-model test

Addresses review feedback: the preference is process-wide, so a
teardown-only reset let the first test in this module inherit values
set by earlier tests in the session.

---------

Co-authored-by: emmanuel <emmanuel@musta.lan>

Author: eman

CHANGELOG.rst

@@ -7,6 +7,44 @@ Unreleased
 
 Bug Fixes
 ---------
+- **Run MQTT message dispatch on the event loop**: JSON parsing, pydantic
+  model validation, and user callbacks all executed directly on the AWS
+  CRT network thread. A slow or blocking callback (e.g. the CLI monitor's
+  CSV writes) stalled all MQTT message processing, user callbacks ran on
+  an undocumented SDK thread where asyncio operations are unsafe, and the
+  handler registry could be mutated on the event loop while the CRT
+  thread iterated it (``RuntimeError: dictionary changed size during
+  iteration``, dropping the message — most likely during the
+  reconnect/resubscribe window). The awscrt callback now only marshals
+  the raw payload onto the event loop; parsing and dispatch run there,
+  iterating snapshots of the handler registries.
+- **Stop one raising handler from aborting message delivery**: handler
+  dispatch caught only ``(TypeError, AttributeError, KeyError)``; a user
+  callback raising anything else (e.g. ``ValueError``) escaped into the
+  awscrt callback machinery and skipped the remaining handlers for that
+  message. Individual handler failures are now logged and isolated.
+- **Make the unit system preference process-wide**:
+  ``set_unit_system()`` stored the preference in a ``ContextVar`` set in
+  the caller's task. Tasks scheduled from AWS CRT callback threads never
+  inherit that context, so the preference silently reverted to
+  auto-detect for all MQTT-delivered data — ``nwp-cli --unit-system
+  metric monitor`` logged temperatures in the device's native unit while
+  labeling them °C. The preference is now a process-wide setting visible
+  from every task and thread.
+- **Fix once-listeners firing more than once**: ``EventEmitter.emit()``
+  removed one-time listeners only after invoking them, so a callback
+  that raised stayed registered forever, and two overlapping emits could
+  both fire the same once-listener. Once-listeners are now removed
+  before invocation.
+- **Fix wait_for() leaking its listener on cancellation**:
+  ``EventEmitter.wait_for()`` removed its listener only on timeout; a
+  cancelled waiter left the listener registered until the event next
+  fired, setting a result on a dead future. Cleanup now happens in a
+  ``finally`` block.
+- **Fix wait_for() docstring examples**: examples showed
+  ``args, _ = await emitter.wait_for(...)``, but ``wait_for`` returns
+  just the args tuple — following the documented pattern raised
+  ``ValueError`` or silently mis-assigned.
 - **Serialize concurrent token refresh**: ``ensure_valid_token()`` and
   ``refresh_token()`` had no lock, so concurrent callers (API 401 retry,
   MQTT reconnect, periodic requests) at token expiry fired parallel

src/nwp500/events.py

@@ -245,7 +245,21 @@ async def emit(self, event: str, *args: Any, **kwargs: Any) -> int:
         listeners: list[EventListener] = self._listeners[event].copy()
         # Copy to allow modification during iteration
         called_count = 0
-        listeners_to_remove: list[EventListener] = []
+
+        # Remove one-time listeners BEFORE invoking them (synchronously,
+        # in the same event-loop slice as the snapshot):
+        # - a raising callback must not stay registered forever
+        # - concurrent emits each snapshot the list; deferring removal
+        #   until after (awaited) iteration double-fires once-listeners
+        for listener in listeners:
+            if listener.once:
+                if listener in self._listeners.get(event, []):
+                    self._listeners[event].remove(listener)
+                self._once_callbacks.discard((event, listener.callback))
+
+        # Clean up if no listeners left
+        if event in self._listeners and not self._listeners[event]:
+            del self._listeners[event]
 
         for listener in listeners:
             try:
@@ -257,11 +271,6 @@ async def emit(self, event: str, *args: Any, **kwargs: Any) -> int:
 
                 called_count += 1
 
-                # Check if this is a once listener
-                if listener.once:
-                    listeners_to_remove.append(listener)
-                    self._once_callbacks.discard((event, listener.callback))
-
             except Exception as e:
                 # Catch all exceptions from user callbacks to ensure
                 # resilience. We intentionally catch Exception here because:
@@ -273,15 +282,6 @@ async def emit(self, event: str, *args: Any, **kwargs: Any) -> int:
                     exc_info=True,
                 )
 
-        # Remove one-time listeners after iteration
-        for listener in listeners_to_remove:
-            if listener in self._listeners[event]:
-                self._listeners[event].remove(listener)
-
-        # Clean up if no listeners left
-        if not self._listeners[event]:
-            del self._listeners[event]
-
         # Track event count
         self._event_counts[event] += 1
 
@@ -391,7 +391,7 @@ async def wait_for(
             await emitter.wait_for('device_ready', timeout=30)
 
             # Wait for specific condition
-            args, _ = await emitter.wait_for('temperature_changed')
+            args = await emitter.wait_for('temperature_changed')
             temperature_event = args[0]
             current_temp = temperature_event.new_temperature
         """
@@ -416,7 +416,9 @@ def handler(*args: Any, **kwargs: Any) -> None:
             # Return just args for simplicity (most common case)
             return args_tuple
 
-        except TimeoutError:
-            # Remove the listener on timeout
+        finally:
+            # Always unregister: on timeout AND on cancellation. A leaked
+            # once-listener would otherwise stay registered until the
+            # event next fires (setting a result on a dead future).
+            # No-op when the listener already fired (once() removed it).
             self.off(event, handler)
-            raise

src/nwp500/mqtt/subscriptions.py

@@ -122,37 +122,61 @@ def update_connection(self, connection: Any) -> None:
     def _on_message_received(
         self, topic: str, payload: bytes, **kwargs: Any
     ) -> None:
-        """Handle received MQTT messages.
+        """Handle received MQTT messages (AWS CRT callback).
 
-        Parses JSON payload and routes to registered handlers.
+        This runs on the AWS CRT network thread. It must stay trivial:
+        JSON parsing, pydantic validation, and user callbacks are
+        marshaled onto the asyncio event loop so that
+
+        - a slow or blocking user callback cannot stall all MQTT message
+          processing,
+        - handler registries mutated on the event loop are never iterated
+          concurrently from this thread, and
+        - user callbacks run where asyncio operations are safe, instead
+          of on an undocumented SDK thread.
 
         Args:
             topic: MQTT topic the message was received on
             payload: Raw message payload (JSON bytes)
             **kwargs: Additional MQTT metadata
         """
+        self._schedule_coroutine(self._dispatch_message(topic, payload))
+
+    async def _dispatch_message(self, topic: str, payload: bytes) -> None:
+        """Parse a message and route it to registered handlers.
+
+        Runs on the event loop.
+
+        Args:
+            topic: MQTT topic the message was received on
+            payload: Raw message payload (JSON bytes)
+        """
         try:
-            # Parse JSON payload
             message = json.loads(payload.decode("utf-8"))
+        except (json.JSONDecodeError, UnicodeDecodeError) as e:
+            _logger.error(f"Failed to parse message payload: {e}")
+            return
+
+        if _logger.isEnabledFor(logging.DEBUG):
             _logger.debug("Received message on topic: %s", redact_topic(topic))
 
-            # Call registered handlers that match this topic
-            # Need to match against subscription patterns with wildcards
-            for (
-                subscription_pattern,
-                handlers,
-            ) in self._message_handlers.items():
-                if topic_matches_pattern(topic, subscription_pattern):
-                    for handler in handlers:
-                        try:
-                            handler(topic, message)
-                        except (TypeError, AttributeError, KeyError) as e:
-                            _logger.error(f"Error in message handler: {e}")
-
-        except json.JSONDecodeError as e:
-            _logger.error(f"Failed to parse message payload: {e}")
-        except (AttributeError, KeyError, TypeError) as e:
-            _logger.error(f"Error processing message: {e}")
+        # Call registered handlers that match this topic.
+        # Iterate over snapshots so handlers can subscribe/unsubscribe
+        # (mutating the registries) during dispatch.
+        for subscription_pattern, handlers in list(
+            self._message_handlers.items()
+        ):
+            if topic_matches_pattern(topic, subscription_pattern):
+                for handler in list(handlers):
+                    try:
+                        handler(topic, message)
+                    except Exception:
+                        # One bad handler must not abort delivery to the
+                        # remaining handlers for this message.
+                        _logger.exception(
+                            "Error in message handler for %s",
+                            redact_topic(topic),
+                        )
 
     async def subscribe(
         self,

src/nwp500/mqtt_events.py

@@ -204,7 +204,7 @@ class MqttClientEvents:
         )
 
         # Wait for a specific event
-        args, _ = await mqtt_client.wait_for(
+        args = await mqtt_client.wait_for(
             MqttClientEvents.CONNECTION_RESUMED
         )
         connection_event = args[0]

src/nwp500/unit_system.py

@@ -1,16 +1,18 @@
 """Unit system management for temperature, flow rate, and volume conversions.
 
-This module provides context-based unit system management, allowing applications
-to override the device's temperature_type setting and specify a preferred
-measurement system (Metric or Imperial).
+This module provides process-wide unit system management, allowing
+applications to override the device's temperature_type setting and specify a
+preferred measurement system (Metric or Imperial).
 
 The unit system preference can be set at library initialization and is used
-during model validation to convert device values to the user's preferred units.
+during model validation to convert device values to the user's preferred
+units. The preference is a process-wide setting: it is visible from every
+task and thread, including model validation triggered by MQTT callbacks that
+run outside the task that configured it.
 """
 
 from __future__ import annotations
 
-import contextvars
 import logging
 from typing import Literal
 
@@ -25,21 +27,25 @@
 # Type alias for unit system preference
 UnitSystemType = Literal["metric", "us_customary"] | None
 
-# Context variable to store the preferred unit system
+# Process-wide preferred unit system.
 # None means auto-detect from device
 # "metric" means Celsius, "us_customary" means Fahrenheit
-_unit_system_context: contextvars.ContextVar[
-    Literal["metric", "us_customary"] | None
-] = contextvars.ContextVar("unit_system", default=None)
+#
+# This is intentionally NOT a contextvars.ContextVar: MQTT message handling
+# runs in tasks scheduled from AWS CRT callback threads whose context never
+# inherits from the application task, so a context-local preference silently
+# reverted to auto-detect for all real-time data.
+_unit_system: Literal["metric", "us_customary"] | None = None
 
 
 def set_unit_system(
     unit_system: Literal["metric", "us_customary"] | None,
 ) -> None:
     """Set preferred unit system for temperature, flow, and volume conversions.
 
-    This setting overrides the device's temperature_type setting and applies to
-    all subsequent model validation operations in the current async context.
+    This setting overrides the device's temperature_type setting and applies
+    process-wide to all subsequent model validation operations, including
+    those triggered by MQTT callbacks running in other tasks or threads.
 
     Args:
         unit_system: Preferred unit system:
@@ -52,12 +58,9 @@ def set_unit_system(
         >>> set_unit_system("us_customary")
         >>> # All values now in F, GPM, Gallons
         >>> set_unit_system(None)  # Reset to auto-detect
-
-    Note:
-        This is context-aware and works with async code. Each async task
-        maintains its own unit system preference.
     """
-    _unit_system_context.set(unit_system)
+    global _unit_system
+    _unit_system = unit_system
 
 
 def get_unit_system() -> Literal["metric", "us_customary"] | None:
@@ -69,7 +72,7 @@ def get_unit_system() -> Literal["metric", "us_customary"] | None:
             - "us_customary": Fahrenheit, GPM, Gallons
             - None: Auto-detect from device (default)
     """
-    return _unit_system_context.get()
+    return _unit_system
 
 
 def reset_unit_system() -> None:
@@ -78,7 +81,7 @@ def reset_unit_system() -> None:
     This is useful for tests or when switching between different
     device configurations.
     """
-    _unit_system_context.set(None)
+    set_unit_system(None)
 
 
 def unit_system_to_temperature_type(
@@ -108,8 +111,8 @@ def is_metric_preferred(
 ) -> bool:
     """Check if metric (Celsius) is preferred.
 
-    Checks the override first, then falls back to the context-configured
-    unit system. Used during validation to determine preferred units.
+    Checks the override first, then falls back to the configured unit
+    system. Used during validation to determine preferred units.
 
     Args:
         override: Optional override value. If provided, this takes precedence
@@ -123,7 +126,7 @@ def is_metric_preferred(
     if override is not None:
         return override == "metric"
 
-    # Otherwise check context
+    # Otherwise check the configured preference
     unit_system = get_unit_system()
     if unit_system is not None:
         return unit_system == "metric"