From fe3c6de441b9f2bb5fe294194503d0cbf421e6fe Mon Sep 17 00:00:00 2001
From: yrasool <111371895+yrasool@users.noreply.github.com>
Date: Mon, 22 Jun 2026 15:37:11 -0400
Subject: [PATCH] Fix TensorBoard observer shutdown

---
 .../tensorboard/tensorboard_interceptor.py    |  7 ++---
 tests/adapters/test_tensorboard_lifecycle.py  | 31 +++++++++++++++++++
 2 files changed, 34 insertions(+), 4 deletions(-)
 create mode 100644 tests/adapters/test_tensorboard_lifecycle.py

diff --git a/src/flowcept/flowceptor/adapters/tensorboard/tensorboard_interceptor.py b/src/flowcept/flowceptor/adapters/tensorboard/tensorboard_interceptor.py
index f0acf721..124625f6 100644
--- a/src/flowcept/flowceptor/adapters/tensorboard/tensorboard_interceptor.py
+++ b/src/flowcept/flowceptor/adapters/tensorboard/tensorboard_interceptor.py
@@ -106,11 +106,10 @@ def stop(self, check_safe_stops: bool = True) -> bool:
         self.logger.debug("Interceptor stopping...")
         if self._observer is not None:
             self._observer.stop()
+            self._observer.join()
+            self._observer = None
         try:
-            intercepted = self.callback()
-            if intercepted == 0:
-                sleep(self.settings.watch_interval_sec)
-                self.callback()
+            self.callback()
         except Exception as e:
             self.logger.exception(e)
         super().stop(check_safe_stops)
diff --git a/tests/adapters/test_tensorboard_lifecycle.py b/tests/adapters/test_tensorboard_lifecycle.py
new file mode 100644
index 00000000..a36397d3
--- /dev/null
+++ b/tests/adapters/test_tensorboard_lifecycle.py
@@ -0,0 +1,31 @@
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import pytest
+
+tensorboard_module = pytest.importorskip("flowcept.flowceptor.adapters.tensorboard.tensorboard_interceptor")
+TensorboardInterceptor = tensorboard_module.TensorboardInterceptor
+
+
+def _noop(*_args, **_kwargs):
+    pass
+
+
+def test_stop_joins_observer_before_final_callback():
+    events = []
+    interceptor = TensorboardInterceptor.__new__(TensorboardInterceptor)
+    interceptor._observer = SimpleNamespace(
+        stop=lambda: events.append("stop"),
+        join=lambda: events.append("join"),
+    )
+    interceptor.callback = lambda: events.append("callback")
+    interceptor.logger = SimpleNamespace(debug=_noop, exception=_noop)
+    interceptor._interceptor_instance_id = "test-interceptor"
+    interceptor._bundle_exec_id = "test-bundle"
+    interceptor._mq_dao = SimpleNamespace(stop=lambda **_kwargs: events.append("mq_stop"))
+
+    with patch.object(tensorboard_module, "sleep", _noop):
+        assert TensorboardInterceptor.stop(interceptor, check_safe_stops=False)
+
+    assert events == ["stop", "join", "callback", "mq_stop"]
+    assert interceptor._observer is None