return subscription handler, add msg chunks, signing commit

Signed-off-by: nadine.loepfe <[email protected]>

Author: nadineloepfe

examples/query_topic_message.py

@@ -3,11 +3,7 @@
 from datetime import datetime
 from dotenv import load_dotenv
 
-from hiero_sdk_python import (
-    Network,
-    Client,
-    TopicMessageQuery,
-)
+from hiero_sdk_python import Network, Client, TopicMessageQuery
 
 load_dotenv()
 
@@ -20,9 +16,6 @@ def on_message_handler(topic_message):
 
     def on_error_handler(e):
         print(f"Subscription error: {e}")
-        if "Stream removed" in str(e):
-            print("Reconnecting due to stream removal...")
-            query_topic_messages()
 
     query = TopicMessageQuery(
         topic_id=os.getenv('TOPIC_ID'),
@@ -31,15 +24,21 @@ def on_error_handler(e):
         chunking_enabled=True
     )
 
-    query.subscribe(
+    handle = query.subscribe(
         client,
         on_message=on_message_handler,
         on_error=on_error_handler
     )
 
-    print("Subscription started. Waiting for messages...")
-    while True:
-        time.sleep(10)
+    print("Subscription started. Press Ctrl+C to cancel...")
+    try:
+        while True:
+            time.sleep(10)
+    except KeyboardInterrupt:
+        print("Cancelling subscription...")
+        handle.cancel() 
+        handle.join()    
+        print("Subscription cancelled. Exiting.")
 
 if __name__ == "__main__":
     query_topic_messages()

src/hiero_sdk_python/client/client.py

@@ -10,18 +10,13 @@
 
 from hiero_sdk_python.hapi.mirror import (
     consensus_service_pb2_grpc as mirror_consensus_grpc,
-    mirror_network_service_pb2_grpc as mirror_network_grpc
 )
 
 from .network import Network
 from hiero_sdk_python.response_code import ResponseCode
 from hiero_sdk_python.query.transaction_get_receipt_query import TransactionGetReceiptQuery
 from hiero_sdk_python.transaction.transaction_id import TransactionId
 
-import random
-import requests
-from hiero_sdk_python.account.account_id import AccountId
-
 Operator = namedtuple('Operator', ['account_id', 'private_key'])
 
 class Client:

src/hiero_sdk_python/client/network.py

@@ -97,7 +97,6 @@ def _fetch_nodes_from_mirror_node(self):
         url = f"{base_url}/api/v1/network/nodes?limit=100&order=desc"
 
         try:
-            import requests
             response = requests.get(url)
             response.raise_for_status()
             data = response.json()

src/hiero_sdk_python/consensus/topic_message.py

@@ -1,61 +1,160 @@
-from hiero_sdk_python.hapi.mirror import consensus_service_pb2 as mirror_proto
 from datetime import datetime
+from typing import Optional, List, Union
+from hiero_sdk_python.hapi.mirror import consensus_service_pb2 as mirror_proto
+
+def _to_datetime(ts_proto) -> datetime:
+    """
+    Convert a protobuf Timestamp to a Python datetime (UTC).
+    """
+    return datetime.utcfromtimestamp(ts_proto.seconds + ts_proto.nanos / 1e9)
+
+
+class TopicMessageChunk:
+    """
+    Represents a single chunk within a chunked topic message.
+    Mirrors the Java 'TopicMessageChunk'.
+    """
+
+    def __init__(self, response: mirror_proto.ConsensusTopicResponse):
+        self.consensus_timestamp = _to_datetime(response.consensusTimestamp)
+        self.content_size = len(response.message)
+        self.running_hash = response.runningHash
+        self.sequence_number = response.sequenceNumber
+
 
 class TopicMessage:
     """
-    Represents a single message returned from a Hedera Mirror Node subscription.
+    Represents a Hedera TopicMessage, possibly composed of multiple chunks.
     """
 
-    def __init__(self, consensus_timestamp, message, running_hash, sequence_number, 
-                 running_hash_version=None, chunk_info=None, chunks=None, transaction_id=None):
+    def __init__(
+        self,
+        consensus_timestamp: datetime,
+        contents: bytes,
+        running_hash: bytes,
+        sequence_number: int,
+        chunks: List[TopicMessageChunk],
+        transaction_id: Optional[str] = None,
+    ):
         self.consensus_timestamp = consensus_timestamp
-        self.message = message or b""
-        self.running_hash = running_hash or b""
-        self.sequence_number = sequence_number or 0
-        self.running_hash_version = running_hash_version
-        self.chunk_info = chunk_info
+        self.contents = contents
+        self.running_hash = running_hash
+        self.sequence_number = sequence_number
         self.chunks = chunks
         self.transaction_id = transaction_id
 
     @classmethod
-    def from_proto(cls, response: mirror_proto.ConsensusTopicResponse) -> "TopicMessage":
+    def of_single(cls, response: mirror_proto.ConsensusTopicResponse) -> "TopicMessage":
         """
-        Parse a Mirror Node response into a simpler object.
+        Build a TopicMessage from a single-chunk response.
         """
-        transaction_id = (
-            response.chunkInfo.initialTransactionID
-            if response.HasField("chunkInfo") and response.chunkInfo.HasField("initialTransactionID")
-            else None
+        chunk = TopicMessageChunk(response)
+        consensus_timestamp = chunk.consensus_timestamp
+        contents = response.message
+        running_hash = response.runningHash
+        sequence_number = response.sequence_number
+
+        transaction_id = None
+        if response.HasField("chunkInfo") and response.chunkInfo.HasField("initialTransactionID"):
+            tx_id = response.chunkInfo.initialTransactionID
+            transaction_id = (
+                f"{tx_id.shardNum}.{tx_id.realmNum}.{tx_id.accountNum}-"
+                f"{tx_id.transactionValidStart.seconds}.{tx_id.transactionValidStart.nanos}"
+            )
+
+        return cls(
+            consensus_timestamp,
+            contents,
+            running_hash,
+            sequence_number,
+            [chunk],
+            transaction_id
         )
+
+    @classmethod
+    def of_many(cls, responses: List[mirror_proto.ConsensusTopicResponse]) -> "TopicMessage":
+        """
+        Reassemble multiple chunk responses into a single TopicMessage.
+        """
+        sorted_responses = sorted(responses, key=lambda r: r.chunkInfo.number)
+
+        chunks = []
+        total_size = 0
+        transaction_id = None
+
+        for r in sorted_responses:
+            c = TopicMessageChunk(r)
+            chunks.append(c)
+            total_size += len(r.message)
+
+            if (transaction_id is None
+                and r.HasField("chunkInfo")
+                and r.chunkInfo.HasField("initialTransactionID")):
+                tx_id = r.chunkInfo.initialTransactionID
+                transaction_id = (
+                    f"{tx_id.shardNum}.{tx_id.realmNum}.{tx_id.accountNum}-"
+                    f"{tx_id.transactionValidStart.seconds}.{tx_id.transactionValidStart.nanos}"
+                )
+
+        contents = bytearray(total_size)
+        offset = 0
+        for r in sorted_responses:
+            end = offset + len(r.message)
+            contents[offset:end] = r.message
+            offset = end
+
+        last_r = sorted_responses[-1]
+        consensus_timestamp = _to_datetime(last_r.consensusTimestamp)
+        running_hash = last_r.runningHash
+        sequence_number = last_r.sequenceNumber
+
         return cls(
-            consensus_timestamp=response.consensusTimestamp,
-            message=response.message,
-            running_hash=response.runningHash,
-            sequence_number=response.sequenceNumber,
-            running_hash_version=response.runningHashVersion if response.runningHashVersion != 0 else None,
-            chunk_info=response.chunkInfo if response.HasField("chunkInfo") else None,
-            transaction_id=transaction_id,
+            consensus_timestamp,
+            bytes(contents),
+            running_hash,
+            sequence_number,
+            chunks,
+            transaction_id
         )
 
-    def __str__(self):
+    @classmethod
+    def from_proto(
+        cls,
+        response_or_responses: Union[mirror_proto.ConsensusTopicResponse, List[mirror_proto.ConsensusTopicResponse]],
+        chunking_enabled: bool = False
+    ) -> "TopicMessage":
         """
-        Returns a nicely formatted string representation of the topic message.
+        Creates a TopicMessage from either:
+         - A single ConsensusTopicResponse
+         - A list of responses (for multi-chunk)
+
+        If chunking is enabled and multiple chunks are detected, they are reassembled
+        into one combined TopicMessage. Otherwise, a single chunk is returned as-is.
         """
-        timestamp = datetime.utcfromtimestamp(self.consensus_timestamp.seconds).strftime('%Y-%m-%d %H:%M:%S UTC')
-        message = self.message.decode('utf-8', errors='ignore')
-        running_hash = self.running_hash.hex()
-
-        formatted_message = (
-            f"Received Topic Message:\n"
-            f"  - Timestamp: {timestamp}\n"
-            f"  - Sequence Number: {self.sequence_number}\n"
-            f"  - Message: {message}\n"
-            f"  - Running Hash: {running_hash}\n"
+        if isinstance(response_or_responses, mirror_proto.ConsensusTopicResponse):
+            response = response_or_responses
+            if chunking_enabled and response.HasField("chunkInfo") and response.chunkInfo.total > 1:
+                raise ValueError(
+                    "Cannot handle multi-chunk in a single response. Pass all chunk responses in a list."
+                )
+            return cls.of_single(response)
+        else:
+            if not response_or_responses:
+                raise ValueError("Empty response list provided to from_proto().")
+
+            if not chunking_enabled and len(response_or_responses) == 1:
+                return cls.of_single(response_or_responses[0])
+
+            return cls.of_many(response_or_responses)
+
+    def __str__(self):
+        contents_str = self.contents.decode("utf-8", errors="replace")
+        return (
+            f"TopicMessage("
+            f"consensus_timestamp={self.consensus_timestamp}, "
+            f"sequence_number={self.sequence_number}, "
+            f"contents='{contents_str[:40]}{'...' if len(contents_str) > 40 else ''}', "
+            f"chunk_count={len(self.chunks)}, "
+            f"transaction_id={self.transaction_id}"
+            f")"
         )
-        if self.running_hash_version:
-            formatted_message += f"  - Running Hash Version: {self.running_hash_version}\n"
-        if self.chunk_info:
-            formatted_message += f"  - Chunk Info: {self.chunk_info}\n"
-        if self.transaction_id:
-            formatted_message += f"  - Transaction ID: {self.transaction_id}\n"
-        return formatted_message

src/hiero_sdk_python/query/topic_message_query.py

@@ -1,17 +1,22 @@
-from datetime import datetime
-from typing import Optional, Callable, Union
-import threading
 import time
+import threading
+from datetime import datetime
+from typing import Optional, Callable, Union, Dict, List
 
-from hiero_sdk_python.consensus.topic_message import TopicMessage
 from hiero_sdk_python.hapi.mirror import consensus_service_pb2 as mirror_proto
 from hiero_sdk_python.hapi.services import basic_types_pb2, timestamp_pb2
 from hiero_sdk_python.consensus.topic_id import TopicId
+from hiero_sdk_python.consensus.topic_message import TopicMessage
+from hiero_sdk_python.utils.subscription_handle import SubscriptionHandle
+from hiero_sdk_python import Client
 
 
 class TopicMessageQuery:
     """
     A query to subscribe to messages from a specific HCS topic, via a mirror node.
+
+    If `chunking_enabled=True`, multi-chunk messages are automatically reassembled
+    before invoking `on_message`.
     """
 
     def __init__(
@@ -83,10 +88,11 @@ def set_chunking_enabled(self, enabled: bool):
 
     def subscribe(
         self,
-        client,
+        client: Client,
         on_message: Callable[[TopicMessage], None],
         on_error: Optional[Callable[[Exception], None]] = None,
-    ):
+    ) -> SubscriptionHandle:
+
         if not self._topic_id:
             raise ValueError("Topic ID must be set before subscribing.")
         if not client.mirror_stub:
@@ -100,25 +106,61 @@ def subscribe(
         if self._limit is not None:
             request.limit = self._limit
 
+        subscription_handle = SubscriptionHandle()
+
+        pending_chunks: Dict[str, List[mirror_proto.ConsensusTopicResponse]] = {}
+
         def run_stream():
             attempt = 0
-            while attempt < self._max_attempts:
+            while attempt < self._max_attempts and not subscription_handle.is_cancelled():
                 try:
                     message_stream = client.mirror_stub.subscribeTopic(request)
+
                     for response in message_stream:
-                        msg_obj = TopicMessage.from_proto(response)
-                        on_message(msg_obj)
+                        if subscription_handle.is_cancelled():
+                            return
+
+                        if (not self._chunking_enabled
+                                or not response.HasField("chunkInfo")
+                                or response.chunkInfo.total <= 1):
+                            msg_obj = TopicMessage.of_single(response)
+                            on_message(msg_obj)
+                            continue
+
+                        initial_tx_id = response.chunkInfo.initialTransactionID
+                        tx_id_str = (f"{initial_tx_id.shardNum}."
+                                     f"{initial_tx_id.realmNum}."
+                                     f"{initial_tx_id.accountNum}-"
+                                     f"{initial_tx_id.transactionValidStart.seconds}."
+                                     f"{initial_tx_id.transactionValidStart.nanos}")
+                        if tx_id_str not in pending_chunks:
+                            pending_chunks[tx_id_str] = []
+                        pending_chunks[tx_id_str].append(response)
+
+                        if len(pending_chunks[tx_id_str]) == response.chunkInfo.total:
+                            chunk_list = pending_chunks.pop(tx_id_str)
+                            msg_obj = TopicMessage.of_many(chunk_list)
+                            on_message(msg_obj)
+
                     if self._completion_handler:
                         self._completion_handler()
                     return
+
                 except Exception as e:
+                    if subscription_handle.is_cancelled():
+                        return
+
                     attempt += 1
                     if attempt >= self._max_attempts:
                         if on_error:
                             on_error(e)
                         return
+
                     delay = min(0.5 * (2 ** (attempt - 1)), self._max_backoff)
                     time.sleep(delay)
 
         thread = threading.Thread(target=run_stream, daemon=True)
+        subscription_handle.set_thread(thread)
         thread.start()
+
+        return subscription_handle