787 Update sample status field correctly

src/mx_bluesky/common/external_interaction/callbacks/common/plan_reactive_callback.py

@@ -36,7 +36,7 @@ class will be activated, and on recieving the stop document for the
         super().__init__(emit=emit)
         self.emit_cb = emit  # to avoid GC; base class only holds a WeakRef
         self.active = False
-        self.activity_uid = 0
+        self.activity_uid = ""
         self.log = log
 
     def _run_activity_gated(self, name: str, func, doc, override=False):
@@ -76,7 +76,7 @@ def stop(self, doc: RunStop) -> RunStop | None:
         do_stop = self.active
         if doc.get("run_start") == self.activity_uid:
             self.active = False
-            self.activity_uid = 0
+            self.activity_uid = ""
         return (
             self._run_activity_gated(
                 "stop", self.activity_gated_stop, doc, override=True

src/mx_bluesky/hyperion/external_interaction/callbacks/sample_handling/sample_handling_callback.py

@@ -19,22 +19,27 @@ def __init__(self):
         super().__init__(log=ISPYB_ZOCALO_CALLBACK_LOGGER)
         self._sample_id: int | None = None
         self._descriptor: str | None = None
+        self._run_id: str | None = None
 
     def activity_gated_start(self, doc: RunStart):
-        if not self._sample_id:
+        if not self._sample_id and self.active:
             sample_id = doc.get("metadata", {}).get("sample_id")
             self.log.info(f"Recording sample ID at run start {sample_id}")
             self._sample_id = sample_id
+            self._run_id = self.activity_uid
 
     def activity_gated_stop(self, doc: RunStop) -> RunStop:
-        if doc["exit_status"] != "success":
-            exception_type, message = SampleException.type_and_message_from_reason(
-                doc.get("reason", "")
-            )
-            self.log.info(
-                f"Sample handling callback intercepted exception of type {exception_type}: {message}"
-            )
-            self._record_exception(exception_type)
+        if self._run_id == doc.get("run_start"):
+            if doc["exit_status"] != "success":
+                exception_type, message = SampleException.type_and_message_from_reason(
+                    doc.get("reason", "")
+                )
+                self.log.info(
+                    f"Sample handling callback intercepted exception of type {exception_type}: {message}"
+                )
+                self._record_exception(exception_type)
+            self._sample_id = None
+            self._run_id = None
         return doc
 
     def _record_exception(self, exception_type: str):

tests/unit_tests/hyperion/external_interaction/callbacks/sample_handling/test_sample_handling_callback.py

@@ -1,5 +1,6 @@
 from unittest.mock import MagicMock, patch
 
+import bluesky.preprocessors as bpp
 import pytest
 from bluesky.preprocessors import run_decorator
 from bluesky.run_engine import RunEngine
@@ -27,6 +28,69 @@ def plan_with_general_exception(exception_type: type, msg: str):
     raise exception_type(msg)
 
 
+def plan_for_sample_id(sample_id):
+    def plan_with_exception():
+        yield from []
+        raise SampleException(f"Test exception for sample_id {sample_id}")
+
+    yield from bpp.run_wrapper(
+        plan_with_exception(),
+        md={
+            "metadata": {"sample_id": sample_id},
+            "activate_callbacks": ["SampleHandlingCallback"],
+        },
+    )
+
+
+def plan_with_exception_from_inner_plan():
+    @run_decorator(
+        md={
+            "metadata": {"sample_id": TEST_SAMPLE_ID},
+        }
+    )
+    def inner_plan():
+        yield from []
+        raise SampleException("Exception from inner plan")
+
+    @run_decorator(
+        md={
+            "metadata": {"sample_id": TEST_SAMPLE_ID},
+            "activate_callbacks": ["SampleHandlingCallback"],
+        }
+    )
+    @bpp.set_run_key_decorator("outer_plan")
+    def outer_plan():
+        yield from inner_plan()
+
+    yield from outer_plan()
+
+
+def plan_with_rethrown_exception():
+    @run_decorator(
+        md={
+            "metadata": {"sample_id": TEST_SAMPLE_ID},
+        }
+    )
+    def inner_plan():
+        yield from []
+        raise AssertionError("Exception from inner plan")
+
+    @run_decorator(
+        md={
+            "metadata": {"sample_id": TEST_SAMPLE_ID},
+            "activate_callbacks": ["SampleHandlingCallback"],
+        }
+    )
+    @bpp.set_run_key_decorator("outer_plan")
+    def outer_plan():
+        try:
+            yield from inner_plan()
+        except AssertionError as e:
+            raise SampleException("Exception from outer plan") from e
+
+    yield from outer_plan()
+
+
 @run_decorator(
     md={
         "metadata": {"sample_id": TEST_SAMPLE_ID},
@@ -80,3 +144,64 @@ def test_sample_handling_callback_closes_run_normally(RE: RunEngine):
         RE(plan_with_normal_completion())
 
     record_exception.assert_not_called()
+
+
+@patch(
+    "mx_bluesky.hyperion.external_interaction.callbacks.sample_handling.sample_handling_callback"
+    ".ExpeyeInteraction",
+)
+def test_sample_handling_callback_resets_sample_id(
+    mock_expeye_cls: MagicMock, RE: RunEngine
+):
+    mock_expeye = mock_expeye_cls.return_value
+    callback = SampleHandlingCallback()
+    RE.subscribe(callback)
+
+    with pytest.raises(SampleException):
+        RE(plan_for_sample_id(TEST_SAMPLE_ID))
+    mock_expeye.update_sample_status.assert_called_once_with(
+        TEST_SAMPLE_ID, BLSampleStatus.ERROR_SAMPLE
+    )
+    mock_expeye.reset_mock()
+
+    with pytest.raises(SampleException):
+        RE(plan_for_sample_id(TEST_SAMPLE_ID + 1))
+    mock_expeye.update_sample_status.assert_called_once_with(
+        TEST_SAMPLE_ID + 1, BLSampleStatus.ERROR_SAMPLE
+    )
+
+
+@patch(
+    "mx_bluesky.hyperion.external_interaction.callbacks.sample_handling.sample_handling_callback"
+    ".ExpeyeInteraction",
+)
+def test_sample_handling_callback_triggered_only_by_outermost_plan_when_exception_thrown_in_inner_plan(
+    mock_expeye_cls: MagicMock, RE: RunEngine
+):
+    mock_expeye = mock_expeye_cls.return_value
+    callback = SampleHandlingCallback()
+    RE.subscribe(callback)
+
+    with pytest.raises(SampleException):
+        RE(plan_with_exception_from_inner_plan())
+    mock_expeye.update_sample_status.assert_called_once_with(
+        TEST_SAMPLE_ID, BLSampleStatus.ERROR_SAMPLE
+    )
+
+
+@patch(
+    "mx_bluesky.hyperion.external_interaction.callbacks.sample_handling.sample_handling_callback"
+    ".ExpeyeInteraction",
+)
+def test_sample_handling_callback_triggered_only_by_outermost_plan_when_exception_rethrown_from_outermost_plan(
+    mock_expeye_cls: MagicMock, RE: RunEngine
+):
+    mock_expeye = mock_expeye_cls.return_value
+    callback = SampleHandlingCallback()
+    RE.subscribe(callback)
+
+    with pytest.raises(SampleException):
+        RE(plan_with_rethrown_exception())
+    mock_expeye.update_sample_status.assert_called_once_with(
+        TEST_SAMPLE_ID, BLSampleStatus.ERROR_SAMPLE
+    )