update rabbitmq code to make it more robust

listener.py

@@ -6,7 +6,6 @@ import aio_pika
 
 Handler = Callable[[Any], Awaitable[None]]  # payload is envelope["data"]
 
-import logging
 logger = logging.getLogger(__name__)
 
 
@@ -31,27 +30,36 @@ class RabbitListenerBase:
             q = await self._base.declare_queue_bind(
                 exchange=exch, queue_name=qname, routing_keys=rks, ttl_ms=ttl
             )
-            await q.consume(self._make_consumer(d["FuncName"]))
+            # explicit manual-ack, parity with .NET (autoAck: false)
+            await q.consume(self._make_consumer(d["FuncName"]), no_ack=False)
             self._consumers.append(q)
 
     def _make_consumer(self, func_name: str):
         handler = self._handlers.get(func_name)
 
         async def _on_msg(msg: aio_pika.IncomingMessage):
-            async with msg.process():
-                try:
-                    raw_body = msg.body.decode("utf-8", errors="replace")
-                    logger.info(
-                        "Received message for handler '%s': %s",
-                        func_name,
-                        raw_body
-                    )
+            # manual ack after handler completes; no nack/requeue loops
+            try:
+                raw_body = msg.body.decode("utf-8", errors="replace")
+                logger.info("Received message for handler '%s': %s", func_name, raw_body)
 
+                # safe JSON parse to mirror .NET ConvertToObject (no throw)
+                try:
                     envelope = json.loads(raw_body)
-                    data = envelope.get("data", None)
-                    if handler:
-                        await handler(data)
-                except Exception as e:
-                    logger.exception("Error processing message for '%s'", func_name)
+                except Exception:
+                    logger.exception("Invalid JSON for '%s'", func_name)
+                    envelope = {"data": None}
+
+                data = envelope.get("data", None)
+
+                if handler:
+                    await handler(data)
+                else:
+                    logger.error("No handler bound for '%s'", func_name)
+
+                await msg.ack()  # ack on success path
+            except Exception:
+                # match .NET: on exception, do not ack or nack; connection loss will requeue
+                logger.exception("Error processing message for '%s'", func_name)
 
         return _on_msg

oa_server.py

@@ -34,12 +34,12 @@ class OpenAIServers:
     """
 
     def __init__(self, publisher: RabbitRepo,
-                 *, chat_backend=None, images_backend=None):
+                 *, chat_backend: Optional[ChatBackend] = None,
+                 images_backend: Optional[ImagesBackend] = None):
         self._pub = publisher
         self._chat = chat_backend
         self._img  = images_backend
 
-
     # -------- Chat Completions --------
     async def handle_chat_create(self, data: Dict[str, Any]) -> None:
         """
@@ -57,10 +57,17 @@ class OpenAIServers:
 
         try:
             async for chunk in self._chat.stream(data):
-                # CloudEvent-wrapped OpenAI chunk to oa.chat.reply
-                await self._pub.publish("oa.chat.reply", chunk, routing_key=reply_key)
+                try:
+                    await self._pub.publish("oa.chat.reply", chunk, routing_key=reply_key)
+                except Exception:
+                    logger.exception("oaChatCreate: publish failed")
+                    break  # stop streaming on publish failure
+
             # Optional sentinel
-            await self._pub.publish("oa.chat.reply", {"object": "stream.end"}, routing_key=reply_key)
+            try:
+                await self._pub.publish("oa.chat.reply", {"object": "stream.end"}, routing_key=reply_key)
+            except Exception:
+                logger.exception("oaChatCreate: publish sentinel failed")
         except Exception:
             logger.exception("oaChatCreate: streaming failed")
 
@@ -80,10 +87,14 @@ class OpenAIServers:
         try:
             b64 = await self._img.generate_b64(data)
             resp = {"created": _now(), "data":[{"b64_json": b64}]}
-            await self._pub.publish("oa.images.reply", resp, routing_key=reply_key)
+            try:
+                await self._pub.publish("oa.images.reply", resp, routing_key=reply_key)
+            except Exception:
+                logger.exception("oaImagesGenerate: publish failed")
         except Exception:
             logger.exception("oaImagesGenerate: generation failed")
 
+
 # --- at the bottom of oa_server.py ---
 # Provide aliases expected by vllm_backend.py
 try:
@@ -103,4 +114,3 @@ except NameError:
         ImagesBackend = ImageGenerationsBackend       # noqa: F821
     except Exception:
         pass
-

rabbit_base.py

@@ -8,11 +8,10 @@ import aio_pika
 
 from config import settings
 
-ExchangeResolver = Callable[[str], str]  # exchangeName -> exchangeType
-
-import logging
+ExchangeResolver = Callable[[str], str]
 logger = logging.getLogger(__name__)
 
+
 def _normalize_exchange_type(val: str) -> aio_pika.ExchangeType:
     if isinstance(val, str):
         name = val.upper()
@@ -24,6 +23,7 @@ def _normalize_exchange_type(val: str) -> aio_pika.ExchangeType:
             pass
     return aio_pika.ExchangeType.TOPIC
 
+
 def _parse_amqp_url(url: str) -> dict:
     parts = urlsplit(url)
     return {
@@ -46,7 +46,6 @@ class RabbitBase:
             lambda _: settings.RABBIT_EXCHANGE_TYPE
         )
 
-    # -------- Status helpers --------
     def is_connected(self) -> bool:
         return bool(
             self._conn and not self._conn.is_closed and
@@ -68,15 +67,12 @@ class RabbitBase:
             self._conn = None
         logger.info("AMQP connection closed")
 
-    # -------- Core ops --------
     async def connect(self) -> None:
-        if self._conn and not self._conn.is_closed:
-            if self._chan and not self._chan.is_closed:
-                return
+        if self._conn and not self._conn.is_closed and self._chan and not self._chan.is_closed:
+            return
 
         conn_kwargs = _parse_amqp_url(str(settings.AMQP_URL))
 
-        # Log connection target (mask password)
         safe_target = {
             "scheme": conn_kwargs["scheme"],
             "host": conn_kwargs["host"],
@@ -87,7 +83,6 @@ class RabbitBase:
         }
         logger.info("AMQP connect -> %s", json.dumps(safe_target))
 
-        # TLS (intentionally disabling verification if requested)
         ssl_ctx = None
         if conn_kwargs.get("ssl"):
             ssl_ctx = ssl.create_default_context()
@@ -104,10 +99,15 @@ class RabbitBase:
                 virtualhost=conn_kwargs["virtualhost"],
                 ssl=conn_kwargs["ssl"],
                 ssl_context=ssl_ctx,
+                heartbeat=60,              # keepalive during long CPU work
+                timeout=30,
+                client_properties={"connection_name": "hf_backend_publisher"},
             )
             logger.info("AMQP connection established")
+
             self._chan = await self._conn.channel()
             logger.info("AMQP channel created")
+
             await self._chan.set_qos(prefetch_count=settings.RABBIT_PREFETCH)
             logger.info("AMQP QoS set (prefetch=%s)", settings.RABBIT_PREFETCH)
         except Exception:
@@ -117,9 +117,13 @@ class RabbitBase:
     async def ensure_exchange(self, name: str) -> aio_pika.Exchange:
         await self.connect()
         if name in self._exchanges:
-            return self._exchanges[name]
+            ex = self._exchanges[name]
+            if ex.channel and not ex.channel.is_closed:
+                return ex
+            # drop stale cache and recreate
+            self._exchanges.pop(name, None)
 
-        ex_type_str = self._exchange_type_resolver(name)  # e.g. "direct"
+        ex_type_str = self._exchange_type_resolver(name)
         ex_type = _normalize_exchange_type(ex_type_str)
 
         try:

rabbit_repo.py

@@ -1,38 +1,65 @@
 # rabbit_repo.py
 import uuid
+import asyncio
+import logging
 from typing import Any, Optional
 
+import aiormq
 import aio_pika
 
 from config import settings
 from models import CloudEvent
 from rabbit_base import RabbitBase
 from utils import to_json, json_compress_str
-import logging
+
 logger = logging.getLogger(__name__)
 
+
 class RabbitRepo(RabbitBase):
     def __init__(self, external_source: str):
         super().__init__(exchange_type_resolver=self._resolve_type)
         self._source = external_source
 
     def _resolve_type(self, exch: str) -> str:
-        # First check for oa.* exchanges
         if exch.lower().startswith("oa."):
-            return "direct"  # Default for oa.* exchanges
-    
-        # Then check EXCHANGE_TYPES if present
+            return "direct"
         if hasattr(settings, 'EXCHANGE_TYPES') and settings.EXCHANGE_TYPES:
-            matches = [k for k in settings.EXCHANGE_TYPES.keys() 
-                      if exch.lower().startswith(k.lower())]
+            matches = [k for k in settings.EXCHANGE_TYPES.keys()
+                       if exch.lower().startswith(k.lower())]
             if matches:
                 return settings.EXCHANGE_TYPES[max(matches, key=len)]
-    
-        # Default fallback
         return "fanout"
 
+    async def _publish_with_retry(self, exchange: str, body: bytes, routing_key: str = "") -> None:
+        attempts, delay = 0, 0.5
+        while True:
+            try:
+                ex = await self.ensure_exchange(exchange)
+                msg = aio_pika.Message(
+                    body=body,
+                    delivery_mode=aio_pika.DeliveryMode.PERSISTENT,
+                )
+                await ex.publish(msg, routing_key=routing_key)
+                return
+            except (asyncio.CancelledError,
+                    aiormq.exceptions.ChannelInvalidStateError,
+                    aiormq.exceptions.ConnectionClosed,
+                    aio_pika.exceptions.AMQPError,
+                    RuntimeError) as e:
+                attempts += 1
+                logger.warning("publish failed attempt=%d exchange=%s rk=%s err=%r",
+                               attempts, exchange, routing_key, e)
+                try:
+                    await self.close()
+                except Exception:
+                    pass
+                if attempts >= 5:
+                    logger.exception("publish giving up after %d attempts", attempts)
+                    raise
+                await asyncio.sleep(delay)
+                delay = min(delay * 2, 5.0)
+
     async def publish(self, exchange: str, obj: Any, routing_key: str = "") -> None:
-        ex = await self.ensure_exchange(exchange)
         payload = obj if not hasattr(obj, "model_dump") else obj.model_dump(by_alias=True)
         evt = CloudEvent.wrap(
             event_id=str(uuid.uuid4()),
@@ -41,7 +68,7 @@ class RabbitRepo(RabbitBase):
             data=payload,
         )
         body = evt.model_dump_json(exclude_none=True).encode("utf-8")
-        await ex.publish(aio_pika.Message(body=body), routing_key=routing_key)
+        await self._publish_with_retry(exchange, body, routing_key)
 
     async def publish_jsonz(
         self,
@@ -50,7 +77,6 @@ class RabbitRepo(RabbitBase):
         routing_key: str = "",
         with_id: Optional[str] = None,
     ) -> str:
-        ex = await self.ensure_exchange(exchange)
         payload = obj if not hasattr(obj, "model_dump") else obj.model_dump(by_alias=True)
         datajson = to_json(payload)
         datajsonZ = json_compress_str(datajson)
@@ -63,5 +89,5 @@ class RabbitRepo(RabbitBase):
             data=wrapped,
         )
         body = evt.model_dump_json(exclude_none=True).encode("utf-8")
-        await ex.publish(aio_pika.Message(body=body), routing_key=routing_key)
+        await self._publish_with_retry(exchange, body, routing_key)
         return datajsonZ