diff --git a/notebooks/test/test-logging.ipynb b/notebooks/test/test-logging.ipynb
index 606e6300c5dec36f66f0d88ed3e8be370a28c35b..0d033847c8069d2d99134e3d8fe4c77e7d786a46 100644
--- a/notebooks/test/test-logging.ipynb
+++ b/notebooks/test/test-logging.ipynb
@@ -36,6 +36,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "# This wont be logged\n",
     "logging.info(\"Logging some (INFO)rmation\")"
    ]
   },
@@ -55,9 +56,9 @@
    "outputs": [],
    "source": [
     "import warnings\n",
-    "from cal_tools.warnings import CalWarning\n",
+    "from cal_tools.warnings import CalibrationWarning\n",
     "\n",
-    "class TestCalWarning(CalWarning):\n",
+    "class TestCalWarning(CalibrationWarning):\n",
     "    \"\"\"Base class for custom user warnings\"\"\"\n",
     "    pass\n",
     "\n",
@@ -77,6 +78,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "# This wont be logged\n",
     "logging.warning(\"This is a warning message using logging standard library.\")"
    ]
   },
@@ -95,6 +97,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "# This wont be logged\n",
     "logging.error(\"Logging some (ERROR) without failing the notebook\")"
    ]
   },
@@ -105,9 +108,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from cal_tools.exceptions import CalError\n",
+    "from cal_tools.exceptions import CalibrationError\n",
     "\n",
-    "raise CalError('Calibration Failure')"
+    "raise CalibrationError('Calibration Failure')"
    ]
   }
  ],
diff --git a/src/xfel_calibrate/setup_logging.py b/src/xfel_calibrate/setup_logging.py
index 24d1c7f775a2fa1e21121673d1ffc562dc66c55c..3c6b065ef8cc5d21161ba7dc9dd9169efdba6361 100644
--- a/src/xfel_calibrate/setup_logging.py
+++ b/src/xfel_calibrate/setup_logging.py
@@ -12,9 +12,22 @@ JOB_ID = os.getenv('SLURM_JOB_ID', 'local')
 
 class ContextFilter(logging.Filter):
     def filter(self, record):
-        record.notebook = NOTEBOOK_NAME
-        record.job_id = JOB_ID
-        return True
+        # Only allow records that come from exception handlers
+        if getattr(record, 'from_exception_handler', False):
+            record.notebook = NOTEBOOK_NAME
+            record.job_id = JOB_ID
+            return True
+        return False
+
+
+def get_class_hierarchy(cls):
+    """Helper function to get the full class hierarchy"""
+    class_hierarchy = []
+    current_class = cls
+    while current_class and current_class != object:
+        class_hierarchy.append(current_class.__name__)
+        current_class = current_class.__base__
+    return '.'.join(reversed(class_hierarchy))
 
 
 class CustomJsonFormatter(jsonlogger.JsonFormatter):
@@ -25,23 +38,15 @@ class CustomJsonFormatter(jsonlogger.JsonFormatter):
         log_record['level'] = record.levelname
         log_record['filename'] = record.filename
         log_record['lineno'] = record.lineno
-        log_record['class'] = getattr(record, 'class', 'DefaultClass')
-
-        exc_info = getattr(record, 'exc_info', None)
-        if exc_info and exc_info[0]:
-            exc_class = exc_info[0]
-            class_hierarchy = []
-            current_class = exc_class
-            while current_class and current_class != Exception:
-                class_hierarchy.append(current_class.__name__)
-                current_class = current_class.__base__
-            class_hierarchy.append('Exception')
-            log_record['class'] = '.'.join(reversed(class_hierarchy))
-        else:
-            log_record['class'] = getattr(record, 'class', 'DefaultClass')
+
+        # Get log_class from extra parameters (set in our warning/error handlers)
+        if hasattr(record, 'log_class'):
+            log_record['log_class'] = record.log_class
 
         if record.exc_info:
             log_record['exc_info'] = self.formatException(record.exc_info)
+            exc_class = record.exc_info[0]
+            log_record['class'] = get_class_hierarchy(exc_class)
 
 
 # Create a logger
@@ -51,7 +56,7 @@ logger.setLevel(logging.INFO)
 # Define a custom JSON format
 formatter = CustomJsonFormatter(
     '%(timestamp)s %(level)s %(filename)s %(lineno)d '
-    '%(notebook)s %(job_id)s %(class)s %(message)s')
+    '%(notebook)s %(job_id)s %(log_class)s %(message)s')
 
 
 # Function to create a file handler with job-specific JSON log file
@@ -91,6 +96,12 @@ handling_error = False
 
 def safe_handle_error(exc_type, exc_value, exc_traceback):
     global handling_error
+
+    # Added this block to skip sys.exit()
+    if exc_type in (SystemExit, KeyboardInterrupt):
+        sys.__excepthook__(exc_type, exc_value, exc_traceback)
+        return
+
     if handling_error:  # Avoid infinite loop of errors.
         sys.stderr.write("Recursive error detected!\n")
         traceback.print_exception(
@@ -103,7 +114,8 @@ def safe_handle_error(exc_type, exc_value, exc_traceback):
             extra={
                 'notebook': NOTEBOOK_NAME,
                 'job_id': JOB_ID,
-                'class': exc_type.__name__ if exc_type else 'DefaultErrorClass',  # noqa
+                'log_class': get_class_hierarchy(exc_type),
+                'from_exception_handler': True
             },
             exc_info=(exc_type, exc_value, exc_traceback)
         )
@@ -123,7 +135,8 @@ def handle_warning(message, category, filename, lineno, file=None, line=None):
             extra={
                 'notebook': NOTEBOOK_NAME,
                 'job_id': JOB_ID,
-                'class': category.__name__ if category else 'DefaultWarningClass',  # noqa
+                'log_class': get_class_hierarchy(category),
+                'from_exception_handler': True
             }
         )
     except Exception as log_error: