fix thread leaks

cosmos-framework/cosmos_framework/data/vfm/action/action_viz/app.py

@@ -356,6 +356,7 @@ def launch_action_viz(config: AppConfig) -> None:
         renderer = UnifiedRenderer(client)
         sample_state: dict[str, LoadedSample | None] = {"loaded": None}
         playing = {"enabled": False, "last_frame_time": _time.time()}
+        play_stop_event = threading.Event()
         generation_thread = {"thread": None}
         generation_lock = threading.Lock()
         control_handles: dict[int, Any] = {}
@@ -1161,10 +1162,8 @@ def launch_action_viz(config: AppConfig) -> None:
                 if loaded is not None:
                     _set_action_plot(plot_handle, loaded.baked_action, channel_visible, loaded.adapter.channel_names)
 
-        load_sample(int(datasets[config.default_dataset].initial_index))
-
         def _play_loop() -> None:
-            while True:
+            while not play_stop_event.wait(timeout=0.02):
                 loaded = sample_state["loaded"]
                 if loaded is not None and playing["enabled"]:
                     now = _time.time()
@@ -1172,20 +1171,73 @@ def launch_action_viz(config: AppConfig) -> None:
                     if now - playing["last_frame_time"] >= frame_period:
                         time_slider.value = (int(time_slider.value) + 1) % max(len(loaded.baked_action), 1)
                         playing["last_frame_time"] = now
-                _time.sleep(0.02)
 
-        threading.Thread(target=_play_loop, daemon=True).start()
+        play_thread = threading.Thread(
+            target=_play_loop,
+            name=f"action-viz-playback-client-{client.client_id}",
+            daemon=True,
+        )
+        # Register the session before starting the thread so a disconnect that
+        # races ahead of this handler can always find the session, set the stop
+        # event, and stop the workers. Otherwise the play thread and IK worker
+        # leak permanently for that client.
         with sessions_lock:
-            sessions[client.client_id] = {"client": client, "scene_update_worker": scene_update_worker}
+            sessions[client.client_id] = {
+                "client": client,
+                "scene_update_worker": scene_update_worker,
+                "play_stop_event": play_stop_event,
+                "play_thread": play_thread,
+            }
+        play_thread.start()
+
+        if not play_stop_event.is_set():
+            load_sample(int(datasets[config.default_dataset].initial_index))
 
     @server.on_client_disconnect
     def _(client) -> None:
         with sessions_lock:
             session = sessions.pop(client.client_id, None)
         if session is not None:
+            play_stop_event = session.get("play_stop_event")
+            if isinstance(play_stop_event, threading.Event):
+                play_stop_event.set()
             worker = session.get("scene_update_worker")
             if isinstance(worker, _LatestOnlyIkWorker):
                 worker.stop()
+            play_thread = session.get("play_thread")
+            # is_alive() guards against a disconnect that races in before the
+            # connect handler called start(): join() would raise RuntimeError on
+            # an unstarted thread. Once started, the thread exits promptly since
+            # the stop event is already set, so it needs no join here.
+            if isinstance(play_thread, threading.Thread) and play_thread.is_alive():
+                play_thread.join(timeout=1.0)
+
+        # Orphan-leak watchdog. A disconnect that lands before the connect
+        # handler registers its session (sessions.pop -> None above) can't stop
+        # the per-client IK/playback threads it would otherwise own, leaking an
+        # idle daemon thread. We chose not to close that narrow race, so instead
+        # surface accumulation: count live "action-viz-*client-<id>" threads whose
+        # client_id has no active session. A handful are normal transients from
+        # clients mid-connect (the IK thread starts before registration); a
+        # sustained count above the threshold means the race is being hit often
+        # enough to revisit the two-phase-registration fix. Runs after the
+        # joins above so this client's own threads are gone, not miscounted.
+        orphan_warn_threshold = 8
+        with sessions_lock:
+            active_ids = {str(client_id) for client_id in sessions}
+        orphan_threads = [
+            thread.name
+            for thread in threading.enumerate()
+            if thread.is_alive()
+            and thread.name.startswith("action-viz-")
+            and thread.name.rsplit("client-", 1)[-1] not in active_ids
+        ]
+        if len(orphan_threads) >= orphan_warn_threshold:
+            log.warning(
+                f"action-viz: {len(orphan_threads)} orphaned client threads alive with "
+                f"no active session (e.g. {orphan_threads[:4]}) across {len(active_ids)} "
+                f"active sessions; connect/disconnect race thread leak may be accumulating."
+            )
 
     log.info(f"Action viz ready at http://0.0.0.0:{config.port}")
     if config.share: