diff --git a/src/judas_protocol/message.py b/src/judas_protocol/message.py index 8167e00..42591bd 100644 --- a/src/judas_protocol/message.py +++ b/src/judas_protocol/message.py @@ -18,6 +18,9 @@ class Message: category: Category, action: Enum, payload: dict[str, Any] | None = None, + *, + acknowledged: bool = False, + ack_required: bool = False, ) -> None: """Initialize a Message. @@ -26,19 +29,28 @@ class Message: category (Category): The category of the message. action (Enum): The action of the message. payload (dict[str, Any] | None): The payload of the message. + acknowledged (bool): Whether the message has been acknowledged. Defaults to False. + ack_required (bool): Whether the message requires an acknowledgment. Defaults to False. """ self.id: str = id_ or str(uuid.uuid4()) self.category: Category = category self.action: Enum = action self.payload: dict[str, Any] = payload or {} + self.acknowledged: bool = acknowledged + self.ack_required: bool = ack_required + @override def __str__(self) -> str: - return f"Message ({self.id}) [{self.category.name}] {self.action.name} - {self.payload}" + return f"Message ({self.id}) ack'd: {self.acknowledged}, req: {self.ack_required} [{self.category.name}] {self.action.name} - {self.payload}" @override def __repr__(self) -> str: - return f"Message(category={self.category}, action={self.action}, payload={self.payload})" + return ( + f"Message(id_={self.id}, category={self.category}, action={self.action}" + f", payload={self.payload}, acknowledged={self.acknowledged}" + f", ack_required={self.ack_required})" + ) def to_dict(self) -> dict[str, Any]: """Convert the message to a dictionary. @@ -51,6 +63,8 @@ class Message: "category": self.category, "action": self.action, "payload": self.payload, + "acknowledged": self.acknowledged, + "ack_required": self.ack_required, } def to_json(self) -> str: @@ -85,6 +99,8 @@ class Message: id_: str | None = data["id"] category = Category(data["category"]) action_str = data["action"] + acknowledged = data["acknowledged"] + ack_required = data["ack_required"] match category: case Category.CONTROL: @@ -93,7 +109,14 @@ class Message: raise ValueError(f"Unknown category: {category}") payload = data.get("payload", {}) - return cls(id_=id_, category=category, action=action, payload=payload) + return cls( + id_=id_, + category=category, + action=action, + payload=payload, + acknowledged=acknowledged, + ack_required=ack_required, + ) @classmethod def from_json(cls, data: str) -> Message: @@ -137,6 +160,8 @@ class Message: category=Category.CONTROL, action=ControlAction.ACK, payload={"target_id": target_id}, + acknowledged=False, + ack_required=False, ) @classmethod @@ -153,4 +178,6 @@ class Message: category=Category.CONTROL, action=ControlAction.HELLO, payload={"mac": mac}, + acknowledged=False, + ack_required=True, )