update sqs

Author: eyw520

servers/fai-lambda/fai-scribe/src/utils/sqs_client.py

@@ -13,6 +13,7 @@ def __init__(self, queue_url: str):
         self.queue_url = queue_url
         self.sqs_client = boto3.client("sqs")
         self._last_checked_count = 0
+        self._cached_messages: list[dict[str, Any]] = []
 
     def receive_messages(self, max_messages: int = 10) -> list[dict[str, Any]]:
         try:
@@ -74,19 +75,28 @@ def has_interrupt_message(self) -> tuple[bool, str | None]:
             if body.get("type") == "INTERRUPT":
                 LOGGER.info(f"Found INTERRUPT message in queue")
                 return True, msg["receipt_handle"]
+            else:
+                self._cached_messages.append(msg)
 
         return False, None
 
     def get_resume_messages(self) -> list[dict[str, Any]]:
-        messages = self.receive_messages(max_messages=10)
         resume_messages = []
 
+        for msg in self._cached_messages:
+            body = msg["body"]
+            if body.get("type") == "RESUME":
+                resume_messages.append({"body": body, "receipt_handle": msg["receipt_handle"]})
+
+        messages = self.receive_messages(max_messages=10)
         for msg in messages:
             body = msg["body"]
             if body.get("type") == "RESUME":
                 resume_messages.append({"body": body, "receipt_handle": msg["receipt_handle"]})
 
         if resume_messages:
-            LOGGER.info(f"Found {len(resume_messages)} RESUME messages in queue")
+            LOGGER.info(f"Found {len(resume_messages)} RESUME messages ({len([m for m in self._cached_messages if m['body'].get('type') == 'RESUME'])} cached, {len([m for m in messages if m['body'].get('type') == 'RESUME'])} new)")
+
+        self._cached_messages.clear()
 
         return resume_messages