From 37eec56afeb2697955702aba04476ccdde7435e4 Mon Sep 17 00:00:00 2001
From: Artur Borecki <me@pufereq.pl>
Date: Sun, 5 Oct 2025 13:34:28 +0200
Subject: [PATCH] fix(message.py): allow providing ID via argument

---
 src/judas_protocol/message.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/src/judas_protocol/message.py b/src/judas_protocol/message.py
index dc221e5..8548f59 100644
--- a/src/judas_protocol/message.py
+++ b/src/judas_protocol/message.py
@@ -14,6 +14,7 @@ class Message:
 
     def __init__(
         self,
+        id_: str | None,
         category: Category,
         action: Enum,
         payload: dict[str, Any] | None = None,
@@ -21,11 +22,12 @@ class Message:
         """Initialize a Message.
 
         Args:
+            id_ (str | None): The ID of the message. If None, a new UUID will be generated.
             category (Category): The category of the message.
             action (Enum): The action of the message.
             payload (dict[str, Any] | None): The payload of the message.
         """
-        self.id: str = str(uuid.uuid4())
+        self.id: str = id_ or str(uuid.uuid4())
         self.category: Category = category
         self.action: Enum = action
         self.payload: dict[str, Any] = payload or {}
@@ -77,6 +79,10 @@ class Message:
         Returns:
             Message: The created message.
         """
+        if "id" not in data:
+            data["id"] = None
+
+        id_: str | None = data["id"]
         category = Category(data["category"])
         action_str = data["action"]
 
@@ -87,7 +93,7 @@ class Message:
                 raise ValueError(f"Unknown category: {category}")
 
         payload = data.get("payload", {})
-        return cls(category=category, action=action, payload=payload)
+        return cls(id_=id_, category=category, action=action, payload=payload)
 
     @classmethod
     def from_json(cls, data: str) -> Message: