A running calibration can be aborted

gently/hardware/dispim/device_layer.py

@@ -3334,7 +3334,9 @@ async def handle_halt_motion(self, request):
         """
         halted: list[str] = []
         errors: dict[str, str] = {}
-        for key in ("fdrive", "xy_stage", "z_stage"):
+        # The piezo is a stage to MMCore too, and an aborted calibration can
+        # leave it mid-sweep. The galvo is not a positioner and has no stop.
+        for key in ("fdrive", "xy_stage", "z_stage", "piezo"):
             dev = self.devices.get(key)
             if dev is None:
                 continue

gently/ui/web/routes/data.py

@@ -1,5 +1,6 @@
 """Data routes - calibration, snapshots, embryos, sequence, status, events."""
 
+import asyncio
 import logging
 from datetime import datetime
 from pathlib import Path
@@ -42,6 +43,42 @@ def _json_safe(obj):
     return obj
 
 
+class CalibrationAborted(Exception):
+    """The operator pressed Abort while a calibration was running."""
+
+
+async def _run_cancellable_calibration(agent, coro, what: str):
+    """Run a calibration as a task that POST /calibrate/abort can cancel.
+
+    The routine runs here, agent-side — Claude-vision edge search, adaptive
+    focus sweeps — as a long series of SHORT device-layer plans (a snap, a
+    move, a snap). There is nothing on the device layer to abort: each plan
+    is over in a fraction of a second. What has to stop is this coroutine,
+    and cancelling its task stops it at its next await; the plan in flight
+    completes on its own and the halt that follows stops any motion.
+
+    A CancelledError that is OURS — the HTTP client went away — is re-raised
+    unchanged, so a closed tab still cancels the way it always did.
+    """
+    task = asyncio.create_task(coro)
+    agent._calibration_task = task
+    agent._calibration_what = what
+    try:
+        return await task
+    except asyncio.CancelledError:
+        # Task.cancelling() is 3.11+; the deps-less mypy run types against an
+        # older stdlib, so it is reached for by name.
+        current = asyncio.current_task()
+        cancelling = getattr(current, "cancelling", None)
+        if callable(cancelling) and cancelling():
+            raise
+        raise CalibrationAborted(what) from None
+    finally:
+        if getattr(agent, "_calibration_task", None) is task:
+            agent._calibration_task = None
+            agent._calibration_what = None
+
+
 def _parse_dic_config(raw) -> dict | None:
     """The DIC overview block of a plan, validated. None when absent or off."""
     if raw is None:
@@ -1442,11 +1479,17 @@ async def calibrate_embryo_route(embryo_id: str, payload: dict = Body(default={}
 
         registry = get_tool_registry()
         try:
-            message = await registry.execute(
-                "calibrate_embryo",
-                args,
-                {"agent": agent, "client": client},
+            message = await _run_cancellable_calibration(
+                agent,
+                registry.execute("calibrate_embryo", args, {"agent": agent, "client": client}),
+                embryo_id,
             )
+        except CalibrationAborted:
+            # 409, like the pre-flight refusal: nothing is broken, the operator
+            # chose this. The detail says which 409 it is.
+            raise HTTPException(
+                status_code=409, detail=f"Calibration of {embryo_id} aborted by operator"
+            ) from None
         except Exception as exc:
             logger.exception("Calibration failed for %s", embryo_id)
             raise HTTPException(status_code=502, detail=f"calibration failed: {exc}") from exc
@@ -1546,6 +1589,30 @@ async def restore_spim_alignment(payload: dict = Body(...)):  # noqa: B008
             raise HTTPException(status_code=404, detail="no alignment with that timestamp")
         return record.to_dict()
 
+    @router.post("/api/devices/calibrate/abort", dependencies=[Depends(require_control)])
+    async def calibrate_abort():
+        """Stop the calibration that is running, and any motion with it.
+
+        Cancels the task the calibrate route registered, then halts every
+        positioner. Answers {aborted: false} when nothing is running rather
+        than erroring: an Abort pressed a moment too late is not a fault.
+        """
+        agent = _require_agent_with_experiment()
+        task = getattr(agent, "_calibration_task", None)
+        what = getattr(agent, "_calibration_what", None)
+        if task is None or task.done():
+            return {"success": True, "aborted": False, "detail": "No calibration running"}
+        task.cancel()
+        halted = None
+        client = _resolve_client()
+        if client is not None:
+            try:
+                halted = await client.halt_motion()
+            except Exception as exc:
+                halted = {"success": False, "error": str(exc)}
+        logger.warning("Calibration of %s aborted by operator (halt: %s)", what, halted)
+        return {"success": True, "aborted": True, "what": what, "halted": halted}
+
     @router.post("/api/devices/calibrate/all", dependencies=[Depends(require_control)])
     async def calibrate_all_route(payload: dict = Body(default={})):  # noqa: B008
         """Calibrate several embryos in one go. Body: {scope, ...cal settings}.
@@ -1592,9 +1659,13 @@ async def calibrate_all_route(payload: dict = Body(default={})):  # noqa: B008
         args = {"embryo_ids": targets, **_calibration_args(payload)}
         registry = get_tool_registry()
         try:
-            message = await registry.execute(
-                "calibrate_all_embryos", args, {"agent": agent, "client": client}
+            message = await _run_cancellable_calibration(
+                agent,
+                registry.execute("calibrate_all_embryos", args, {"agent": agent, "client": client}),
+                "all",
             )
+        except CalibrationAborted:
+            raise HTTPException(status_code=409, detail="Calibration aborted by operator") from None
         except Exception as exc:
             logger.exception("Batch calibration failed")
             raise HTTPException(status_code=502, detail=f"calibration failed: {exc}") from exc

gently/ui/web/static/js/operate.js

@@ -1329,6 +1329,8 @@ const OperateManager = (function () {
             if (b) b.textContent = `Calibrating… ${Math.round((Date.now() - t0) / 1000)}s`;
         }, 1000);
         if (b) { b.disabled = true; b.textContent = 'Calibrating… 0s'; }
+        const ab = $('op-cal-abort');
+        if (ab) { ab.hidden = false; ab.disabled = false; ab.textContent = 'Abort'; }
         if (out) out.textContent = 'sweeping…';
         // The frames this run is about to take are already broadcast; the
         // progress panel shows them as they land.
@@ -1359,7 +1361,13 @@ const OperateManager = (function () {
             // 409 is the pre-flight check declining, not a crash: nothing was
             // spent past one frame and the operator can say "anyway". Anything
             // else is a failure and reads as one.
-            if (e && e.status === 409) {
+            const detail = String((e && e.data && e.data.detail) || '');
+            if (e && e.status === 409 && /aborted/i.test(detail)) {
+                // The operator's own doing: not a failure, not a refusal.
+                if (out) out.textContent = 'aborted';
+                if (typeof CalProgressPanel !== 'undefined') CalProgressPanel.finish(false, 'aborted by operator');
+                toast('Calibration aborted');
+            } else if (e && e.status === 409) {
                 if (out) out.textContent = 'nothing there';
                 _refusedFor = _selected;
                 showRefusal(e);
@@ -1376,10 +1384,24 @@ const OperateManager = (function () {
         } finally {
             clearInterval(tick);
             if (b) b.disabled = false;
+            if (ab) ab.hidden = true;
             renderCalTarget();   // restores the verb: Calibrate / Recalibrate
         }
     }
 
+    /** Stop the running calibration. The calibrate request then answers 409 "aborted". */
+    async function abortCalibration() {
+        const ab = $('op-cal-abort');
+        if (ab) { ab.disabled = true; ab.textContent = 'Aborting…'; }
+        try {
+            const d = await postJSON('/api/devices/calibrate/abort', {});
+            if (!d.aborted) toast('Nothing to abort — the calibration had already finished');
+        } catch (e) {
+            toastFail(`Abort failed (${why(e)})`);
+            if (ab) { ab.disabled = false; ab.textContent = 'Abort'; }
+        }
+    }
+
     // The calibration pane names its subject and reports the fit it has, if any
     // — the same field the server-side gate checks, so the pane shows what a run
     // would refuse rather than leaving it to be discovered at Start.
@@ -2565,6 +2587,7 @@ const OperateManager = (function () {
         // back-off button. Restored, and pinned by a test that counts them.
         const sp = $('op-spim-toggle'); if (sp) sp.addEventListener('click', toggleSpim);
         const cal = $('op-calibrate'); if (cal) cal.addEventListener('click', calibrateSelected);
+        const abort = $('op-cal-abort'); if (abort) abort.addEventListener('click', abortCalibration);
         const borrow = $('op-cal-borrow');
         if (borrow) borrow.addEventListener('click', borrowCalibration);
         const all = $('op-cal-all');

gently/ui/web/templates/index.html

@@ -824,6 +824,9 @@ <h2 class="devices-title"><span class="devices-title-script">Device</span> <em c
                            the numbers it drops cost ~70 exposures to measure. #}
                         <button class="op-btn op-btn-quiet" id="op-cal-clear" type="button" hidden>Clear fit</button>
                         <button class="op-btn op-btn-primary" id="op-calibrate" type="button">Calibrate</button>
+                        {# Shown only while a calibration runs. It cancels the
+                           routine at its next step and halts every positioner. #}
+                        <button class="op-btn op-btn-warn" id="op-cal-abort" type="button" hidden>Abort</button>
                       </div>
                       <p class="op-cap" id="op-cal-borrow-note" hidden></p>
 

tests/test_calibration_can_be_aborted.py

@@ -0,0 +1,201 @@
+"""A running calibration can be aborted.
+
+"calibration routine -- while it is running, has no abort feature. in
+devices > operate > calibration"
+
+The routine runs agent-side as a long series of short device-layer plans, so
+there is nothing on the device layer to abort: what stops is the tool's task
+in the web process, and every positioner is halted after it. Everything
+here runs against fakes.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import re
+from pathlib import Path
+from types import SimpleNamespace
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+pytest.importorskip("httpx")
+
+import httpx  # noqa: E402
+from fastapi import FastAPI  # noqa: E402
+
+# The calibrate routes import the calibration tools, whose @tool decorators
+# register against the live registry at import. Imported here, once, against
+# the REAL registry — before any test swaps in a fake — so the route's import
+# is a no-op and the fake never has to be a registry.
+import gently.app.tools.calibration_tools  # noqa: E402, F401
+import gently.ui.web.auth as auth  # noqa: E402
+from gently.ui.web.routes.data import create_router  # noqa: E402
+
+WEB = Path(__file__).resolve().parents[1] / "gently" / "ui" / "web"
+HTML = (WEB / "templates" / "index.html").read_text(encoding="utf-8")
+OPERATE = (WEB / "static" / "js" / "operate.js").read_text(encoding="utf-8")
+DEVICE_LAYER = (
+    Path(__file__).resolve().parents[1] / "gently" / "hardware" / "dispim" / "device_layer.py"
+).read_text(encoding="utf-8")
+
+
+class _SlowRegistry:
+    """A calibration that takes as long as the test lets it."""
+
+    def __init__(self):
+        self.started = asyncio.Event()
+        self.finished = False
+        self.cancelled = False
+
+    def register(self, *a, **k):  # a tool module re-registering against us
+        return None
+
+    async def execute(self, name, args, context):
+        self.started.set()
+        try:
+            await asyncio.sleep(30)
+        except asyncio.CancelledError:
+            self.cancelled = True
+            raise
+        self.finished = True
+        return "calibrated"
+
+
+def _server():
+    server = MagicMock()
+    agent = SimpleNamespace(
+        experiment=SimpleNamespace(embryos={"embryo_1": SimpleNamespace(calibration={})}),
+        client=MagicMock(),
+        session_id="s1",
+    )
+    agent.client.is_connected = True
+    agent.client.halt_motion = AsyncMock(return_value={"success": True, "halted": ["xy_stage"]})
+    server.agent_bridge.agent = agent
+    return server, agent
+
+
+async def _client(server, registry, monkeypatch):
+    import gently.harness.tools.registry as reg_mod
+
+    monkeypatch.setattr(reg_mod, "get_tool_registry", lambda: registry)
+    app = FastAPI()
+    app.include_router(create_router(server))
+    app.dependency_overrides[auth.require_control] = lambda: True
+    transport = httpx.ASGITransport(app=app)
+    return httpx.AsyncClient(transport=transport, base_url="http://t")
+
+
+def test_abort_stops_the_running_calibration_and_halts_motion(monkeypatch):
+    async def scenario():
+        server, agent = _server()
+        registry = _SlowRegistry()
+        async with await _client(server, registry, monkeypatch) as c:
+            run = asyncio.create_task(c.post("/api/devices/embryos/embryo_1/calibrate", json={}))
+            await asyncio.wait_for(registry.started.wait(), 5)
+            assert getattr(agent, "_calibration_task", None) is not None, (
+                "the run was not registered"
+            )
+
+            r = await c.post("/api/devices/calibrate/abort", json={})
+            assert r.status_code == 200, r.text
+            assert r.json()["aborted"] is True and r.json()["what"] == "embryo_1"
+            agent.client.halt_motion.assert_awaited_once()
+
+            res = await asyncio.wait_for(run, 5)
+            assert res.status_code == 409, res.text
+            assert "aborted" in res.json()["detail"]
+            assert registry.cancelled and not registry.finished
+            assert getattr(agent, "_calibration_task", None) is None, (
+                "the finished task was left registered"
+            )
+
+    asyncio.run(scenario())
+
+
+def test_abort_with_nothing_running_is_not_an_error(monkeypatch):
+    async def scenario():
+        server, agent = _server()
+        async with await _client(server, _SlowRegistry(), monkeypatch) as c:
+            r = await c.post("/api/devices/calibrate/abort", json={})
+            assert r.status_code == 200
+            assert r.json()["aborted"] is False
+            agent.client.halt_motion.assert_not_awaited()
+
+    asyncio.run(scenario())
+
+
+def test_the_batch_calibration_is_abortable_too(monkeypatch):
+    async def scenario():
+        server, agent = _server()
+        agent.experiment.embryos["embryo_1"].calibration = {}
+        agent.experiment.embryos["embryo_1"].should_skip = False
+        agent.experiment.notify_embryos_changed = lambda: None
+        registry = _SlowRegistry()
+        async with await _client(server, registry, monkeypatch) as c:
+            run = asyncio.create_task(c.post("/api/devices/calibrate/all", json={}))
+            try:
+                await asyncio.wait_for(registry.started.wait(), 5)
+            except asyncio.TimeoutError:
+                res = await run
+                pytest.skip(
+                    f"batch route never reached the tool here: {res.status_code} {res.text[:60]}"
+                )
+            r = await c.post("/api/devices/calibrate/abort", json={})
+            assert r.json()["aborted"] is True and r.json()["what"] == "all"
+            res = await asyncio.wait_for(run, 5)
+            assert res.status_code == 409 and "aborted" in res.json()["detail"]
+
+    asyncio.run(scenario())
+
+
+def test_a_client_that_goes_away_still_cancels_as_before(monkeypatch):
+    """Our own cancellation (the tab closed) must not be dressed up as an abort."""
+
+    async def scenario():
+        server, agent = _server()
+        registry = _SlowRegistry()
+        from gently.ui.web.routes.data import _run_cancellable_calibration
+
+        async def route_like():
+            return await _run_cancellable_calibration(
+                agent, registry.execute("x", {}, {}), "embryo_1"
+            )
+
+        outer = asyncio.create_task(route_like())
+        await asyncio.wait_for(registry.started.wait(), 5)
+        outer.cancel()
+        with pytest.raises(asyncio.CancelledError):
+            await outer
+        assert registry.cancelled
+
+    asyncio.run(scenario())
+
+
+# ---------------------------------------------------------------------------
+# The pane, and the halt
+# ---------------------------------------------------------------------------
+
+
+def _wire_body() -> str:
+    start = OPERATE.index("if (_wired) return;")
+    end = OPERATE.index("\n    async function ", start)
+    return OPERATE[start:end]
+
+
+def test_abort_is_on_the_pane_only_while_it_runs_and_is_wired():
+    assert 'id="op-cal-abort"' in HTML and "hidden>Abort</button>" in HTML
+    assert "$('op-cal-abort')" in _wire_body() and "abortCalibration" in _wire_body()
+    fn = OPERATE[OPERATE.index("async function calibrateSelected(") :][:4200]
+    assert "ab.hidden = false" in fn and "if (ab) ab.hidden = true" in fn
+    assert "/aborted/i.test(detail)" in fn, "an abort reads as a failure or a refusal"
+    assert fn.index("/aborted/i.test(detail)") < fn.index("showRefusal(e)"), (
+        "the abort branch must be checked before the pre-flight refusal, both are 409"
+    )
+
+
+def test_halt_now_covers_the_piezo():
+    m = re.search(r'for key in \(("fdrive", "xy_stage", "z_stage"[^)]*)\):', DEVICE_LAYER)
+    assert m and '"piezo"' in m.group(1), (
+        "an aborted sweep leaves the piezo wherever it was heading"
+    )