Fix exit on hook failure and exit code handling

Author: kirkbrauer

packages/jumpstarter-cli/jumpstarter_cli/run.py

@@ -76,11 +76,22 @@ async def signal_handler():
                 except* Exception as excgroup:
                     _handle_exporter_exceptions(excgroup)
 
+                # Check if exporter set an exit code (e.g., from hook failure with on_failure='exit')
+                exporter_exit_code = exporter.exit_code
+
             # Cancel the signal handler after exporter completes
             signal_tg.cancel_scope.cancel()
 
-        # Return signal number if received, otherwise 0 for immediate restart
-        return received_signal if received_signal else 0
+        # Return exit code in priority order:
+        # 1. Signal number if received (for signal-based termination)
+        # 2. Exporter's exit code if set (for hook failure with on_failure='exit')
+        # 3. 0 for immediate restart (normal exit without signal or explicit exit code)
+        if received_signal:
+            return received_signal
+        elif exporter_exit_code is not None:
+            return exporter_exit_code
+        else:
+            return 0
 
     sys.exit(anyio.run(serve_with_graceful_shutdown))
 

packages/jumpstarter/jumpstarter/exporter/exporter.py

@@ -139,6 +139,14 @@ class Exporter(AsyncContextManagerMixin, Metadata):
     determine when to trigger before-lease and after-lease hooks.
     """
 
+    _exit_code: int | None = field(init=False, default=None)
+    """Exit code to use when the exporter shuts down.
+
+    When set to a non-zero value, the exporter should terminate permanently
+    (not restart). This is used by hooks with on_failure='exit' to signal
+    that the exporter should shut down and not be restarted by the CLI.
+    """
+
     _lease_context: LeaseContext | None = field(init=False, default=None)
     """Encapsulates all resources associated with the current lease.
 
@@ -157,13 +165,17 @@ class Exporter(AsyncContextManagerMixin, Metadata):
     a reference holder and doesn't manage resource lifecycles directly.
     """
 
-    def stop(self, wait_for_lease_exit=False, should_unregister=False):
+    def stop(self, wait_for_lease_exit=False, should_unregister=False, exit_code: int | None = None):
         """Signal the exporter to stop.
 
         Args:
             wait_for_lease_exit (bool): If True, wait for the current lease to exit before stopping.
             should_unregister (bool): If True, unregister from controller. Otherwise rely on heartbeat.
+            exit_code (int | None): If set, the exporter will exit with this code (non-zero means no restart).
         """
+        # Set exit code if provided
+        if exit_code is not None:
+            self._exit_code = exit_code
 
         # Stop immediately if not started yet or if immediate stop is requested
         if (not self._started or not wait_for_lease_exit) and self._tg is not None:
@@ -178,6 +190,15 @@ def stop(self, wait_for_lease_exit=False, should_unregister=False):
             self._stop_requested = True
             logger.info("Exporter marked for stop upon lease exit")
 
+    @property
+    def exit_code(self) -> int | None:
+        """Get the exit code for the exporter.
+
+        Returns:
+            The exit code if set, or None if the exporter should restart.
+        """
+        return self._exit_code
+
     async def _get_controller_stub(self) -> jumpstarter_pb2_grpc.ControllerServiceStub:
         """Create and return a controller service stub."""
         return jumpstarter_pb2_grpc.ControllerServiceStub(await self.channel_factory())

packages/jumpstarter/jumpstarter/exporter/hooks.py

@@ -274,7 +274,7 @@ async def run_before_lease_hook(
         self,
         lease_scope: "LeaseContext",
         report_status: Callable[["ExporterStatus", str], Awaitable[None]],
-        shutdown: Callable[[], None],
+        shutdown: Callable[..., None],
     ) -> None:
         """Execute before-lease hook with full orchestration.
 
@@ -288,7 +288,7 @@ async def run_before_lease_hook(
         Args:
             lease_scope: LeaseScope containing session, socket_path, and sync event
             report_status: Async callback to report status changes to controller
-            shutdown: Callback to trigger exporter shutdown on critical failures
+            shutdown: Callback to trigger exporter shutdown (accepts optional exit_code kwarg)
         """
         try:
             # Wait for lease scope to be fully populated by handle_lease
@@ -334,7 +334,8 @@ async def run_before_lease_hook(
                     f"beforeLease hook failed (on_failure=exit, shutting down): {e}",
                 )
                 logger.error("Shutting down exporter due to beforeLease hook failure with on_failure='exit'")
-                shutdown()
+                # Exit code 1 tells the CLI not to restart the exporter
+                shutdown(exit_code=1)
             else:
                 # on_failure='endLease' - just block this lease, exporter stays available
                 logger.error("beforeLease hook failed with on_failure='endLease': %s", e)
@@ -360,7 +361,7 @@ async def run_after_lease_hook(
         self,
         lease_scope: "LeaseContext",
         report_status: Callable[["ExporterStatus", str], Awaitable[None]],
-        shutdown: Callable[[], None],
+        shutdown: Callable[..., None],
     ) -> None:
         """Execute after-lease hook with full orchestration.
 
@@ -374,7 +375,7 @@ async def run_after_lease_hook(
         Args:
             lease_scope: LeaseScope containing session, socket_path, and client info
             report_status: Async callback to report status changes to controller
-            shutdown: Callback to trigger exporter shutdown on critical failures
+            shutdown: Callback to trigger exporter shutdown (accepts optional exit_code kwarg)
         """
         try:
             # Verify lease scope is ready - for after-lease this should always be true
@@ -412,7 +413,8 @@ async def run_after_lease_hook(
                     f"afterLease hook failed (on_failure=exit, shutting down): {e}",
                 )
                 logger.error("Shutting down exporter due to afterLease hook failure with on_failure='exit'")
-                shutdown()
+                # Exit code 1 tells the CLI not to restart the exporter
+                shutdown(exit_code=1)
             else:
                 # on_failure='endLease' - lease already ended, just report the failure
                 # The exporter remains available for new leases