fix: thread safety, graceful shutdown, health endpoint, and dynamic obs fallback

SemSorter/server/agent_bridge.py

@@ -12,9 +12,12 @@ import asyncio
 import logging
 import os
 import sys
+import threading
 from pathlib import Path
 from typing import Any, Callable, Dict, List, Optional
 
+from dotenv import load_dotenv
+
 logger = logging.getLogger(__name__)
 
 # ── Path setup ────────────────────────────────────────────────────────────────
@@ -30,6 +33,8 @@ for _plugin in ("gemini", "deepgram", "elevenlabs", "getstream"):
     if _plugin_path.exists():
         sys.path.insert(0, str(_plugin_path))
 
+load_dotenv(_PROJECT_ROOT / ".env")
+
 # ── Quota-tracking state ──────────────────────────────────────────────────────
 _quota_exceeded: Dict[str, bool] = {
     "gemini": False,
@@ -55,6 +60,7 @@ _bridge = None
 _llm = None
 _tts = None
 _notify_cb: Optional[Callable[[Dict], None]] = None  # Push events to WebSocket
+_sim_lock = threading.RLock()
 
 
 def set_notify_callback(cb: Callable[[Dict], None]) -> None:
@@ -73,9 +79,10 @@ def _push(event: Dict) -> None:
 
 
 def _check_quota_error(exc: Exception) -> Optional[str]:
-    """Return service name if the exception indicates API quota exhaustion."""
+    """Return service name if the exception indicates quota/auth API failures."""
     msg = str(exc).lower()
-    if "resource_exhausted" in msg or "429" in msg or "quota" in msg:
+    if ("resource_exhausted" in msg or "429" in msg or "quota" in msg
+            or "invalid api key" in msg or "unauthorized" in msg or "401" in msg):
         if "gemini" in msg or "google" in msg:
             return "gemini"
         if "deepgram" in msg:
@@ -105,23 +112,25 @@ def _mark_quota_exceeded(service: str) -> None:
 
 def get_simulation():
     global _sim
-    if _sim is None:
-        os.environ.setdefault("MUJOCO_GL", "egl")
-        from controller import SemSorterSimulation
-        logger.info("Initialising MuJoCo simulation…")
-        _sim = SemSorterSimulation()
-        _sim.load_scene()
-        _sim.step(300)
-        logger.info("Simulation ready: %d items", len(_sim.items))
+    with _sim_lock:
+        if _sim is None:
+            os.environ.setdefault("MUJOCO_GL", "egl")
+            from controller import SemSorterSimulation
+            logger.info("Initialising MuJoCo simulation…")
+            _sim = SemSorterSimulation()
+            _sim.load_scene()
+            _sim.step(300)
+            logger.info("Simulation ready: %d items", len(_sim.items))
     return _sim
 
 
 def get_bridge():
     global _bridge
-    if _bridge is None:
-        from vlm_bridge import VLMSimBridge
-        _bridge = VLMSimBridge(simulation=get_simulation(), use_direct=True)
-        logger.info("VLM bridge ready")
+    with _sim_lock:
+        if _bridge is None:
+            from vlm_bridge import VLMSimBridge
+            _bridge = VLMSimBridge(simulation=get_simulation(), use_direct=True)
+            logger.info("VLM bridge ready")
     return _bridge
 
 
@@ -171,9 +180,7 @@ async def _scan_hazards_impl() -> Dict[str, Any]:
     try:
         bridge = get_bridge()
         loop = asyncio.get_event_loop()
-        detections = await loop.run_in_executor(
-            None, bridge.processor.detect_hazards)
-        matched = bridge.match_detections_to_items(detections)
+        detections, matched = await loop.run_in_executor(None, _detect_and_match_impl)
         return _format_scan(matched, demo=False)
     except Exception as exc:
         svc = _check_quota_error(exc)
@@ -216,14 +223,15 @@ async def _pick_place_impl(item_name: str, bin_type: str) -> Dict[str, Any]:
         return {"success": False, "error": f"{item_name} already sorted"}
 
     loop = asyncio.get_event_loop()
-    success = await loop.run_in_executor(None, sim.pick_and_place, item_name, target)
+    success = await loop.run_in_executor(None, _pick_place_sync, sim, item_name, target)
     return {"success": success, "item": item_name, "bin": bin_type,
             "total_sorted": sim._items_sorted}
 
 
 def _state_impl() -> Dict[str, Any]:
-    sim = get_simulation()
-    state = sim.get_state()
+    with _sim_lock:
+        sim = get_simulation()
+        state = sim.get_state()
     return {
         "time": round(state.time, 2),
         "arm_busy": state.arm_busy,
@@ -263,6 +271,31 @@ async def _sort_all_impl() -> Dict[str, Any]:
             "items_sorted": sorted_count, "details": details, "demo_mode": demo}
 
 
+def render_frame(camera: str = "overview"):
+    """Thread-safe simulation frame render for the video WS endpoint."""
+    with _sim_lock:
+        sim = get_simulation()
+        return sim.render_frame(camera=camera)
+
+
+def close_resources() -> None:
+    """Best-effort shutdown for long-running server process."""
+    global _bridge, _sim
+    with _sim_lock:
+        if _bridge is not None:
+            try:
+                _bridge.close()
+            except Exception:
+                pass
+            _bridge = None
+        if _sim is not None and hasattr(_sim, "close"):
+            try:
+                _sim.close()
+            except Exception:
+                pass
+            _sim = None
+
+
 # ── Text → agent response ─────────────────────────────────────────────────────
 
 async def process_text_command(text: str) -> str:
@@ -288,6 +321,20 @@ async def process_text_command(text: str) -> str:
         return f"Error processing command: {exc}"
 
 
+def _detect_and_match_impl():
+    """Run detect+match atomically to avoid simulation/render race conditions."""
+    with _sim_lock:
+        bridge = get_bridge()
+        detections = bridge.processor.detect_hazards()
+        matched = bridge.match_detections_to_items(detections)
+        return detections, matched
+
+
+def _pick_place_sync(sim, item_name: str, target) -> bool:
+    with _sim_lock:
+        return sim.pick_and_place(item_name, target)
+
+
 async def _llm_demo_response(text: str) -> str:
     """Return a plausible demo response when Gemini quota is exhausted."""
     t = text.lower()

SemSorter/server/app.py

@@ -24,12 +24,11 @@ import json
 import logging
 import os
 from pathlib import Path
-from typing import Set
+from typing import Optional, Set
 
 from fastapi import FastAPI, WebSocket, WebSocketDisconnect, UploadFile, File
 from fastapi.responses import HTMLResponse, JSONResponse
 from fastapi.staticfiles import StaticFiles
-import numpy as np
 from PIL import Image
 
 # ── Local imports ─────────────────────────────────────────────────────────────
@@ -48,10 +47,12 @@ _STATIC.mkdir(exist_ok=True)
 # ── Connected WebSocket clients ───────────────────────────────────────────────
 _chat_clients: Set[WebSocket] = set()
 _video_clients: Set[WebSocket] = set()
+_main_loop: Optional[asyncio.AbstractEventLoop] = None
 
 
 async def _broadcast_chat(event: dict) -> None:
     """Push a JSON event to all connected chat WebSocket clients."""
+    global _chat_clients
     payload = json.dumps(event)
     dead = set()
     for ws in list(_chat_clients):
@@ -59,17 +60,20 @@ async def _broadcast_chat(event: dict) -> None:
             await ws.send_text(payload)
         except Exception:
             dead.add(ws)
-    _chat_clients -= dead
+    for ws in dead:
+        _chat_clients.discard(ws)
 
 
 def _sync_broadcast(event: dict) -> None:
     """Thread-safe push called from sync code (bridge callbacks)."""
+    if _main_loop is None:
+        return
     try:
-        loop = asyncio.get_event_loop()
-        if loop.is_running():
-            asyncio.create_task(_broadcast_chat(event))
+        _main_loop.call_soon_threadsafe(
+            asyncio.create_task, _broadcast_chat(event)
+        )
     except Exception:
-        pass
+        logger.exception("Failed to schedule chat broadcast")
 
 
 # Register the broadcast callback so agent_bridge can push quota warnings
@@ -79,12 +83,21 @@ bridge.set_notify_callback(_sync_broadcast)
 # ── Startup: pre-warm simulation ──────────────────────────────────────────────
 @app.on_event("startup")
 async def startup():
+    global _main_loop
+    _main_loop = asyncio.get_running_loop()
     logger.info("Pre-warming MuJoCo simulation…")
-    loop = asyncio.get_event_loop()
-    await loop.run_in_executor(None, bridge.get_simulation)
+    await _main_loop.run_in_executor(None, bridge.get_simulation)
     logger.info("Simulation ready")
 
 
+@app.on_event("shutdown")
+async def shutdown():
+    logger.info("Shutting down SemSorter resources…")
+    loop = asyncio.get_running_loop()
+    await loop.run_in_executor(None, bridge.close_resources)
+    logger.info("Shutdown complete")
+
+
 # ── REST endpoints ────────────────────────────────────────────────────────────
 
 @app.get("/", response_class=HTMLResponse)
@@ -100,6 +113,11 @@ async def api_state():
     return JSONResponse(state)
 
 
+@app.get("/health")
+async def health():
+    return JSONResponse({"ok": True})
+
+
 @app.post("/api/sort")
 async def api_sort():
     """Trigger the full detect-match-sort pipeline."""
@@ -178,8 +196,7 @@ async def ws_chat(ws: WebSocket):
 
 def _render_frame_jpeg(quality: int = 75) -> bytes:
     """Render a MuJoCo frame and encode as JPEG bytes."""
-    sim = bridge.get_simulation()
-    frame = sim.render_frame(camera="overview")         # numpy H×W×3
+    frame = bridge.render_frame(camera="overview")         # numpy H×W×3
     img = Image.fromarray(frame)
     buf = io.BytesIO()
     img.save(buf, format="JPEG", quality=quality)

SemSorter/vision/vision_pipeline.py

@@ -14,13 +14,11 @@ Usage:
 
 import os
 import sys
-import cv2
 import json
-import time
 import logging
 import google.generativeai as genai
 from PIL import Image
-from typing import List, Dict, Optional
+from typing import List, Dict
 
 logger = logging.getLogger(__name__)
 
@@ -98,6 +96,14 @@ class HazardDetectionProcessor:
 
     def _capture_from_obs(self) -> Image.Image:
         """Capture a frame from the OBS Virtual Camera."""
+        try:
+            import cv2
+        except ImportError as exc:
+            raise RuntimeError(
+                "OpenCV is required for OBS capture mode. "
+                "Install opencv-python or opencv-python-headless."
+            ) from exc
+
         if self._video_cap is None or not self._video_cap.isOpened():
             self._video_cap = cv2.VideoCapture(self.device_id)
             if not self._video_cap.isOpened():

render.yaml

@@ -17,5 +17,5 @@ services:
         sync: false
       - key: STREAM_API_SECRET
         sync: false
-    healthCheckPath: /api/state
+    healthCheckPath: /health
     autoDeploy: true