From 5b0b936b6a866cc00cf28a8e5cf1bce554f98a20 Mon Sep 17 00:00:00 2001
From: Andres Morfin Veytia <78442543+amorfinv@users.noreply.github.com>
Date: Tue, 25 Aug 2026 17:53:22 +0200
Subject: [PATCH 1/4] Consolidate proxy teardown state-clearing into one shared
 helper

The three teardown paths (close, stop_client, _handle_disconnection) each
maintained their own copy of "forget all cached server state" and had
drifted apart. Extract DataManager._reset_cached_state as the single
implementation and route all paths through it, so browsers always receive
the cleared node/traffic picture on disconnect.

Also remove the dead reconnect() and _ensure_clean_zmq_context() code
paths (the app reconnects by building a fresh BlueSkyProxy), and add unit
tests covering the teardown paths.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 WebATM/proxy/core.py                        |   8 --
 WebATM/proxy/managers/connection_manager.py | 135 +++-----------------
 WebATM/proxy/managers/data_manager.py       |  34 ++---
 tests/test_connection_manager.py            | 133 +++++++++++++++++++
 4 files changed, 171 insertions(+), 139 deletions(-)
 create mode 100644 tests/test_connection_manager.py

diff --git a/WebATM/proxy/core.py b/WebATM/proxy/core.py
index c42f0b6..546e772 100644
--- a/WebATM/proxy/core.py
+++ b/WebATM/proxy/core.py
@@ -131,10 +131,6 @@ def is_connected(self) -> bool:
         """
         return self.was_connected and self.running and len(self.tracked_nodes) > 0
 
-    def _ensure_clean_zmq_context(self):
-        """Ensure we have a clean environment for ZMQ connections."""
-        return self.connection_mgr._ensure_clean_zmq_context()
-
     def _connect_bluesky_client_signals(self):
         """Connect BlueSky client signals to our handlers."""
         return self.connection_mgr._connect_bluesky_client_signals()
@@ -176,10 +172,6 @@ def _close_bluesky_client(self):
         """Close network client following ZMQ pattern: close sockets first, then context."""
         return self.connection_mgr._close_bluesky_client()
 
-    def reconnect(self, hostname=None):
-        """Reconnect to BlueSky server following ZMQ pattern."""
-        return self.connection_mgr.reconnect(hostname)
-
     def close(self):
         """Close all network connections and clear state like BlueSky's close() method."""
         return self.connection_mgr.close()
diff --git a/WebATM/proxy/managers/connection_manager.py b/WebATM/proxy/managers/connection_manager.py
index 14e2da4..fd17ce3 100644
--- a/WebATM/proxy/managers/connection_manager.py
+++ b/WebATM/proxy/managers/connection_manager.py
@@ -1,6 +1,5 @@
 """Connection management for the BlueSky proxy."""
 
-import gc
 import threading
 import time
 
@@ -26,23 +25,6 @@ def __init__(self, proxy):
         """
         self.proxy = proxy
 
-    def _ensure_clean_zmq_context(self):
-        """Ensure we have a clean environment for ZMQ connections."""
-        try:
-            logger.debug("Preparing clean environment for ZMQ connection...")
-
-            # Force garbage collection to clean up any lingering references
-            gc.collect()
-
-            # Small delay to ensure cleanup completes
-            time.sleep(0.2)
-
-            logger.info("Environment cleanup complete - ready for new FixedClient")
-
-        except Exception as e:
-            logger.warning(f" Warning during cleanup: {e}")
-            # Continue anyway - the FixedClient might still work
-
     def _connect_bluesky_client_signals(self):
         """Connect BlueSky client signals to our node-manager handlers."""
         client = self.proxy.bluesky_client
@@ -276,26 +258,12 @@ def _handle_disconnection(self, reason="Unknown"):
         # Don't try to reconnect - just stop running and close
         self.proxy.running = False
         self.proxy.allow_reconnection = False
-
-        # Reset connection failure counter
         self.proxy.connection_failures = 0
 
-        # Clear active node reference immediately to prevent showing corrupted data
-        if self.proxy.bluesky_client and hasattr(self.proxy.bluesky_client, "act_id"):
-            self.proxy.bluesky_client.act_id = None
-
-        # Clear all tracked nodes and servers immediately
-        self.proxy.tracked_nodes.clear()
-        self.proxy.tracked_servers.clear()
+        # Clear the cached state BEFORE emitting below, so browsers receive
+        # the disconnected (empty) node/traffic picture, not the stale one.
+        self.proxy.data_mgr._reset_cached_state()
 
-        # Clear all cached data
-        self.proxy.traffic_data = {}
-        self.proxy.sim_data = {}
-        self.proxy.echo_data = {}
-        self.proxy.poly_data_by_node.clear()
-        self.proxy.polyline_data_by_node.clear()
-
-        # Clear all screen data and emit updates to show disconnected state
         if self.proxy.socketio and self.proxy.connected_clients > 0:
             try:
                 # Emit cleared data to remove all aircraft and simulation info from screen
@@ -307,66 +275,33 @@ def _handle_disconnection(self, reason="Unknown"):
             except Exception as e:
                 logger.warning(f" Error sending disconnection updates: {e}")
 
-        # Close connections and clear state (we might reconnect with same client)
+        # Close the network client's sockets and clear remaining state
         self.close()
 
         logger.info("Disconnection cleanup complete - Ready for new connection")
         logger.info("Use web interface settings to reconnect to BlueSky server")
 
     def close(self):
-        """Close all network connections and clear cached proxy state.
+        """Close the network client's sockets and clear cached proxy state.
 
-        Mirrors BlueSky's own ``close()``: shuts the network client's sockets
-        (the ZMQ context is left to the client), resets connection monitoring,
-        and clears tracked nodes/servers, data caches, emission timestamps,
-        and the pending command dictionary.
+        The client instance itself is kept (only ``stop_client`` destroys it);
+        the app reconnects by creating a fresh ``BlueSkyProxy``, so this only
+        has to release the sockets and forget the cached server state.
         """
-        # Disable reconnection first
         self.proxy.allow_reconnection = False
 
-        # Just close the network client - don't destroy ZMQ context
-        # The app creates a completely new BlueSkyProxy instance for reconnection
-        try:
-            logger.debug(" Closing network client...")
-            if self.proxy.bluesky_client:
-                self.proxy.bluesky_client.close()
-            logger.info(" Network client closed successfully")
-        except Exception as e:
-            logger.error(f" Error closing network client: {e}")
-
-        # We reuse the same network client instance - just close its sockets
-
-        # Reset connection monitoring
-        self.proxy.was_connected = False
-        self.proxy.last_successful_update = time.time()
-
-        # Clear all tracked state
-        self.proxy.tracked_nodes.clear()
-        self.proxy.tracked_servers.clear()
-
-        # Clear active node reference to prevent showing corrupted data
-        if hasattr(self.proxy.bluesky_client, "act_id"):
-            self.proxy.bluesky_client.act_id = None
-
-        # Clear data caches
-        self.proxy.traffic_data = {}
-        self.proxy.sim_data = {}
-        self.proxy.echo_data = {}
-
-        # Reset emission timestamps
-        self.proxy.last_siminfo_emit = 0
-        self.proxy.last_acdata_emit = 0
-        self.proxy.last_node_info_emit = 0
-
-        # Clear current map bounds
-        self.proxy.current_bbox = None
-
-        # Clear command dictionary
-        self.proxy.cmddict.clear()
-
-        # Following ZMQ pattern: clear client reference after closing
-        # (new client will be created when reconnecting)
+        client = self.proxy.bluesky_client
+        if client is not None:
+            try:
+                client.close()
+                logger.info(" Network client closed successfully")
+            except Exception as e:
+                logger.error(f" Error closing network client: {e}")
+            # Clear the active node so no stale data can be attributed to it.
+            if hasattr(client, "act_id"):
+                client.act_id = None
 
+        self.proxy.data_mgr._reset_cached_state()
         logger.debug(" Client state cleared and connections closed")
 
     def stop_client(self, context="disconnect"):
@@ -439,35 +374,3 @@ def _close_bluesky_client(self):
                 # Following ZMQ pattern: destroy client instance after closing sockets
                 self.proxy.bluesky_client = None
                 logger.info("Network client instance destroyed")
-
-    def reconnect(self, hostname=None):
-        """Reconnect to the BlueSky server with fresh ZMQ resources.
-
-        Stops the current client, clears state, then starts a new connection
-        via ``start_client``.
-
-        Args:
-            hostname (str | None): BlueSky server hostname/IP to reconnect
-                to. When None, the previously configured host is reused.
-
-        Raises:
-            Exception: Propagated from ``start_client`` if reconnection fails.
-        """
-        logger.info("Reconnecting to BlueSky server...")
-
-        # Following ZMQ pattern: close sockets and destroy context first
-        self.stop_client("disconnect")
-
-        # Wait briefly for ZMQ cleanup to complete
-        time.sleep(0.2)
-
-        # Clear state and prepare for fresh connection
-        self.proxy.data_mgr._clear_state()
-
-        # Following ZMQ pattern: create fresh context and sockets
-        try:
-            self.start_client(hostname=hostname)
-            logger.info(" Reconnection successful with fresh ZMQ resources")
-        except Exception as e:
-            logger.error(f" Reconnection failed: {e}")
-            raise
diff --git a/WebATM/proxy/managers/data_manager.py b/WebATM/proxy/managers/data_manager.py
index bf6e89a..3c7ebf3 100644
--- a/WebATM/proxy/managers/data_manager.py
+++ b/WebATM/proxy/managers/data_manager.py
@@ -122,41 +122,45 @@ def backup_data_emit(self):
         # Schedule next backup emission
         self.start_backup_timer()
 
-    def _clear_state(self, context="disconnect"):
-        """Clear all cached client state after a stop or disconnect.
-
-        Args:
-            context (str): Cleanup context — ``"disconnect"`` for
-                reconnection, ``"manual"`` for a user-initiated disconnect,
-                ``"shutdown"`` for app termination. Only affects the final
-                log message.
+    def _reset_cached_state(self):
+        """Reset connection monitoring and drop all cached BlueSky state.
+
+        The single implementation of "forget everything we knew about the
+        server": tracked nodes/servers, data caches, emission throttles, map
+        bounds and the command dictionary. Shared by every teardown path
+        (``stop_client``, ``ConnectionManager.close`` and
+        ``_handle_disconnection``) so the paths cannot drift apart.
         """
-        # Reset connection monitoring
         self.proxy.was_connected = False
         self.proxy.last_successful_update = time.time()
 
-        # Clear all tracked state
         self.proxy.tracked_nodes.clear()
         self.proxy.tracked_servers.clear()
 
-        # Clear data caches
         self.proxy.traffic_data = {}
         self.proxy.sim_data = {}
         self.proxy.echo_data = {}
         self.proxy.poly_data_by_node.clear()
         self.proxy.polyline_data_by_node.clear()
 
-        # Reset emission timestamps
         self.proxy.last_siminfo_emit = 0
         self.proxy.last_acdata_emit = 0
         self.proxy.last_node_info_emit = 0
 
-        # Clear current map bounds
         self.proxy.current_bbox = None
-
-        # Clear command dictionary
         self.proxy.cmddict.clear()
 
+    def _clear_state(self, context="disconnect"):
+        """Clear all cached client state after a stop or disconnect.
+
+        Args:
+            context (str): Cleanup context — ``"disconnect"`` for
+                reconnection, ``"manual"`` for a user-initiated disconnect,
+                ``"shutdown"`` for app termination. Only affects the final
+                log message.
+        """
+        self._reset_cached_state()
+
         # Emit updated node info to show disconnection
         if self.proxy.socketio and self.proxy.connected_clients > 0:
             try:
diff --git a/tests/test_connection_manager.py b/tests/test_connection_manager.py
new file mode 100644
index 0000000..90538f0
--- /dev/null
+++ b/tests/test_connection_manager.py
@@ -0,0 +1,133 @@
+"""Tests for WebATM.proxy.managers.connection_manager.ConnectionManager.
+
+Covers the teardown paths (``close``, ``stop_client``,
+``_handle_disconnection``), which all funnel their cached-state clearing
+through ``DataManager._reset_cached_state`` so the paths cannot drift apart.
+"""
+
+
+def _seed_cached_state(proxy):
+    """Populate every cache a teardown is expected to clear."""
+    proxy.tracked_nodes["n1"] = {"node_id": b"NODE\x81"}
+    proxy.tracked_servers[b"SRV\x80\x80"] = {"server_id": b"SRV\x80\x80"}
+    proxy.traffic_data = {"id": ["AC1"]}
+    proxy.sim_data = {"scenname": "demo"}
+    proxy.echo_data = {"text": "hello"}
+    proxy.poly_data_by_node["n1"] = {"polys": {}}
+    proxy.polyline_data_by_node["n1"] = {"polys": {}}
+    proxy.last_siminfo_emit = 123.0
+    proxy.last_acdata_emit = 123.0
+    proxy.last_node_info_emit = 123.0
+    proxy.current_bbox = (0.0, 0.0, 1.0, 1.0)
+    proxy.was_connected = True
+
+
+def _assert_cached_state_cleared(proxy):
+    assert proxy.tracked_nodes == {}
+    assert proxy.tracked_servers == {}
+    assert proxy.traffic_data == {}
+    assert proxy.sim_data == {}
+    assert proxy.echo_data == {}
+    assert proxy.poly_data_by_node == {}
+    assert proxy.polyline_data_by_node == {}
+    assert proxy.last_siminfo_emit == 0
+    assert proxy.last_acdata_emit == 0
+    assert proxy.last_node_info_emit == 0
+    assert proxy.current_bbox is None
+    assert proxy.cmddict == {}
+    assert proxy.was_connected is False
+
+
+class TestClose:
+    def test_closes_client_and_clears_cached_state(self, proxy, fake_client):
+        proxy.bluesky_client = fake_client
+        fake_client.act_id = b"NODE\x81"
+        _seed_cached_state(proxy)
+
+        proxy.connection_mgr.close()
+
+        assert fake_client.closed is True
+        assert fake_client.act_id is None
+        assert proxy.allow_reconnection is False
+        _assert_cached_state_cleared(proxy)
+        # close() keeps the client instance; only stop_client destroys it.
+        assert proxy.bluesky_client is fake_client
+
+    def test_without_client_does_not_raise(self, proxy):
+        _seed_cached_state(proxy)
+        proxy.connection_mgr.close()
+        _assert_cached_state_cleared(proxy)
+
+    def test_client_close_error_still_clears_state(self, proxy, fake_client):
+        def boom():
+            raise RuntimeError("simulated socket failure")
+
+        fake_client.close = boom
+        proxy.bluesky_client = fake_client
+        _seed_cached_state(proxy)
+
+        proxy.connection_mgr.close()  # must not raise
+
+        _assert_cached_state_cleared(proxy)
+
+
+class TestStopClient:
+    def test_destroys_client_and_clears_state(self, proxy, fake_client):
+        proxy.bluesky_client = fake_client
+        proxy.running = True
+        _seed_cached_state(proxy)
+
+        proxy.connection_mgr.stop_client("manual")
+
+        assert proxy.running is False
+        assert proxy.allow_reconnection is False
+        assert fake_client.closed is True
+        # Unlike close(), stop_client destroys the client instance.
+        assert proxy.bluesky_client is None
+        _assert_cached_state_cleared(proxy)
+
+
+class TestHandleDisconnection:
+    def test_stops_clears_and_notifies_browsers(
+        self, proxy, fake_socketio, fake_client
+    ):
+        proxy.bluesky_client = fake_client
+        proxy.running = True
+        _seed_cached_state(proxy)
+
+        proxy.connection_mgr._handle_disconnection("test reason")
+
+        assert proxy.running is False
+        assert proxy.allow_reconnection is False
+        assert proxy.connection_failures == 0
+        assert fake_client.closed is True
+        _assert_cached_state_cleared(proxy)
+
+        # Browsers get connection_status(False), the map-clearing payloads,
+        # and a node_info reflecting the already-cleared (empty) state.
+        assert fake_socketio.last("connection_status")["connected"] is False
+        assert fake_socketio.last("acdata")["id"] == []
+        assert fake_socketio.count("server_disconnected") == 1
+        node_info = fake_socketio.last("node_info")
+        assert node_info["nodes"] == {}
+        assert node_info["total_nodes"] == 0
+        assert node_info["active_node"] is None
+
+    def test_no_emits_without_web_clients(self, proxy, fake_socketio, fake_client):
+        proxy.bluesky_client = fake_client
+        proxy.running = True
+        proxy.connected_clients = 0
+        _seed_cached_state(proxy)
+
+        proxy.connection_mgr._handle_disconnection()
+
+        assert fake_socketio.emitted == []
+        assert proxy.running is False
+        _assert_cached_state_cleared(proxy)
+
+    def test_already_disconnected_does_not_emit_status(self, proxy, fake_socketio):
+        proxy.was_connected = False
+        proxy.connection_mgr._handle_disconnection()
+        # No connection_status flip when we were never connected; the cleared
+        # payloads still go out so browsers converge on the empty state.
+        assert fake_socketio.count("connection_status") == 0

From 2e3d6dee64d13b8311d633559d48e54054ebfad7 Mon Sep 17 00:00:00 2001
From: Andres Morfin Veytia <78442543+amorfinv@users.noreply.github.com>
Date: Tue, 25 Aug 2026 17:53:29 +0200
Subject: [PATCH 2/4] Share app bootstrap between dev server and WSGI entry
 points

start_WebATM, script/wsgi.py, and script/wsgi_integrated.py each
duplicated the create_app + BLUESKY_SERVER_HOST wiring. Extract
create_configured_app() in WebATM.main and use it everywhere, add
allow_unsafe_werkzeug to the python-run fallback of the WSGI scripts,
and fix a stale module name in the app.py docstring. Tests cover the
new bootstrap and import both WSGI scripts.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 WebATM/app.py             | 19 +++++--------
 WebATM/main.py            | 33 +++++++++++++++------
 script/wsgi.py            | 31 ++++++--------------
 script/wsgi_integrated.py | 36 +++++++----------------
 tests/test_app.py         | 60 +++++++++++++++++++++++++++++++++++++++
 5 files changed, 109 insertions(+), 70 deletions(-)

diff --git a/WebATM/app.py b/WebATM/app.py
index 7bdc540..1f73514 100644
--- a/WebATM/app.py
+++ b/WebATM/app.py
@@ -7,7 +7,7 @@
 
 - server/session_manager.py: Session tracking
 - server/routes.py: Basic Flask routes (index, commands, config, health)
-- server/server_status.py: BlueSky server status
+- server/bluesky_server_status.py: BlueSky server status
 - server/socket_handlers.py: Socket.IO event handlers
 """
 
@@ -84,12 +84,9 @@ def create_app():
     bluesky_proxy.socketio = socketio
     set_bluesky_proxy(bluesky_proxy)  # Set it globally for the subscriber callbacks
 
-    # NB: subscribers are NOT registered here. The proxy creates its network
-    # client lazily on connect (ZMQ pattern), so at app-creation time
-    # bluesky_client is still None and there is nothing to attach to.
-    # register_subscribers() is therefore called on connect instead -- by the
-    # /api/server/config route (standalone) and by the auto-start hook
-    # (integrated), both right after start_client() builds the client.
+    # NB: subscribers are NOT registered here — the network client doesn't
+    # exist yet. register_subscribers() runs on connect instead (the
+    # /api/server/config route, or the integrated auto-start hook).
 
     # Store proxy reference in app for access in routes
     app.bluesky_proxy = bluesky_proxy
@@ -121,11 +118,9 @@ def handle_exception(e):
     # Harmless and unused in the default build.
     app.session_manager = session_manager
 
-    # Optional integrated extensions: BlueSky server lifecycle control and
-    # live log streaming. This is a no-op in the default build -- the
-    # WEBATM_INTEGRATED env var is unset and the webatm_integrated package is
-    # not installed, so the import is skipped or caught. The core package
-    # never imports webatm_integrated; the dependency points the other way.
+    # Optional integrated extensions (server lifecycle control, live log
+    # streaming). A no-op in the default build: the env var is unset and the
+    # webatm_integrated package is not installed.
     if os.environ.get("WEBATM_INTEGRATED") == "1":
         try:
             import webatm_integrated
diff --git a/WebATM/main.py b/WebATM/main.py
index 0f1eb3b..1812997 100644
--- a/WebATM/main.py
+++ b/WebATM/main.py
@@ -8,6 +8,29 @@
 logger = get_logger()
 
 
+def create_configured_app(bluesky_host=None):
+    """Create the app with the default BlueSky server IP set on the proxy.
+
+    Shared bootstrap for every entry point: the dev server
+    (:func:`start_WebATM`) and the gunicorn WSGI scripts
+    (``script/wsgi.py``, ``script/wsgi_integrated.py``). The proxy is only
+    configured, not connected.
+
+    Args:
+        bluesky_host (str | None): BlueSky server hostname/IP. Falls back to
+            the ``BLUESKY_SERVER_HOST`` environment variable, then
+            ``"localhost"``.
+
+    Returns:
+        tuple: The ``(app, socketio)`` pair from :func:`WebATM.app.create_app`.
+    """
+    bluesky_host = bluesky_host or os.environ.get("BLUESKY_SERVER_HOST", "localhost")
+    app, socketio = create_app()
+    app.bluesky_proxy.server_ip = bluesky_host
+    logger.info(f"Default BlueSky server IP set to: {bluesky_host}")
+    return app, socketio
+
+
 def start_WebATM(hostname=None, port=8082, debug=False):
     """Start the WebATM web server.
 
@@ -26,18 +49,10 @@ def start_WebATM(hostname=None, port=8082, debug=False):
             set, takes precedence.
         debug (bool): Whether to run the Socket.IO server in debug mode.
     """
-    # Get BlueSky server hostname from environment variable or parameter
-    bluesky_host = hostname or os.environ.get("BLUESKY_SERVER_HOST", "localhost")
     web_port = int(os.environ.get("WEB_PORT", port))
     web_host = os.environ.get("WEB_HOST", "localhost")
 
-    # create the app
-    app, socketio = create_app()
-
-    # Set default server IP on the client but don't connect - wait for user to configure
-    app.bluesky_proxy.server_ip = bluesky_host
-    logger.info("BlueSky Proxy initialized (not connected to BlueSky server)")
-    logger.info(f"Default BlueSky server IP set to: {bluesky_host}")
+    app, socketio = create_configured_app(hostname)
     logger.info("Ready - Connect to BlueSky server via WebATM")
 
     try:
diff --git a/script/wsgi.py b/script/wsgi.py
index d5a08b1..a573fa8 100644
--- a/script/wsgi.py
+++ b/script/wsgi.py
@@ -1,40 +1,25 @@
 #!/usr/bin/env python
-"""
-WSGI entry point for production deployment with gunicorn.
+"""WSGI entry point for production deployment with gunicorn.
 
 Run it with a threaded worker — the SocketIO instance is created with
-``async_mode="threading"`` (see ``WebATM.app.create_app``), and gunicorn 26+
-no longer ships an eventlet worker:
+``async_mode="threading"`` (see ``WebATM.app.create_app``):
 
     gunicorn --worker-class gthread --threads 4 -w 1 --bind 0.0.0.0:8082 wsgi:app
 """
 
 import os
 
-from WebATM.app import create_app
 from WebATM.logger import get_logger
+from WebATM.main import create_configured_app
 
 logger = get_logger()
 
-# Get configuration from environment variables
-bluesky_host = os.environ.get("BLUESKY_SERVER_HOST", "localhost")
-web_port = int(os.environ.get("WEB_PORT", 8082))
-web_host = os.environ.get("WEB_HOST", "0.0.0.0")
-
-# Create the Flask app and SocketIO instance
-app, socketio = create_app()
-
-# Set default server IP on the client but don't connect
-app.bluesky_proxy.server_ip = bluesky_host
-logger.info("BlueSky Proxy initialized (not connected to BlueSky server)")
-logger.info(f"Default BlueSky server IP set to: {bluesky_host}")
+app, socketio = create_configured_app()
 logger.info("Ready - Connect to BlueSky server via WebATM")
 
-# Note: Under gunicorn's threaded worker, Flask-SocketIO (threading mode,
-# WebSocket via simple-websocket) handles everything automatically. We just
-# expose the Flask app, not socketio.
-
 if __name__ == "__main__":
-    # This won't be used by gunicorn, but allows testing with python wsgi.py
+    # Fallback for testing without gunicorn: python wsgi.py
+    web_host = os.environ.get("WEB_HOST", "0.0.0.0")
+    web_port = int(os.environ.get("WEB_PORT", 8082))
     logger.info(f"Starting WebATM on http://{web_host}:{web_port}")
-    socketio.run(app, host=web_host, port=web_port)
+    socketio.run(app, host=web_host, port=web_port, allow_unsafe_werkzeug=True)
diff --git a/script/wsgi_integrated.py b/script/wsgi_integrated.py
index e0bf672..ada6708 100644
--- a/script/wsgi_integrated.py
+++ b/script/wsgi_integrated.py
@@ -1,45 +1,29 @@
 #!/usr/bin/env python
 """WSGI entry point for the ``webatm-integrated`` build.
 
-Like ``wsgi.py``, this entry point uses plain threading (no monkey patching):
-the integrated build reads a blocking subprocess pipe in a background thread.
-Run it with a threaded worker, e.g.::
+Same threaded-worker setup as ``wsgi.py`` (no monkey patching; the
+integrated build reads a blocking subprocess pipe in a background thread):
 
     gunicorn --worker-class gthread --threads 4 -w 1 --bind 0.0.0.0:8082 wsgi_integrated:app
-
-The core SocketIO is created with ``async_mode="threading"``, which is correct
-under a gthread worker.
 """
 
 import os
 
-from WebATM.app import create_app
 from WebATM.logger import get_logger
+from WebATM.main import create_configured_app
 
 logger = get_logger()
 
-# Ensure the integrated hook in WebATM.app.create_app() fires even if the
-# orchestrator forgot to set it. create_app() reads this at call time (below),
-# so setting it here is sufficient.
+# create_configured_app() -> create_app() reads this at call time, so setting
+# it here ensures the integrated hook fires even if the orchestrator forgot.
 os.environ.setdefault("WEBATM_INTEGRATED", "1")
 
-# Get configuration from environment variables
-bluesky_host = os.environ.get("BLUESKY_SERVER_HOST", "localhost")
-web_port = int(os.environ.get("WEB_PORT", 8082))
-web_host = os.environ.get("WEB_HOST", "0.0.0.0")
-
-# Create the Flask app and SocketIO instance (integrated extensions register here)
-app, socketio = create_app()
-
-# Set the BlueSky host. The integrated build auto-starts the bundled BlueSky
-# server and connects the proxy on boot (see webatm_integrated.auto_start); set
-# WEBATM_AUTO_START=0 to instead start it manually from the web UI.
-app.bluesky_proxy.server_ip = bluesky_host
-logger.info("WebATM (integrated) initialized")
-logger.info(f"Default BlueSky server IP set to: {bluesky_host}")
+app, socketio = create_configured_app()
 logger.info("Ready - BlueSky server auto-starting (WEBATM_AUTO_START=0 to disable)")
 
 if __name__ == "__main__":
-    # Not used by gunicorn, but allows testing with `python wsgi_integrated.py`.
+    # Fallback for testing without gunicorn: python wsgi_integrated.py
+    web_host = os.environ.get("WEB_HOST", "0.0.0.0")
+    web_port = int(os.environ.get("WEB_PORT", 8082))
     logger.info(f"Starting WebATM (integrated) on http://{web_host}:{web_port}")
-    socketio.run(app, host=web_host, port=web_port)
+    socketio.run(app, host=web_host, port=web_port, allow_unsafe_werkzeug=True)
diff --git a/tests/test_app.py b/tests/test_app.py
index d9d99ed..8d36dac 100644
--- a/tests/test_app.py
+++ b/tests/test_app.py
@@ -68,6 +68,66 @@ def test_integrated_hook_enabled_without_package_is_safe(self, monkeypatch):
             set_bluesky_proxy(None)
 
 
+class TestCreateConfiguredApp:
+    """The shared bootstrap used by start_WebATM and the wsgi entry points."""
+
+    def test_defaults_to_localhost(self, monkeypatch):
+        from WebATM.main import create_configured_app
+
+        monkeypatch.delenv("BLUESKY_SERVER_HOST", raising=False)
+        app, socketio = create_configured_app()
+        try:
+            assert app.bluesky_proxy.server_ip == "localhost"
+            assert socketio is not None
+        finally:
+            set_bluesky_proxy(None)
+
+    def test_env_var_sets_server_ip(self, monkeypatch):
+        from WebATM.main import create_configured_app
+
+        monkeypatch.setenv("BLUESKY_SERVER_HOST", "10.1.2.3")
+        app, _ = create_configured_app()
+        try:
+            assert app.bluesky_proxy.server_ip == "10.1.2.3"
+        finally:
+            set_bluesky_proxy(None)
+
+    def test_explicit_host_beats_env_var(self, monkeypatch):
+        from WebATM.main import create_configured_app
+
+        monkeypatch.setenv("BLUESKY_SERVER_HOST", "10.1.2.3")
+        app, _ = create_configured_app("192.168.0.9")
+        try:
+            assert app.bluesky_proxy.server_ip == "192.168.0.9"
+        finally:
+            set_bluesky_proxy(None)
+
+
+class TestWsgiEntryPoints:
+    """The gunicorn entry scripts must build a configured app on import."""
+
+    @pytest.mark.parametrize("script_name", ["wsgi.py", "wsgi_integrated.py"])
+    def test_script_builds_app(self, script_name, monkeypatch, tmp_path):
+        import importlib.util
+        from pathlib import Path
+
+        # Pin WEBATM_INTEGRATED to "0" (setdefault in wsgi_integrated.py then
+        # leaves it alone, and monkeypatch restores it afterwards) so neither
+        # script registers the integrated extensions during the test.
+        monkeypatch.setenv("WEBATM_INTEGRATED", "0")
+        monkeypatch.setenv("BLUESKY_SERVER_HOST", "wsgi.test.host")
+
+        script = Path(__file__).parent.parent / "script" / script_name
+        spec = importlib.util.spec_from_file_location(f"test_{script.stem}", script)
+        module = importlib.util.module_from_spec(spec)
+        try:
+            spec.loader.exec_module(module)
+            assert module.app.bluesky_proxy.server_ip == "wsgi.test.host"
+            assert module.socketio is not None
+        finally:
+            set_bluesky_proxy(None)
+
+
 class TestHealthAndStatus:
     def test_health_returns_200(self, client):
         resp = client.get("/health")

From 96ad8407980866a7b12d89552fa766ed89919c15 Mon Sep 17 00:00:00 2001
From: Andres Morfin Veytia <78442543+amorfinv@users.noreply.github.com>
Date: Tue, 25 Aug 2026 17:53:35 +0200
Subject: [PATCH 3/4] Decode BlueSky numpy arrays with np.frombuffer instead of
 a dtype map

Replace the hand-rolled struct-format table (which only knew five dtypes
and fell back to hex for everything else, e.g. float16, uint64, string
arrays) with np.frombuffer on the transmitted dtype string, matching
bluesky.network.npcodec. Hex fallback is kept for invalid dtypes and
truncated buffers, with tests for the new and fallback paths.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 WebATM/utils.py     | 55 ++++++++++-----------------------------------
 tests/test_utils.py | 44 ++++++++++++++++++++++++++++++++++--
 2 files changed, 54 insertions(+), 45 deletions(-)

diff --git a/WebATM/utils.py b/WebATM/utils.py
index 859a233..25d37ac 100644
--- a/WebATM/utils.py
+++ b/WebATM/utils.py
@@ -33,54 +33,23 @@ def make_json_serializable(obj):
     elif isinstance(obj, np.floating):
         return float(obj)
     elif isinstance(obj, dict):
-        # Handle BlueSky's serialized numpy arrays
         if b"numpy" in obj and b"data" in obj and b"type" in obj and b"shape" in obj:
+            # BlueSky msgpack-encoded numpy array; decode it the same way as
+            # bluesky.network.npcodec.decode_ndarray, flattened for JSON.
             try:
-                # This is a BlueSky serialized numpy array - deserialize it
-                import struct
-
-                dtype = (
-                    obj[b"type"].decode()
-                    if isinstance(obj[b"type"], bytes)
-                    else obj[b"type"]
-                )
-                shape = obj[b"shape"]
-                data_bytes = obj[b"data"]
-
-                # Convert numpy dtype string to struct format
-                dtype_map = {
-                    "<f8": "d",  # double
-                    "<f4": "f",  # float
-                    "<i8": "q",  # long long
-                    "<i4": "i",  # int
-                    "|b1": "?",  # bool
-                }
-
-                if dtype in dtype_map:
-                    format_char = dtype_map[dtype]
-                    num_elements = 1
-                    for dim in shape:
-                        num_elements *= dim
-
-                    # Unpack the binary data
-                    values = list(
-                        struct.unpack(f"<{num_elements}{format_char}", data_bytes)
-                    )
-
-                    # Return as list for JSON serialization
-                    return values
-                else:
-                    logger.warning(
-                        f"Utils: Unknown numpy dtype {dtype}, returning raw data"
-                    )
-                    return obj[b"data"].hex()  # Return as hex string if we can't parse
-
+                dtype = obj[b"type"]
+                if isinstance(dtype, bytes):
+                    dtype = dtype.decode()
+                return np.frombuffer(obj[b"data"], dtype=np.dtype(dtype)).tolist()
             except Exception as e:
                 logger.warning(f"Utils: Error deserializing numpy array: {e}")
-                # Fall back to converting dict normally
-                pass
+                data_bytes = obj[b"data"]
+                return (
+                    data_bytes.hex()
+                    if isinstance(data_bytes, bytes)
+                    else str(data_bytes)
+                )
 
-        # Normal dict processing
         return {
             (key.decode() if isinstance(key, bytes) else key): make_json_serializable(
                 value
diff --git a/tests/test_utils.py b/tests/test_utils.py
index 869fe10..6f1c51b 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -122,16 +122,56 @@ def test_bluesky_serialized_numpy_multidim_shape(self):
         # The deserializer flattens to a single list of values.
         assert make_json_serializable(obj) == [1.0, 2.0, 3.0, 4.0]
 
-    def test_bluesky_serialized_unknown_dtype_falls_back_to_hex(self):
+    def test_bluesky_serialized_numpy_half_float_array(self):
+        arr = np.array([1.5, 2.5], dtype=np.float16)
+        obj = {
+            b"numpy": True,
+            b"data": arr.tobytes(),
+            b"type": arr.dtype.str,  # "<f2"
+            b"shape": [2],
+        }
+        assert make_json_serializable(obj) == [1.5, 2.5]
+
+    def test_bluesky_serialized_numpy_uint64_array(self):
+        arr = np.array([1, 2, 3], dtype=np.uint64)
+        obj = {
+            b"numpy": True,
+            b"data": arr.tobytes(),
+            b"type": arr.dtype.str,  # "<u8"
+            b"shape": [3],
+        }
+        assert make_json_serializable(obj) == [1, 2, 3]
+
+    def test_bluesky_serialized_numpy_string_array(self):
+        arr = np.array(["KL01", "KL02"])
+        obj = {
+            b"numpy": True,
+            b"data": arr.tobytes(),
+            b"type": arr.dtype.str,  # "<U4"
+            b"shape": [2],
+        }
+        assert make_json_serializable(obj) == ["KL01", "KL02"]
+
+    def test_bluesky_serialized_invalid_dtype_falls_back_to_hex(self):
         data_bytes = b"\x01\x02\x03\x04"
         obj = {
             b"numpy": True,
             b"data": data_bytes,
-            b"type": "<f2",  # not in the dtype map
+            b"type": "not-a-dtype",
             b"shape": [2],
         }
         assert make_json_serializable(obj) == data_bytes.hex()
 
+    def test_bluesky_serialized_truncated_buffer_falls_back_to_hex(self):
+        data_bytes = b"\x01\x02\x03"  # not a multiple of 8 bytes
+        obj = {
+            b"numpy": True,
+            b"data": data_bytes,
+            b"type": "<f8",
+            b"shape": [1],
+        }
+        assert make_json_serializable(obj) == data_bytes.hex()
+
     def test_plain_dict_without_all_numpy_keys_is_treated_normally(self):
         # Has b"numpy" but is missing b"data"/b"type"/b"shape": normal dict path.
         obj = {b"numpy": True, b"other": 1}

From bb12d8c64e2c5bccfe89e9f8a6c72858c8791390 Mon Sep 17 00:00:00 2001
From: Andres Morfin Veytia <78442543+amorfinv@users.noreply.github.com>
Date: Tue, 25 Aug 2026 17:53:44 +0200
Subject: [PATCH 4/4] Fix 3D shape extrusion and style-change re-rendering

Polygons without altitude data were extruded to a hard-coded 1000 m; they
now stay flat in the 2D layers, matching how BlueSky treats shapes with
no vertical extent. The new extrusionBounds() helper (with vitest
coverage) also clamps the base into [0, top] so malformed data cannot
produce an inside-out extrusion.

On map style changes, ShapeRenderer.onStyleChange() re-initialized both
renderers and then rendered shapes again on top of the subscription-driven
render, and Shape3DRenderer re-initialized itself independently. The 3D
renderer now just marks its layers as lost and lets the owning
ShapeRenderer rebuild everything once.

The duplicated "recreate source then retry setData" logic moves into a
shared updateSourceWithRecovery() helper in utils/maplibre.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 frontend/src/ui/map/shapes/Shape3DRenderer.ts | 59 ++++++++-----------
 frontend/src/ui/map/shapes/ShapeRenderer.ts   | 41 ++++---------
 frontend/src/ui/map/shapes/extrusion.test.ts  | 37 ++++++++++++
 frontend/src/ui/map/shapes/extrusion.ts       | 25 ++++++++
 frontend/src/utils/maplibre.ts                | 21 +++++++
 5 files changed, 116 insertions(+), 67 deletions(-)
 create mode 100644 frontend/src/ui/map/shapes/extrusion.test.ts
 create mode 100644 frontend/src/ui/map/shapes/extrusion.ts

diff --git a/frontend/src/ui/map/shapes/Shape3DRenderer.ts b/frontend/src/ui/map/shapes/Shape3DRenderer.ts
index ac5d679..a61d85e 100644
--- a/frontend/src/ui/map/shapes/Shape3DRenderer.ts
+++ b/frontend/src/ui/map/shapes/Shape3DRenderer.ts
@@ -1,15 +1,16 @@
-import { GeoJSONSource } from 'maplibre-gl';
 import { PolygonShape, DisplayOptions } from '../../../data/types';
 import type { MapDisplay } from '../MapDisplay';
 import type { StateManager } from '../../../core/StateManager';
-import { featureCollection, polygonFeature, toLngLatCoords } from '../../../utils/geojson';
+import { polygonFeature, toLngLatCoords } from '../../../utils/geojson';
 import { logger } from '../../../utils/Logger';
+import { extrusionBounds } from './extrusion';
 import {
     ensureGeoJSONSource,
     ensureLayer,
     setLayerVisibility,
     safeRemoveLayer,
-    safeRemoveSource
+    safeRemoveSource,
+    updateSourceWithRecovery
 } from '../../../utils/maplibre';
 
 /**
@@ -72,47 +73,31 @@ export class Shape3DRenderer {
     }
 
     /**
-     * Render polygons as extruded 3D shapes.
-     * Polygons with altitude data use their actual top/bottom values.
-     * Polygons without altitude data get a default extrusion height.
+     * Render polygons with altitude data as extruded 3D volumes. Polygons
+     * without altitude data are skipped - they stay flat in the 2D layers,
+     * matching how BlueSky treats shapes with no vertical extent.
      */
     public renderExtrudedPolygons(polygons: PolygonShape[]): void {
         const map = this.mapDisplay.getMap();
         if (!map) return;
 
         const displayOptions = this.stateManager.getDisplayOptions();
-        const DEFAULT_EXTRUSION_HEIGHT = 1000; // meters - default height for polygons without altitude
-
-        const features = polygons.map(poly => {
-            const hasAltitude = poly.topAltitude !== undefined && poly.topAltitude !== null;
-            const topAlt = hasAltitude
-                ? (poly.topAltitude || 0)
-                : DEFAULT_EXTRUSION_HEIGHT;
-            const bottomAlt = hasAltitude
-                ? (poly.bottomAltitude || 0)
-                : 0;
-
-            return polygonFeature(toLngLatCoords(poly.coordinates), {
+
+        const features: GeoJSON.Feature[] = [];
+        for (const poly of polygons) {
+            const bounds = extrusionBounds(poly);
+            if (!bounds) continue;
+            features.push(polygonFeature(toLngLatCoords(poly.coordinates), {
                 name: poly.name,
                 fillColor: poly.fillColor || displayOptions.shapeFillColor || '#ff00ff',
-                extrusionHeight: topAlt,
-                extrusionBase: bottomAlt
-            });
-        });
-
-        // Re-create the layers once if the source is missing (initial load
-        // or a style change removed it), then retry.
-        let source = map.getSource(this.SOURCE_ID) as GeoJSONSource | undefined;
-        if (!source) {
-            logger.debug('Shape3DRenderer', 'Source not found, re-creating layers...');
-            this.setupMapLayers();
-            source = map.getSource(this.SOURCE_ID) as GeoJSONSource | undefined;
+                extrusionHeight: bounds.top,
+                extrusionBase: bounds.base
+            }));
         }
 
-        if (source) {
-            source.setData(featureCollection(features));
-        } else {
-            logger.warn('Shape3DRenderer', `Failed to create source - cannot render ${polygons.length} extruded polygons`);
+        const ok = updateSourceWithRecovery(map, this.SOURCE_ID, features, () => this.setupMapLayers());
+        if (!ok) {
+            logger.warn('Shape3DRenderer', `Failed to create source - cannot render ${features.length} extruded polygons`);
         }
     }
 
@@ -122,10 +107,12 @@ export class Shape3DRenderer {
         setLayerVisibility(map, this.LAYER_ID, displayOptions.show3DOverlay && displayOptions.showShapes);
     }
 
+    /**
+     * Mark the layers as lost after a map style change; the next
+     * initialize() (driven by ShapeRenderer) re-creates them.
+     */
     public onStyleChange(): void {
-        logger.debug('Shape3DRenderer', 'Map style changed - recreating layers');
         this.initialized = false;
-        this.initialize();
     }
 
     public destroy(): void {
diff --git a/frontend/src/ui/map/shapes/ShapeRenderer.ts b/frontend/src/ui/map/shapes/ShapeRenderer.ts
index da20d12..2b622c6 100644
--- a/frontend/src/ui/map/shapes/ShapeRenderer.ts
+++ b/frontend/src/ui/map/shapes/ShapeRenderer.ts
@@ -1,14 +1,14 @@
-import { GeoJSONSource } from 'maplibre-gl';
 import { Shape, PolygonShape, PolylineShape, DisplayOptions } from '../../../data/types';
 import type { MapDisplay } from '../MapDisplay';
 import type { StateManager } from '../../../core/StateManager';
 import { Shape3DRenderer } from './Shape3DRenderer';
-import { featureCollection, lineStringFeature, pointFeature, polygonFeature, toLngLatCoords } from '../../../utils/geojson';
+import { lineStringFeature, pointFeature, polygonFeature, toLngLatCoords } from '../../../utils/geojson';
 import { logger } from '../../../utils/Logger';
 import {
     ensureGeoJSONSource,
     ensureLayer,
-    setLayerVisibility
+    setLayerVisibility,
+    updateSourceWithRecovery
 } from '../../../utils/maplibre';
 
 /**
@@ -39,10 +39,8 @@ export class ShapeRenderer {
         this.stateManager = stateManager;
         this.shape3DRenderer = new Shape3DRenderer(mapDisplay, stateManager);
 
-        // Subscribe once for the renderer's lifetime. initialize() re-runs on
-        // every style change and source recovery, so subscriptions must not
-        // live there - each re-init would stack another listener and multiply
-        // the render work per update.
+        // Subscribe once for the renderer's lifetime; initialize() re-runs on
+        // every style change, so subscribing there would stack listeners.
         this.unsubscribers.push(
             this.stateManager.subscribeToShapes((shapes) => {
                 this.renderShapes(shapes);
@@ -189,16 +187,8 @@ export class ShapeRenderer {
         const map = this.mapDisplay.getMap();
         if (!map) return;
 
-        let source = map.getSource(sourceId) as GeoJSONSource | undefined;
-        if (!source) {
-            logger.debug('ShapeRenderer', `Source ${sourceId} not found, re-creating layers...`);
-            this.setupMapLayers();
-            source = map.getSource(sourceId) as GeoJSONSource | undefined;
-        }
-
-        if (source) {
-            source.setData(featureCollection(features));
-        } else {
+        const ok = updateSourceWithRecovery(map, sourceId, features, () => this.setupMapLayers());
+        if (!ok) {
             logger.warn('ShapeRenderer', `Failed to create source ${sourceId} - cannot render ${features.length} features`);
         }
     }
@@ -311,25 +301,14 @@ export class ShapeRenderer {
     }
 
     /**
-     * Handle map style changes - re-add layers and re-render.
+     * Handle map style changes: mark both renderers' layers as lost, then
+     * let initialize() rebuild them and re-render the stored shapes.
      */
     public onStyleChange(): void {
         logger.debug('ShapeRenderer', 'Map style changed - recreating layers');
         this.initialized = false;
-        this.initialize();
-
         this.shape3DRenderer.onStyleChange();
-
-        // Resize map after it settles from re-adding fill-extrusion layer
-        const map = this.mapDisplay.getMap();
-        if (map) {
-            map.once('idle', () => {
-                this.mapDisplay.resize();
-            });
-        }
-
-        const shapes = this.stateManager.getAllShapes();
-        this.renderShapes(shapes);
+        this.initialize();
     }
 
     public destroy(): void {
diff --git a/frontend/src/ui/map/shapes/extrusion.test.ts b/frontend/src/ui/map/shapes/extrusion.test.ts
new file mode 100644
index 0000000..a63ec9f
--- /dev/null
+++ b/frontend/src/ui/map/shapes/extrusion.test.ts
@@ -0,0 +1,37 @@
+/**
+ * Tests for the 3D extrusion bounds of polygon shapes.
+ */
+import { describe, it, expect } from 'vitest';
+import { extrusionBounds } from './extrusion';
+
+describe('extrusionBounds', () => {
+    it('returns the vertical extent of an altitude-bounded polygon', () => {
+        expect(extrusionBounds({ topAltitude: 2438.4, bottomAltitude: 609.6 }))
+            .toEqual({ top: 2438.4, base: 609.6 });
+    });
+
+    it('defaults a missing bottom altitude to ground level', () => {
+        expect(extrusionBounds({ topAltitude: 1500 })).toEqual({ top: 1500, base: 0 });
+    });
+
+    it('returns null for a polygon without altitude data, so it renders flat', () => {
+        expect(extrusionBounds({})).toBeNull();
+        expect(extrusionBounds({ topAltitude: undefined, bottomAltitude: undefined })).toBeNull();
+    });
+
+    it('returns null when only a bottom bound exists (top is unbounded)', () => {
+        expect(extrusionBounds({ bottomAltitude: 300 })).toBeNull();
+    });
+
+    it('returns null for a non-positive top altitude', () => {
+        expect(extrusionBounds({ topAltitude: 0 })).toBeNull();
+        expect(extrusionBounds({ topAltitude: -500, bottomAltitude: -1000 })).toBeNull();
+    });
+
+    it('clamps the base into [0, top]', () => {
+        expect(extrusionBounds({ topAltitude: 1000, bottomAltitude: -200 }))
+            .toEqual({ top: 1000, base: 0 });
+        expect(extrusionBounds({ topAltitude: 1000, bottomAltitude: 5000 }))
+            .toEqual({ top: 1000, base: 1000 });
+    });
+});
diff --git a/frontend/src/ui/map/shapes/extrusion.ts b/frontend/src/ui/map/shapes/extrusion.ts
new file mode 100644
index 0000000..10f2503
--- /dev/null
+++ b/frontend/src/ui/map/shapes/extrusion.ts
@@ -0,0 +1,25 @@
+import { PolygonShape } from '../../../data/types';
+
+/** Vertical extent of an extruded polygon, in metres above ground. */
+export interface ExtrusionBounds {
+    base: number;
+    top: number;
+}
+
+/**
+ * Vertical extent for a polygon's 3D extrusion, or null when the shape has
+ * no renderable extent and must stay flat. The backend only forwards finite
+ * altitude bounds (BlueSky's +/-1e9 "unbounded" sentinels are stripped), so
+ * a missing top altitude means the shape is a plain 2D area.
+ *
+ * A missing bottom defaults to ground level, and the base is clamped into
+ * [0, top] so malformed data can never produce an inside-out extrusion.
+ */
+export function extrusionBounds(
+    poly: Pick<PolygonShape, 'topAltitude' | 'bottomAltitude'>
+): ExtrusionBounds | null {
+    const top = poly.topAltitude;
+    if (typeof top !== 'number' || top <= 0) return null;
+    const bottom = typeof poly.bottomAltitude === 'number' ? poly.bottomAltitude : 0;
+    return { top, base: Math.min(Math.max(bottom, 0), top) };
+}
diff --git a/frontend/src/utils/maplibre.ts b/frontend/src/utils/maplibre.ts
index fc0b262..bc679a7 100644
--- a/frontend/src/utils/maplibre.ts
+++ b/frontend/src/utils/maplibre.ts
@@ -51,6 +51,27 @@ export function updateSourceFeatures(
     source.setData({ type: 'FeatureCollection', features });
 }
 
+/**
+ * Push features to a GeoJSON source, calling `recreate` once to rebuild the
+ * owning layers when the source is missing (initial load, or a style change
+ * removed it). Returns false when the source still doesn't exist afterwards.
+ */
+export function updateSourceWithRecovery(
+    map: MapLibreMap,
+    sourceId: string,
+    features: GeoJSON.Feature[],
+    recreate: () => void
+): boolean {
+    let source = map.getSource(sourceId) as GeoJSONSource | undefined;
+    if (!source) {
+        recreate();
+        source = map.getSource(sourceId) as GeoJSONSource | undefined;
+    }
+    if (!source) return false;
+    source.setData({ type: 'FeatureCollection', features });
+    return true;
+}
+
 export function setLayerVisibility(
     map: MapLibreMap,
     layerId: string,