fixed startup

app.py

@@ -1,5 +1,6 @@
 # app.py
 import asyncio, logging
+from contextlib import asynccontextmanager
 import gradio as gr
 
 from config import settings
@@ -56,18 +57,69 @@ DECLS = [
 
 listener = RabbitListenerBase(base, instance_name=settings.RABBIT_INSTANCE_NAME, handlers=handlers)
 
-# ----------------- Startup init -----------------
-async def _startup_init():
+# ----------------- Service lifecycle -----------------
+_startup_lock = asyncio.Lock()
+_service_status = "Starting OpenAI MQ + TimesFM"
+_listener_started = False
+
+
+async def _start_listener() -> None:
+    """Connect and register consumers once for this process."""
+    global _listener_started, _service_status
+
+    async with _startup_lock:
+        if _listener_started and base.is_connected():
+            return
+
+        _service_status = "Connecting OpenAI MQ + TimesFM"
+        try:
+            await base.connect()
+            await listener.start(DECLS)
+        except Exception as exc:
+            _listener_started = False
+            _service_status = f"ERROR: {exc}"
+            # A partial declaration must not leave a stale connection around if
+            # the hosting platform retries this process.
+            await base.close()
+            log.exception("RabbitMQ listener startup failed")
+            raise
+
+        _listener_started = True
+        _service_status = "OpenAI MQ + TimesFM: ready"
+        log.info(_service_status)
+
+
+async def _on_startup() -> None:
+    # Fail the application startup if RabbitMQ is unavailable. Serving the UI
+    # without its consumer makes a healthy-looking but non-functional service.
+    await _start_listener()
+
+
+async def _on_shutdown() -> None:
+    global _listener_started, _service_status
+    _listener_started = False
+    _service_status = "Stopped"
+    await base.close()
+
+
+@asynccontextmanager
+async def service_lifespan(app):
+    # Pass this to launch: Gradio replaces demo.app and composes this lifespan
+    # with its own lifecycle (including MCP) on the application it actually serves.
+    await _on_startup()
     try:
-        await base.connect()          # connect to RabbitMQ
-        await listener.start(DECLS)   # start queue listeners
-        return "OpenAI MQ + TimesFM: ready"
-    except Exception as e:
-        log.exception("Startup init failed")
-        return f"ERROR: {e}"
+        yield
+    finally:
+        await _on_shutdown()
+
+
+async def service_status() -> str:
+    """Read-only UI status; lifecycle startup is handled by FastAPI."""
+    return _service_status
+
 
 async def ping():
-    return "ok"
+    return "ok" if _listener_started and base.is_connected() else _service_status
 
 # ----------------- Gradio UI -----------------
 with gr.Blocks(title="OpenAI over RabbitMQ (local vLLM)", theme=gr.themes.Soft()) as demo:
@@ -78,12 +130,17 @@ with gr.Blocks(title="OpenAI over RabbitMQ (local vLLM)", theme=gr.themes.Soft()
             out = gr.Textbox(label="Ping result")
             btn.click(ping, inputs=None, outputs=out)
             init_status = gr.Textbox(label="Startup status", interactive=False)
-            demo.load(fn=_startup_init, inputs=None, outputs=init_status)
+            demo.load(fn=service_status, inputs=None, outputs=init_status)
 
         with gr.Tab("@spaces.GPU Probe"):
             gpu_btn = gr.Button("GPU Ready Probe", variant="primary")
             gpu_out = gr.Textbox(label="GPU Probe Result", interactive=False)
             gpu_btn.click(gpu_entrypoint, inputs=None, outputs=gpu_out)
 
+
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True, debug=True, mcp_server=True)
+    demo.launch(
+        server_name="0.0.0.0", server_port=7860, show_error=True,
+        debug=True, mcp_server=True,
+        app_kwargs={"lifespan": service_lifespan},
+    )

listener.py

@@ -15,6 +15,7 @@ class RabbitListenerBase:
         self._instance_name = instance_name
         self._handlers = handlers
         self._consumers: List[aio_pika.abc.AbstractRobustQueue] = []
+        self._started = False
 
     def _qname(self, exchange: str, routing_keys: List[str]) -> str:
         rk_part = "-".join(sorted([rk for rk in (routing_keys or [""]) if rk])) or ""
@@ -22,6 +23,9 @@ class RabbitListenerBase:
         return f"{self._instance_name}-{exchange}{suffix}"
 
     async def start(self, declarations: List[dict]) -> None:
+        if self._started:
+            return
+
         for d in declarations:
             exch = d["ExchangeName"]
             ttl = d.get("MessageTimeout") or None
@@ -33,6 +37,7 @@ class RabbitListenerBase:
             # manual ack, parity with .NET (autoAck: false)
             await q.consume(self._make_consumer(d["FuncName"]), no_ack=False)
             self._consumers.append(q)
+        self._started = True
 
     def _make_consumer(self, func_name: str):
         handler = self._handlers.get(func_name)

test_app_startup.py

@@ -0,0 +1,107 @@
+"""Exercise real Gradio launch with fake RabbitMQ and forecasting dependencies.
+
+Run with: python -m unittest test_app_startup
+Requires gradio[mcp] (5.42.0, as pinned by the Space) and aio-pika.
+"""
+import asyncio
+from pathlib import Path
+import runpy
+import socket
+import sys
+import types
+import unittest
+from unittest.mock import AsyncMock, Mock, patch
+
+import gradio as gr
+from gradio.routes import App
+from starlette.testclient import TestClient
+
+
+class AppStartupTests(unittest.TestCase):
+    def setUp(self):
+        self.base = Mock()
+        self.connected = False
+
+        async def connect():
+            self.connected = True
+
+        async def close():
+            self.connected = False
+
+        self.base.connect = AsyncMock(side_effect=connect)
+        self.base.close = AsyncMock(side_effect=close)
+        self.base.is_connected = lambda: self.connected
+        self.queue = Mock()
+        self.queue.consume = AsyncMock(return_value="test-consumer")
+        self.base.declare_queue_bind = AsyncMock(return_value=self.queue)
+        servers = Mock(handle_chat_create=AsyncMock(), handle_images_generate=AsyncMock())
+        self.modules = {}
+        for name, attrs in {
+            "config": {"settings": types.SimpleNamespace(
+                RABBIT_EXCHANGE_TYPE="direct", RABBIT_ROUTING_KEY="execute.api",
+                RABBIT_INSTANCE_NAME="startup-test")},
+            "rabbit_base": {"RabbitBase": Mock(return_value=self.base)},
+            "rabbit_repo": {"RabbitRepo": Mock()},
+            "oa_server": {"OpenAIServers": Mock(return_value=servers)},
+            "timesfm_backend": {"TimesFMBackend": Mock(), "StubImagesBackend": Mock()},
+        }.items():
+            module = types.ModuleType(name)
+            module.__dict__.update(attrs)
+            self.modules[name] = module
+        self.modules["spaces"] = None
+
+    def load_app(self, name="startup_test"):
+        with patch.dict(sys.modules, self.modules):
+            return runpy.run_path(str(Path(__file__).with_name("app.py")), run_name=name)
+
+    def test_launch_starts_consumers_without_browser_and_shutdown_closes_connection(self):
+        original_launch = gr.Blocks.launch
+        launched = []
+
+        def launch(blocks, **kwargs):
+            # Execute the production entry point and real launch implementation.
+            # Only change networking/blocking options so this test can finish.
+            launched.append(blocks)
+            with socket.socket() as sock:
+                sock.bind(("127.0.0.1", 0))
+                port = sock.getsockname()[1]
+            kwargs.update(server_name="127.0.0.1", server_port=port,
+                          prevent_thread_lock=True, debug=False, quiet=True,
+                          inbrowser=False, share=False)
+            return original_launch(blocks, **kwargs)
+
+        try:
+            with patch.object(gr.Blocks, "launch", launch):
+                app = self.load_app("__main__")
+            self.base.connect.assert_awaited_once()
+            self.assertEqual(self.queue.consume.await_count, 2)
+            self.assertTrue(app["demo"].mcp_server)
+            for _ in range(2):
+                self.assertIn("ready", asyncio.run(app["service_status"]()))
+                self.assertEqual(asyncio.run(app["ping"]()), "ok")
+            self.assertEqual(self.queue.consume.await_count, 2)
+        finally:
+            for blocks in launched:
+                blocks.close(verbose=False)
+        self.base.close.assert_awaited_once()
+        self.assertFalse(self.connected)
+
+    def test_startup_failure_propagates_and_closes_partial_connection(self):
+        app = self.load_app()
+        self.base.declare_queue_bind.side_effect = RuntimeError("binding failed")
+        # Use Gradio's real application factory and lifespan composition.
+        server = App.create_app(app["demo"],
+                                app_kwargs={"lifespan": app["service_lifespan"]},
+                                mcp_server=True)
+        entered = False
+        with self.assertRaises(Exception):
+            with TestClient(server):
+                entered = True
+        self.assertFalse(entered, "Server became available despite failed startup")
+        self.base.close.assert_awaited_once()
+        self.assertFalse(self.connected)
+        self.assertEqual(asyncio.run(app["service_status"]()), "ERROR: binding failed")
+
+
+if __name__ == "__main__":
+    unittest.main()