From 469e494f54ad2aedd3a1a13313bea7170b66d3d3 Mon Sep 17 00:00:00 2001 From: Jeffy Yu Date: Mon, 6 Jan 2025 13:55:28 -0800 Subject: [PATCH] added fetch message by receiver id --- src/message_bus.py | 20 +++++++++++++++++++- 1 file changed, 19 insertions(+), 1 deletion(-) diff --git a/src/message_bus.py b/src/message_bus.py index 57d714d7..13136d62 100644 --- a/src/message_bus.py +++ b/src/message_bus.py @@ -54,4 +54,22 @@ def fetch_all(self) -> List[MessagePayload]: return current_messages except Exception as e: logger.error(f"Error fetching messages: {e}") - return [] \ No newline at end of file + return [] + + def fetch_for_receiver(self, receiver_id: str) -> List[MessagePayload]: + """ + Retrieve and remove all messages for a given receiver_id. + If receiver_id is None, it means broadcast messages. + """ + with self.lock: + matching = [] + remaining = [] + for m in self.messages: + # Include messages where receiver_id matches or is any broadcast + if m.receiver_id is None or m.receiver_id == receiver_id: + matching.append(m) + else: + remaining.append(m) + + self.messages = remaining + return matching \ No newline at end of file