fix: propagate re-registered node id and make rebind thread-safe

src/server/clients/redis.py

@@ -18,15 +18,13 @@
 def _keepalive_kwargs() -> dict[str, Any]:
     """Connection kwargs that keep an idle Redis socket alive.
 
-    A pub/sub SUBSCRIBE that receives no messages is an idle TCP connection.
-    Intermediaries (load balancers, NAT gateways) cull idle server-side
-    connections on their own timeout, silently dropping a node's dispatch
-    subscription so the publisher sees zero receivers. TCP keepalive probes
-    keep the connection observably alive well below such timeouts, and
-    ``health_check_interval`` makes the client re-establish a connection it
-    finds dead rather than blocking on a stale socket.
+    An idle pub/sub SUBSCRIBE gets culled by intermediaries (load balancers, NAT
+    gateways), dropping a node's dispatch subscription. Keepalive probes hold it open
+    and ``health_check_interval`` reconnects a dead socket instead of blocking on it.
     """
     options: dict[int, int] = {}
+    # These constants are Linux-only; skip any the running platform lacks (e.g.
+    # macOS/Windows dev hosts) so importing this module there doesn't fail.
     for name, value in (
         ("TCP_KEEPIDLE", 60),
         ("TCP_KEEPINTVL", 15),
@@ -148,6 +146,25 @@ def ssh_connection_key(connection_id: str) -> str:
 SSH_CONNECTION_IDS_KEY = "ssh:connections:active"
 
 
+def parse_pubsub_message(msg: dict[str, Any] | None) -> Any | None:
+    """Decode a redis-py pub/sub frame into its JSON payload.
+
+    Returns ``None`` for control frames, empty payloads, and malformed JSON. Payloads
+    are always JSON objects, so a ``None`` return means "nothing to deliver".
+    """
+    if msg is None or msg.get("type") != "message":
+        return None
+    raw = msg.get("data")
+    if raw is None:
+        return None
+    if isinstance(raw, bytes):
+        raw = raw.decode("utf-8", errors="ignore")
+    try:
+        return json.loads(raw)
+    except json.JSONDecodeError:
+        return None
+
+
 def iter_pubsub_messages(pubsub: PubSub) -> Iterable[Any]:
     """Iterate over messages from a Redis PubSub instance.
 
@@ -157,17 +174,10 @@ def iter_pubsub_messages(pubsub: PubSub) -> Iterable[Any]:
     """
     try:
         for msg in pubsub.listen():
-            if msg.get("type") != "message":
-                continue
-            raw = msg.get("data")
-            if raw is None:
-                continue
-            if isinstance(raw, bytes):
-                raw = raw.decode("utf-8", errors="ignore")
-            try:
-                yield json.loads(raw)
-            except json.JSONDecodeError:
+            parsed = parse_pubsub_message(msg)
+            if parsed is None:
                 continue
+            yield parsed
     except (ConnectionError, OSError):
         return
 
@@ -280,6 +290,9 @@ def ttl(self, key: str) -> float | None:
     def incr(self, key: str) -> int:
         return int(_sync(self._control.incr(key)))
 
+    def eval(self, script: str, numkeys: int, *keys_and_args: str) -> str:
+        return _sync(self._control.eval(script, numkeys, *keys_and_args))
+
     def set_value(self, key: str, value: str) -> None:
         self._control.set(key, value)
 
@@ -483,6 +496,9 @@ async def ttl(self, key: str) -> float | None:
     async def incr(self, key: str) -> int:
         return int(await _awaitable(self._control.incr(key)))
 
+    async def eval(self, script: str, numkeys: int, *keys_and_args: str) -> str:
+        return await _awaitable(self._control.eval(script, numkeys, *keys_and_args))
+
     async def set_value(self, key: str, value: str) -> None:
         await _awaitable(self._control.set(key, value))
 

src/server/main.py

@@ -355,11 +355,18 @@ async def _lifespan(_: FastAPI):
         # --- Supervisor (all nodes with worker management) ---
         if SUPERVISOR is not None:
             await SUPERVISOR.start(system_principal)
-            app.state.node_id = SUPERVISOR.node_id
-            # Tell EventMonitor which node this server belongs to so that it can wait
-            # for the supervisor's SV_UNREGISTER event on shutdown.
-            if EVENT_MONITOR is not None:
-                EVENT_MONITOR.set_own_node(SUPERVISOR.node_id)
+
+            def _on_node_id_change(new_node_id: str) -> None:
+                app.state.node_id = new_node_id
+                # Tell EventMonitor which node this server belongs to so that it can
+                # wait for the supervisor's SV_UNREGISTER event on shutdown.
+                if EVENT_MONITOR is not None:
+                    EVENT_MONITOR.set_own_node(new_node_id)
+
+            # Keep the node_id updated for request auth scope + shutdown
+            # self-identification.
+            _on_node_id_change(SUPERVISOR.node_id)
+            SUPERVISOR.add_node_id_listener(_on_node_id_change)
 
         # --- Startup reconcile ---
         # Runs after the supervisor handshake so this node is in NODE_REGISTRY

src/server/supervisor/services/command_listener.py

@@ -5,7 +5,7 @@
 import secrets
 from collections.abc import Callable, Iterable
 from concurrent import futures
-from threading import Thread
+from threading import Event, Lock, Thread
 
 from pydantic import ValidationError
 from redis.client import PubSub
@@ -19,8 +19,8 @@
 from ...clients.redis import (
     NODE_RESPONSE_CHANNEL,
     SyncRedisClient,
-    iter_pubsub_messages,
     node_cmd_channel,
+    parse_pubsub_message,
 )
 from ...utils.concurrent import Sentinel, TaskReceiver
 from ..adapters.docker import DockerWorkerConfig
@@ -36,6 +36,7 @@
 _CREATE_WORKER_TIMEOUT = 600.0
 _DESTROY_WORKER_TIMEOUT = 60.0
 _DESTROY_WORKERS_TIMEOUT = 120.0
+_POLL_TIMEOUT_SEC = 0.25
 
 
 def _cmd_receiver_loop(
@@ -53,25 +54,6 @@ def send_response(resp: CommandResponse) -> None:
         q.put((cmd, make_handler(cmd_id)))
 
 
-def _pubsub_loop(
-    redis: SyncRedisClient,
-    pubsub: PubSub,
-    q: queue.Queue[tuple[CommandMessage, ResponseHandler]],
-    logger: logging.Logger,
-) -> None:
-    for data in iter_pubsub_messages(pubsub):
-        try:
-            cmd = CommandMessage.model_validate(data)
-        except ValidationError as e:
-            logger.error("Invalid command message: %s", e)
-            continue
-
-        def send_response(resp: CommandResponse) -> None:
-            redis.publish_control(NODE_RESPONSE_CHANNEL, resp.model_dump_json())
-
-        q.put((cmd, send_response))
-
-
 class _CommandStream:
     _SENTINEL = Sentinel()
 
@@ -82,34 +64,84 @@ def __init__(
         cmd_receiver: TaskReceiver[CommandMessage, CommandResponse] | None,
         logger: logging.Logger,
     ):
+        # node_id is mutated only by the pubsub reader thread once running.
         self.node_id = node_id
         self.redis = redis
         self.cmd_receiver = cmd_receiver
         self.logger = logger
         self._cmd_queue: (
             queue.Queue[tuple[CommandMessage, ResponseHandler] | Sentinel] | None
         ) = None
-        self._pubsub: PubSub | None = None
+        self._pubsub_running = False
+
+        self._rebind_lock = Lock()
+        self._pending_node_id: str | None = None
+        self._rebind_applied = Event()
 
     def close(self) -> None:
         if self._cmd_queue is not None:
             self._cmd_queue.put(self._SENTINEL)
             self._cmd_queue = None
 
     def rebind(self, node_id: str) -> None:
-        """Move the command subscription to a new node id on the live pubsub."""
+        """Request moving the command subscription to a new node id.
+
+        Records the target under a lock; the reader thread applies the actual
+        ``subscribe``/``unsubscribe`` between polls (mutating redis-py ``PubSub`` is not
+        thread-safe). ``wait_rebound`` blocks until the switch has taken effect.
+        """
         if node_id == self.node_id:
+            self._rebind_applied.set()
             return
-        old_node_id = self.node_id
-        self.node_id = node_id
-        pubsub = self._pubsub
-        if pubsub is None:
-            return
-        pubsub.subscribe(node_cmd_channel(node_id))
-        pubsub.unsubscribe(node_cmd_channel(old_node_id))
+        self._rebind_applied.clear()
+        with self._rebind_lock:
+            self._pending_node_id = node_id
+
+    def wait_rebound(self, timeout: float) -> bool:
+        return self._rebind_applied.wait(timeout)
+
+    def _apply_pending_rebind(self, pubsub: PubSub, current_id: str) -> str:
+        with self._rebind_lock:
+            pending = self._pending_node_id
+            self._pending_node_id = None
+        if pending is None or pending == current_id:
+            return current_id
+        pubsub.subscribe(node_cmd_channel(pending))
+        pubsub.unsubscribe(node_cmd_channel(current_id))
+        self.node_id = pending
+        self._rebind_applied.set()
         self.logger.info(
-            "Command stream rebound from node %s to %s", old_node_id, node_id
+            "Command stream rebound from node %s to %s", current_id, pending
         )
+        return pending
+
+    def _run_pubsub(
+        self,
+        pubsub: PubSub,
+        cmd_queue: queue.Queue[tuple[CommandMessage, ResponseHandler] | Sentinel],
+    ) -> None:
+        def send_response(resp: CommandResponse) -> None:
+            self.redis.publish_control(NODE_RESPONSE_CHANNEL, resp.model_dump_json())
+
+        current_id = self.node_id
+        try:
+            while self._pubsub_running:
+                current_id = self._apply_pending_rebind(pubsub, current_id)
+                msg = pubsub.get_message(timeout=_POLL_TIMEOUT_SEC)
+                data = parse_pubsub_message(msg)
+                if data is None:
+                    continue
+                try:
+                    cmd = CommandMessage.model_validate(data)
+                except ValidationError as e:
+                    self.logger.error("Invalid command message: %s", e)
+                    continue
+                cmd_queue.put((cmd, send_response))
+        except (ConnectionError, OSError):
+            return
+        except Exception as exc:
+            if self._pubsub_running:
+                self.logger.exception("Command pubsub loop error: %s", exc)
 
     def iter_stream(self) -> Iterable[tuple[CommandMessage, ResponseHandler]]:
         if self._cmd_queue is not None:
@@ -127,10 +159,10 @@ def iter_stream(self) -> Iterable[tuple[CommandMessage, ResponseHandler]]:
             task_thread.start()
 
         pubsub = self.redis.subscribe_control(node_cmd_channel(self.node_id))
-        self._pubsub = pubsub
+        self._pubsub_running = True
         pubsub_thread = Thread(
-            target=_pubsub_loop,
-            args=(self.redis, pubsub, cmd_queue, self.logger),
+            target=self._run_pubsub,
+            args=(pubsub, cmd_queue),
             name="CommandPubSubThread",
             daemon=True,
         )
@@ -144,7 +176,8 @@ def iter_stream(self) -> Iterable[tuple[CommandMessage, ResponseHandler]]:
                     break
                 yield item
         finally:
-            self._pubsub = None
+            self._pubsub_running = False
+            pubsub_thread.join()
             pubsub.close()
 
 
@@ -218,12 +251,18 @@ def start(self) -> None:
         )
 
     def rebind(self, node_id: str) -> None:
-        """Move the command subscription to a new node id without a restart."""
+        """Request moving the command subscription to a new node id.
+
+        The reader thread applies the switch; ``wait_rebound`` blocks until it has taken
+        effect.
+        """
         self._node_id = node_id
-        stream = self._cmd_stream
-        if stream is not None:
+        if stream := self._cmd_stream:
             stream.rebind(node_id)
 
+    def wait_rebound(self, timeout: float) -> bool:
+        return stream.wait_rebound(timeout) if (stream := self._cmd_stream) else True
+
     async def stop(self) -> None:
         if self._thread is None or self._cmd_stream is None:
             self.logger.warning("Command listener not started")