Monkey patch Uvicorn's handle_exit function to stop streams

Before the uvicorn handle_exit function is called kill the camera
streams.

src/openflexure_microscope_server/server/__init__.py

@@ -1,15 +1,42 @@
 from __future__ import annotations
 
-from typing import Optional
+from typing import Optional, Callable
 from copy import copy
 
 import labthings_fastapi as lt
 import uvicorn
+import logging
+from uvicorn.main import Server
 from .serve_static_files import add_static_files
 from .legacy_api import add_v2_endpoints
 from ..logging import configure_logging, retrieve_log, retrieve_log_from_file
 
 
+def set_shutdown_function(shutdown_function: Callable[[], None]):
+    """
+    Ensure a function is called before the shutdown
+
+    This monkey patches the Uvicorn Server's handle_exit. This is needed because
+    the uvicorn `lifecycle` events and FastAPI `shutdown` events only fire once
+    background tasks have completed.
+
+    Without this the system exits cleanly only if no client is receiving a
+    StreamingResponse. This patch is used to stop the async generators that
+    send streaming responses.
+
+    :param shutdown_function: A callable with no arguments or outputs. This
+    should stop any async generators that may be sending to streaming responses.
+    """
+
+    original_handler = Server.handle_exit
+
+    def handle_exit(*args, **kwargs):
+        shutdown_function()
+        original_handler(*args, **kwargs)
+
+    Server.handle_exit = handle_exit
+
+
 def customise_server(
     server: lt.ThingServer, log_folder: str, scans_folder: Optional[str]
 ):
@@ -57,6 +84,21 @@ def serve_from_cli(argv: Optional[list[str]] = None):
         scans_folder = _get_scans_dir(config)
         server = lt.cli.server_from_config(config)
         customise_server(server, log_folder, scans_folder)
+
+        def shutdown_call():
+            try:
+                # Kill any mjpeg streams so that StreamingResponses close.
+                server.things["/camera/"].kill_mjpeg_streams()
+            except BaseException as e:
+                # Catch anything and log as it is essential that this
+                # function cannot raise an unhandled exception or Uvicorn
+                # will never get a shutdown signal.
+                logging.error(e, exc_info=True)
+
+        # Monkey patch uvicorn's exit handling to stop the MJPEG Streams
+        # before waiting for background tasks to complete.
+        set_shutdown_function(shutdown_call)
+
         uvicorn.run(
             server.app,
             host=args.host,