Fix/e stop

swarm_gpt/api/server.py

@@ -166,6 +166,18 @@ def events_after(self, job_id: str, after_id: int) -> list[dict[str, Any]]:
             job = self.get(job_id)
             return [event for event in job.events if event["id"] > after_id]
 
+    def emergency_stop_all(self) -> None:
+        """Emergency-stop every job that currently has an active deployment swarm."""
+        with self._lock:
+            jobs = list(self._jobs.values())
+        for job in jobs:
+            try:
+                job.backend.emergency_stop_active_swarm()
+            except RuntimeError:
+                continue  # No active deployment swarm for this job.
+            except Exception:
+                logger.exception("Emergency stop failed for job %s", job.id)
+
 
 def _backend_from_config(config: ApiConfig, provider: LLMProvider, model_id: str) -> AppBackend:
     return AppBackend(
@@ -342,6 +354,7 @@ def create_app(config: ApiConfig | None = None) -> FastAPI:
     config = config or ApiConfig()
     store = JobStore()
     app = FastAPI(title="SwarmGPT Browser API", lifespan=_app_lifespan)
+    app.state.store = store
 
     @app.get("/api/library")
     def library() -> dict[str, Any]:
@@ -472,6 +485,19 @@ def deploy(job_id: str) -> dict[str, Any]:
         _start_thread(job, lambda: _run_deploy_job(store, job))
         return {"jobId": job.id}
 
+    @app.post("/api/jobs/{job_id}/emergency-stop")
+    def emergency_stop(job_id: str) -> dict[str, Any]:
+        try:
+            job = store.get(job_id)
+        except KeyError:
+            raise HTTPException(status_code=404, detail="Job not found") from None
+        try:
+            job.backend.emergency_stop_active_swarm()
+        except RuntimeError as exc:
+            raise HTTPException(status_code=409, detail=str(exc)) from exc
+        store.emit(job, "emergency_stop_sent", {})
+        return {"jobId": job.id, "emergencyStopped": True}
+
     @app.post("/api/jobs/{job_id}/preset")
     def save_preset(job_id: str) -> dict[str, Any]:
         try:

swarm_gpt/core/backend.py

@@ -119,6 +119,7 @@ def __init__(
         self._preset: None | str = None
         self._strict_processing = strict_processing
         self._strict_drone_match = strict_drone_match
+        self._active_swarm: Any | None = None
         if set(self.songs) & set(self.presets):
             raise ValueError("Songs and presets must have unique names")
 
@@ -299,6 +300,7 @@ def deploy(self, drone_ids: list[int] | None = None) -> bool:
             return False
 
         swarm = DroneSwarm(self.choreographer.drones, lighthouse=self.settings["lighthouse"])
+        self._active_swarm = swarm
         logger.info("Swarm connected...")
 
         # generate references
@@ -378,14 +380,25 @@ def deploy(self, drone_ids: list[int] | None = None) -> bool:
                     )
             self.music_manager.stop()
             swarm.goto(final_pos_dict, duration=2.0)  # Transition from ideal point to hover pos
-            if self.settings["land_on_docks"]: # Commented out for demo
+            if self.settings["land_on_docks"]:  # Commented out for demo
                 swarm.goto(final_pos_dict, duration=3.0)  # Hovering
             swarm.land(duration=1.5)  # Landing
         finally:
-            swarm.close()
+            try:
+                swarm.close()
+            finally:
+                self._active_swarm = None
         logger.info("Deployment successful")
         return True
 
+    def emergency_stop_active_swarm(self) -> None:
+        """Emergency-stop the currently active deployment swarm, if one exists."""
+        swarm = self._active_swarm
+        if swarm is None:
+            raise RuntimeError("No active deployment swarm to emergency stop.")
+        swarm.emergency_stop()
+        self.music_manager.stop()
+
     def load_preset(self, preset_id: str) -> str:
         """Load a preset response.
 

swarm_gpt/core/drone_swarm.py

@@ -37,6 +37,10 @@
 _CommanderLevel = Literal["low", "high"]
 
 
+class EmergencyStopActive(RuntimeError):
+    """Raised when a motion command is issued after the emergency-stop latch is set."""
+
+
 class DroneSwarm:
     """Connects, configures, and commands a Crazyflie swarm with cflib2."""
 
@@ -75,6 +79,7 @@ def __init__(
         self._estimator_stop_event: Event | None = None
         self._estimator_future: Future[None] | None = None
         self._closed = False
+        self._estop = threading.Event()
 
         if not lighthouse:
             from drone_estimators.ros_nodes.ros2_connector import ROSConnector
@@ -125,24 +130,27 @@ def is_active(self, uri: str) -> bool:
 
     def takeoff(self, height: float = 1.5, duration: float = 3.0):
         """Take off the drones to a given height over a given duration."""
+        self._ensure_not_estopped()
 
         async def _takeoff(uri: str) -> None:
             cf = self._cf(uri)
             await self._change_commander_level(uri, "high")
             await cf.high_level_commander().take_off(height, None, duration, None)
-            await asyncio.sleep(duration)
+            await self._estop_guarded_sleep(uri, duration)
 
         self._run(self._parallel_by_uri("Taking off", self.uris, _takeoff, timeout=duration + 1.0))
 
     def land(self, height: float = 0.0, duration: float = 3.0):
         """Land the drones at a given height over a given duration."""
+        self._ensure_not_estopped()
 
         async def _land(uri: str) -> None:
             cf = self._cf(uri)
             await self._change_commander_level(uri, "high")
             high_level_commander = cf.high_level_commander()
             await high_level_commander.land(height, None, duration, None)
-            await asyncio.sleep(duration)
+            if await self._estop_guarded_sleep(uri, duration):
+                return
             await high_level_commander.stop(None)
 
         self._run(self._parallel_by_uri("Landing", self.uris, _land, timeout=duration + 1.0))
@@ -154,6 +162,7 @@ def goto(self, target: dict[str, list], duration: float = 3.0):
             target: Position+Yaw references in the form {'uri1': [target], ...}.
             duration: Duration of the motion in seconds.
         """
+        self._ensure_not_estopped()
         self._validate_required_uris("pos", target)
         for uri, setpoint in target.items():
             if len(setpoint) != 4:
@@ -165,7 +174,7 @@ async def _goto(uri: str) -> None:
             await cf.high_level_commander().go_to(
                 *target[uri], duration, relative=False, linear=True, group_mask=None
             )
-            await asyncio.sleep(duration)
+            await self._estop_guarded_sleep(uri, duration)
 
         self._run(self._parallel_by_uri("Goto", self.uris, _goto, timeout=duration + 1.0))
 
@@ -175,6 +184,7 @@ def setpoint(self, target: dict[str, list]):
         Args:
             target: Position+Yaw references in the form {'uri1': [target], ...}.
         """
+        self._ensure_not_estopped()
         self._validate_required_uris("pos", target)
         for uri, setpoint in target.items():
             if len(setpoint) != 4:
@@ -202,6 +212,7 @@ def execute_choreography(
             color_top: Top deck color cues in the form {uri: {time: wrgb}}.
             color_bot: Bottom deck color cues in the form {uri: {time: wrgb}}.
         """
+        self._ensure_not_estopped()
         self._validate_required_uris("choreography", choreography)
         if not color_top and not color_bot:
             logger.warning("No colors provided for choreography.")
@@ -262,13 +273,27 @@ async def _set_param(uri: str) -> None:
         )
 
     def emergency_stop(self, uri: str | None = None):
-        """Send an emergency stop signal to one URI or all drones (default)."""
+        """Send an emergency stop signal to one URI or all drones (default).
+
+        Sets a latch that halts the choreography stream and blocks further motion
+        commands so fresh setpoints cannot override the motor cut.
+        """
+        self._estop.set()
         if uri is None:
             uris = self.uris
         else:
             self._validate_known_uris("uri", {uri: None})
             uris = [uri]
-        self._run(self._parallel_by_uri("Emergency stop", uris, self._emergency_stop, timeout=0.5))
+        # In-flight loop coroutines self-cut on the latch; this also sends the packet directly
+        # to cover idle windows. Bounded so the caller never blocks on a loop that has stopped.
+        coro = self._parallel_by_uri("Emergency stop", uris, self._emergency_stop, timeout=0.5)
+        try:
+            if self._loop_thread is not None or self._loop.is_running():
+                asyncio.run_coroutine_threadsafe(coro, self._loop).result(timeout=1.0)
+            else:
+                self._loop.run_until_complete(coro)
+        except Exception as exc:
+            logger.error(f"Emergency stop encountered errors: {exc}")
 
     def reset(self):
         """Reset all active drones."""
@@ -345,9 +370,32 @@ async def _close() -> None:
             if self.ros_connector is not None:
                 self.ros_connector.close()
 
+    def _ensure_not_estopped(self) -> None:
+        """Block forward-motion commands once the emergency-stop latch is set."""
+        if self._estop.is_set():
+            raise EmergencyStopActive("Swarm is emergency-stopped; motion commands are blocked.")
+
+    async def _estop_guarded_sleep(self, uri: str, duration: float) -> bool:
+        """Wait up to ``duration``, cutting ``uri``'s motors from the loop if the latch trips.
+
+        Returns True if the latch tripped (motors cut), False if it slept the full duration.
+        """
+        loop = asyncio.get_running_loop()
+        deadline = loop.time() + duration
+        while (remaining := deadline - loop.time()) > 0:
+            if self._estop.is_set():
+                await self._emergency_stop(uri)
+                return True
+            await asyncio.sleep(min(0.02, remaining))
+        return False
+
     def _run(self, coroutine: Awaitable[Any]) -> Any:
-        """Run a cflib2 coroutine on the swarm event loop."""
-        if self._loop_thread is not None:
+        """Run a cflib2 coroutine on the swarm event loop.
+
+        Dispatches cross-thread when the loop runs in another thread or is already running,
+        so emergency stops from the request or signal-handler threads still reach the swarm.
+        """
+        if self._loop_thread is not None or self._loop.is_running():
             return asyncio.run_coroutine_threadsafe(coroutine, self._loop).result()
         return self._loop.run_until_complete(coroutine)
 
@@ -606,6 +654,9 @@ async def _stream_reference(
         t_col = -np.inf
 
         while (t_cur := asyncio.get_running_loop().time() - start_time) < duration:
+            if self._estop.is_set():
+                await self._emergency_stop(uri)
+                return
             await commander.send_setpoint_position(*reference(t_cur))
 
             if t_cur - t_col >= color_period:

swarm_gpt/launch.py

@@ -4,6 +4,7 @@
 import os
 import sys
 from pathlib import Path
+from types import FrameType
 
 import fire
 import uvicorn
@@ -52,8 +53,24 @@ def main(
             use_motion_primitives=use_motion_primitives,
         )
     )
+    server = uvicorn.Server(uvicorn.Config(app, host=host, port=port))
+    # uvicorn captures SIGINT/SIGTERM itself (Server.capture_signals installs self.handle_exit)
+    # and only starts a graceful shutdown, which waits on the open events WebSocket while the
+    # deploy thread keeps flying. Wrap handle_exit so the first Ctrl+C cuts the motors before
+    # uvicorn begins draining. capture_signals installs whatever self.handle_exit is at call time.
+    uvicorn_handle_exit = server.handle_exit
+
+    def handle_exit(sig: int, frame: FrameType | None) -> None:
+        logging.warning("Ctrl+C received: emergency-stopping all active swarms.")
+        try:
+            app.state.store.emergency_stop_all()
+        except Exception:
+            logging.exception("Emergency stop during shutdown failed")
+        uvicorn_handle_exit(sig, frame)
+
+    server.handle_exit = handle_exit
     try:
-        uvicorn.run(app, host=host, port=port)
+        server.run()
     finally:
         shutdown_ollama_generation()
 

tests/unit/test_api.py

@@ -1,3 +1,6 @@
+import threading
+import time
+from collections.abc import Generator
 from pathlib import Path
 from types import SimpleNamespace
 from urllib.parse import quote
@@ -6,6 +9,7 @@
 import pytest
 from fastapi.testclient import TestClient
 
+import swarm_gpt.api.server as server
 from swarm_gpt.api.server import ApiConfig, _backend_from_config, create_app, normalize_playback
 from swarm_gpt.utils.llm_providers import DEFAULT_OPENAI_MODEL_CHOICES
 
@@ -79,3 +83,93 @@ def test_library_returns_preset_display_metadata_and_delete(tmp_path: Path):
     delete_response.raise_for_status()
     assert delete_response.json() == {"deleted": preset_id}
     assert not (preset_dir / preset_id).exists()
+
+
+def test_emergency_stop_endpoint_runs_while_deploy_is_active(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    class DeployingBackend:
+        def __init__(self) -> None:
+            self.songs = ["Test Song"]
+            self.presets: list[str] = []
+            self.settings = {"axswarm": {"pos_min": [-1, -1, 0], "pos_max": [1, 1, 2]}}
+            self.music_manager = SimpleNamespace(song="Test Song")
+            self.splines: dict[int, object] = {}
+            self.deploy_entered = threading.Event()
+            self.stop_requested = threading.Event()
+            self.emergency_stop_calls = 0
+
+        def initial_prompt(self, selection: str) -> list[dict[str, str]]:
+            return []
+
+        def simulate(self) -> Generator[None, None, dict[str, object]]:
+            self.splines[0] = object()
+            states = np.zeros((1, 1, 13))
+            states[:, :, 3:7] = [0, 0, 0, 1]
+            if False:
+                yield None
+            return {"timestamps": np.array([0.0]), "states": states, "num_drones": 1}
+
+        def crop_window(self, song: str) -> tuple[float, float]:
+            return (0.0, 60.0)
+
+        def deploy(self) -> bool:
+            self.deploy_entered.set()
+            self.stop_requested.wait(timeout=2.0)
+            return True
+
+        def emergency_stop_active_swarm(self) -> None:
+            self.emergency_stop_calls += 1
+            self.stop_requested.set()
+
+    backends: list[DeployingBackend] = []
+
+    def backend_from_config(config: ApiConfig, provider: str, model_id: str) -> DeployingBackend:
+        backend = DeployingBackend()
+        backends.append(backend)
+        return backend
+
+    (tmp_path / "Test Song.mp3").write_bytes(b"")
+    monkeypatch.setattr(server, "_backend_from_config", backend_from_config)
+    client = TestClient(create_app(ApiConfig(music_dir=tmp_path)))
+
+    create_response = client.post(
+        "/api/jobs", json={"selection": "Test Song", "provider": "openai", "modelId": "gpt"}
+    )
+    create_response.raise_for_status()
+    job_id = create_response.json()["jobId"]
+    backend = backends[0]
+    for _ in range(50):
+        if client.get(f"/api/jobs/{job_id}").json()["status"] == "ready":
+            break
+        time.sleep(0.01)
+
+    deploy_response = client.post(f"/api/jobs/{job_id}/deploy")
+    deploy_response.raise_for_status()
+    assert backend.deploy_entered.wait(timeout=1.0)
+
+    stop_response = client.post(f"/api/jobs/{job_id}/emergency-stop")
+    stop_response.raise_for_status()
+
+    assert stop_response.json() == {"jobId": job_id, "emergencyStopped": True}
+    assert backend.emergency_stop_calls == 1
+
+
+def test_emergency_stop_all_stops_active_swarms_and_skips_idle_jobs() -> None:
+    calls: list[str] = []
+
+    class ActiveBackend:
+        def emergency_stop_active_swarm(self) -> None:
+            calls.append("active")
+
+    class IdleBackend:
+        def emergency_stop_active_swarm(self) -> None:
+            raise RuntimeError("No active deployment swarm to emergency stop.")
+
+    store = server.JobStore()
+    store.create(ActiveBackend())  # type: ignore[arg-type]
+    store.create(IdleBackend())  # type: ignore[arg-type]
+
+    store.emergency_stop_all()
+
+    assert calls == ["active"]