feat: Implement stateful retry and resumable stream logic for TopicMessageQuery

[manishdait]

Description:
This PR introduces fixes for flaky TopicMessageQuery e2e tests by implementing retry logic and error handling.

Changes Made:

• Added SubscriptionState to track last_message and count. on retry, the query resumes from last_message.consensusTimestamp + 1ns to prevent message loss or duplication.
• Added _should_retry logic to identify retryable gRPC errors (e.g., UNAVAILABLE, RESOURCE_EXHAUSTED) and specific RST_STREAM internal errors.
• Updated SubscriptionHandle to safely manage gRPC call cancellation across threads.

Related issue(s):

Fixes #1796

Notes for reviewer:

Checklist

• Documented (Code comments, README, etc.)
• Tested (unit, integration, etc.)

[codecov]

Codecov Report

❌ Patch coverage is 99.11504% with 1 line in your changes missing coverage. Please review.

@@            Coverage Diff             @@
##             main    #2171      +/-   ##
==========================================
+ Coverage   93.99%   95.04%   +1.04%     
==========================================
  Files         163      163              
  Lines       10445    10514      +69     
==========================================
+ Hits         9818     9993     +175     
+ Misses        627      521     -106     

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[codacy-production]

Up to standards ✅

🟢 Issues 0 issues

Results:
0 new issues

View in Codacy

_{NEW Get contextual insights on your PRs based on Codacy's metrics, along with PR and Jira context, without leaving GitHub. Enable AI reviewer
TIP This summary will be updated as you push new changes.}

[github-actions]

Hi @manishdait,

This pull request has had no commit activity for 10 days. Are you still working on it?
To keep the PR active, you can:

• Push a new commit.
• Comment /working on the linked issue (not this PR).

If you're no longer working on this, please comment /unassign on the linked issue to release it for others. Otherwise, this PR may be closed due to inactivity.

Reach out on discord or join our office hours if you need assistance.

• Discord
• Office Hours

From the Python SDK Team

[github-actions]

Hi @manishdait, this is CronInactivityBot 👋

This pull request has had no new commits for 21 days, so I'm closing it and unassigning you from the linked issue to keep the backlog healthy.

If you're no longer interested, no action is needed.

Tip: You can comment /unassign on any issue to proactively step away before this bot kicks in.

If you'd like to continue working on this later, feel free to comment /assign on the issue to get re-assigned, and open a new PR when you're ready. 🚀

[coderabbitai]

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

Walkthrough

TopicMessageQuery subscription now tracks state across retries, resumes from the last received message, handles chunked responses, applies gRPC-aware retry decisions with exponential backoff, and coordinates thread-safe call cancellation via SubscriptionHandle.

Changes

Topic Message Query Subscription Reliability

Layer / File(s)	Summary
Subscription state tracking and constants src/hiero_sdk_python/query/topic_message_query.py	Added gRPC/logging imports, module logger, RST_STREAM regex, and SubscriptionState dataclass tracking retry attempts, message counts, last_message resume point, and buffered chunks.
Query initialization and error handler API src/hiero_sdk_python/query/topic_message_query.py	Refactored __init__ to parse/store TopicId/Timestamp, default _limit handling, tightened validation for max_attempts/max_backoff, initialized internal handlers, and added public set_error_handler().
Subscription lifecycle helpers: request building, response handling, and retry logic src/hiero_sdk_python/query/topic_message_query.py	Added _should_retry() (gRPC codes + RST_STREAM detail matching), _build_query_request() (resume from last_message, adjust limit), and _handle_response() (update counters, last_message, buffer/reassemble chunked responses keyed by initialTransactionID).
Subscription streaming loop and centralized error retry src/hiero_sdk_python/query/topic_message_query.py	Refactored subscribe() to drive the stream with SubscriptionState, rebuild requests per attempt, process responses via _handle_response, enforce cancellation checks, and apply capped exponential backoff with attempt budget enforcement and conditional user error handler invocation.
Thread-safe gRPC call cancellation in SubscriptionHandle src/hiero_sdk_python/utils/subscription_handle.py	Added _call storage and _lock, _set_call(call: Any) to register and immediately cancel if already signalled, and updated cancel() to cancel stored calls under the mutex.
SubscriptionHandle test coverage for call cancellation tests/unit/subscription_handle_test.py	Added/annotated tests for default non-cancelled state, cancel() marking cancelled, join(timeout=...) forwarding to thread join and no-op join without thread, plus tests verifying active-call cancellation and immediate cancellation of calls set after cancellation.
TopicMessageQuery test expansion: initialization, chunk buffering, retry, and cancellation tests/unit/topic_message_query_test.py	Updated imports/fixtures and payload; added initialization and setter validation tests; added missing-config precondition test; added chunk reassembly test; added retryable and non-retryable error tests asserting retry/stop behavior and on_error calls; added subscription cancellation test asserting underlying RPC cancel() is invoked.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title accurately describes the main change: implementing stateful retry and resumable stream logic for TopicMessageQuery, which aligns with the refactoring and new features in the changeset.
Description check	✅ Passed	The description is related to the changeset, explaining the fixes for flaky tests, SubscriptionState tracking, _should_retry logic, and SubscriptionHandle updates that align with the code changes.
Linked Issues check	✅ Passed	The PR implements the core requirements from issue #1796: fixes intermittent message delivery failures through stateful retry logic, resumable streams from last_message.consensusTimestamp + 1ns, retryable gRPC error detection, thread-safe call cancellation, and comprehensive unit tests covering retry/chunking/cancellation scenarios.
Out of Scope Changes check	✅ Passed	All changes are scoped to the linked objectives: TopicMessageQuery retry/resumption logic, SubscriptionState/error handler additions, SubscriptionHandle thread-safe call management, and corresponding unit tests directly addressing the flaky integration test issues.
Docstring Coverage	✅ Passed	Docstring coverage is 84.00% which is sufficient. The required threshold is 80.00%.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

📋 Issue Planner

Let us write the prompt for your AI agent so you can ship faster (with fewer bugs).

View plan for ticket: #1796

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 6

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Pro

Run ID: ea309a58-41ae-4f9a-9ef9-8bb283f770f5

📥 Commits

Reviewing files that changed from the base of the PR and between bc1e285 and 631d51e.

📒 Files selected for processing (4)

• src/hiero_sdk_python/query/topic_message_query.py
• src/hiero_sdk_python/utils/subscription_handle.py
• tests/unit/subscription_handle_test.py
• tests/unit/topic_message_query_test.py

[coderabbitai]

Actionable comments posted: 4

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Pro

Run ID: 1454aa79-c860-4bc3-8f79-1b3fde0bbdc6

📥 Commits

Reviewing files that changed from the base of the PR and between 631d51e and 43f9b2c.

📒 Files selected for processing (4)

• src/hiero_sdk_python/query/topic_message_query.py
• src/hiero_sdk_python/utils/subscription_handle.py
• tests/unit/subscription_handle_test.py
• tests/unit/topic_message_query_test.py

[coderabbitai]

Actionable comments posted: 1

♻️ Duplicate comments (4)

src/hiero_sdk_python/query/topic_message_query.py (1)

61-62: ⚠️ Potential issue | 🟡 Minor | ⚡ Quick win

Type annotation mismatch: _error_handler declared as Callable[[], None] but used with Exception argument.

Line 62 declares _error_handler as Callable[[], None], but:

• set_error_handler (line 85) expects Callable[[Exception], None]
• _on_error (line 118) takes err: Exception
• The handler is invoked with self._error_handler(e) at line 243

This inconsistency means type checkers will report errors, and users following the type hint at line 62 would provide incompatible handlers.

🐛 Proposed fix

         self._completion_handler: Callable[[], None] | None = self._on_complete
-        self._error_handler: Callable[[], None] | None = self._on_error
+        self._error_handler: Callable[[Exception], None] | None = self._on_error

tests/unit/topic_message_query_test.py (3)

152-154: ⚠️ Potential issue | 🟡 Minor | ⚡ Quick win

Chunk ordering is not verified — test would pass even if chunks are assembled in wrong order.

Using in only checks presence, not position. If the implementation incorrectly orders chunks, this test won't detect it.

🐛 Proposed fix to verify correct ordering

     assert len(received_messages) == 1
-    assert b"chunk-1" in received_messages[0].contents
-    assert b"chunk-2" in received_messages[0].contents
+    # Verify chunks are assembled in correct order
+    contents = received_messages[0].contents
+    chunk1_pos = contents.find(b"chunk-1")
+    chunk2_pos = contents.find(b"chunk-2")
+    assert chunk1_pos != -1, "chunk-1 not found in assembled message"
+    assert chunk2_pos != -1, "chunk-2 not found in assembled message"
+    assert chunk1_pos < chunk2_pos, "Chunks assembled in wrong order: chunk-1 should come before chunk-2"

As per coding guidelines, tests should catch regressions in chunk ordering behavior.

---

167-178: ⚠️ Potential issue | 🟡 Minor | ⚡ Quick win

Missing assertion that on_error is invoked after retries exhaust.

When max_attempts are exhausted with retryable errors, on_error should be called. The test creates an on_error=MagicMock() but never verifies it was invoked, so regressions in terminal error handling would go undetected.

🐛 Proposed fix

 def test_retry_logic_on_retryable_error(mock_client, error):
     """Test that the query retries on retryable errors but stops after max_attempts."""
     query = TopicMessageQuery(topic_id="0.0.123").set_max_attempts(2).set_max_backoff(0.5)

     mock_client.mirror_stub.subscribeTopic.side_effect = [error, error]

-    handle = query.subscribe(mock_client, on_message=MagicMock(), on_error=MagicMock())
+    on_error = MagicMock()
+    handle = query.subscribe(mock_client, on_message=MagicMock(), on_error=on_error)

     handle._thread.join(timeout=2.0)
+    assert not handle._thread.is_alive(), "Thread should have terminated after retries exhausted"

     assert mock_client.mirror_stub.subscribeTopic.call_count == 2
+    on_error.assert_called_once_with(error)

As per coding guidelines, "Tests must provide useful error messages when they fail for future debugging."

---

220-224: ⚠️ Potential issue | 🟡 Minor | ⚡ Quick win

Timing-dependent sleep can cause flaky tests in CI.

time.sleep(0.2) followed by assert handle._thread.is_alive() is scheduler-dependent. In slow CI environments, the thread may not have started processing yet, or conversely, the sleep may be insufficient. Use an event to synchronize on first message receipt before cancellation.

🐛 Proposed fix using event synchronization

+import threading
+
 def test_subscription_cancellation(mock_client):
     """Test that cancelling a handle stops the subscription thread."""
     query = TopicMessageQuery(topic_id="0.0.123")

     def infinite_stream():
         while True:
             yield mirror_proto.ConsensusTopicResponse(message=b"ping")
             time.sleep(0.1)

     mock_call = MagicMock()
     mock_call.__iter__.return_value = infinite_stream()

     mock_client.mirror_stub.subscribeTopic.return_value = mock_call

-    on_message = MagicMock()
+    first_message_seen = threading.Event()
+    on_message = MagicMock(side_effect=lambda _: first_message_seen.set())
     handle = query.subscribe(mock_client, on_message=on_message)

-    time.sleep(0.2)
-    assert handle._thread.is_alive()
+    assert first_message_seen.wait(timeout=2.0), "Expected at least one streamed message before cancellation"
+    assert handle._thread.is_alive(), "Subscription thread should be alive before cancel"

     handle.cancel()

As per coding guidelines, "No timing-dependent or unseeded random assertions."

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Pro

Run ID: cc2d52be-f082-42a4-b517-4b8426de7c85

📥 Commits

Reviewing files that changed from the base of the PR and between 43f9b2c and 45c80fb.

📒 Files selected for processing (4)

• src/hiero_sdk_python/query/topic_message_query.py
• src/hiero_sdk_python/utils/subscription_handle.py
• tests/unit/subscription_handle_test.py
• tests/unit/topic_message_query_test.py

[github-actions]

Hi, this is WorkflowBot.
Your pull request cannot be merged as it is not passing all our workflow checks.
Please click on each check to review the logs and resolve issues so all checks pass.
To help you:

• DCO signing guide
• Merge conflicts guide
• Rebase guide
• Testing guide
• Discord
• Community Calls
  Thank you for contributing!
  From the Hiero Python SDK Team

[coderabbitai]

Actionable comments posted: 2

♻️ Duplicate comments (4)

tests/unit/topic_message_query_test.py (4)

383-408: ⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Replace timing-dependent sleep with event-based synchronization.

The time.sleep(0.2) at line 400 followed by is_alive() assertion at line 401 is scheduler-dependent and can flap in CI under high load or slow schedulers. Use an event to deterministically wait for the first message delivery before asserting thread liveness and performing cancellation.

⏱️ Proposed fix using threading.Event

+import threading
 import time
 from datetime import datetime, timezone
@@
 def test_subscription_cancellation(mock_client):
     """Test that cancelling a handle stops the subscription thread."""
     query = TopicMessageQuery(topic_id="0.0.123")

     def infinite_stream():
         while True:
             yield mirror_proto.ConsensusTopicResponse(message=b"ping")
             time.sleep(0.1)

     mock_call = MagicMock()
     mock_call.__iter__.return_value = infinite_stream()

     mock_client.mirror_stub.subscribeTopic.return_value = mock_call

-    on_message = MagicMock()
+    first_message_seen = threading.Event()
+    on_message = MagicMock(side_effect=lambda _: first_message_seen.set())
     handle = query.subscribe(mock_client, on_message=on_message)

-    time.sleep(0.2)
-    assert handle._thread.is_alive()
+    assert first_message_seen.wait(timeout=1.0), "Expected at least one message before cancellation"
+    assert handle._thread.is_alive(), "Subscription thread should be alive before cancel()"

     handle.cancel()

     handle._thread.join(timeout=1.0)

     assert not handle._thread.is_alive()
     mock_call.cancel.assert_called()

As per coding guidelines, "No timing-dependent or unseeded random assertions."

---

54-80: ⚠️ Potential issue | 🟡 Minor | ⚡ Quick win

Verify fluent interface contract: setters must return self.

The test validates internal state but doesn't assert that each setter returns self, which is required for method chaining. If a setter is accidentally changed to return None, chaining would break but this test would pass.

🔗 Proposed fix

 def test_topic_message_query_initialization():
     """Test initializing the query with various parameter types and setters."""
     start = datetime(2023, 1, 1, tzinfo=timezone.utc)

     def mock_complete():
         pass

     def mock_error(e):
         pass

-    query = (
-        TopicMessageQuery()
-        .set_topic_id("0.0.123")
-        .set_start_time(start)
-        .set_limit(5)
-        .set_chunking_enabled(True)
-        .set_completion_handler(mock_complete)
-        .set_error_handler(mock_error)
-    )
+    query = TopicMessageQuery()
+    assert query.set_topic_id("0.0.123") is query, "set_topic_id should return self"
+    assert query.set_start_time(start) is query, "set_start_time should return self"
+    assert query.set_limit(5) is query, "set_limit should return self"
+    assert query.set_chunking_enabled(True) is query, "set_chunking_enabled should return self"
+    assert query.set_completion_handler(mock_complete) is query, "set_completion_handler should return self"
+    assert query.set_error_handler(mock_error) is query, "set_error_handler should return self"

     assert query._topic_id.topicNum == 123
     assert query._start_time.seconds == int(start.timestamp())
     assert query._limit == 5
     assert query._chunking_enabled is True
     assert query._completion_handler == mock_complete
     assert query._error_handler == mock_error

As per coding guidelines, "Assert fluent setters return self".

---

304-306: ⚠️ Potential issue | 🟡 Minor | ⚡ Quick win

Chunk ordering is not verified — test passes even if chunks are assembled backwards.

The test uses in at lines 305-306, which only checks presence. If the implementation concatenates chunk-2 before chunk-1, this assertion would still pass.

🔍 Proposed fix to verify correct ordering

     assert len(received_messages) == 1
-    assert b"chunk-1" in received_messages[0].contents
-    assert b"chunk-2" in received_messages[0].contents
+    # Verify chunks are assembled in correct order
+    contents = received_messages[0].contents
+    chunk1_pos = contents.find(b"chunk-1")
+    chunk2_pos = contents.find(b"chunk-2")
+    assert chunk1_pos != -1, "chunk-1 not found in assembled message"
+    assert chunk2_pos != -1, "chunk-2 not found in assembled message"
+    assert chunk1_pos < chunk2_pos, "Chunks assembled in wrong order: chunk-1 must precede chunk-2"

As per coding guidelines, tests must catch regressions in ordering behavior.

---

345-355: ⚠️ Potential issue | 🟡 Minor | ⚡ Quick win

Verify terminal error callback and thread termination after retries exhaust.

The test passes on_error=MagicMock() at line 350 but never asserts it's invoked when max_attempts are exhausted. The implementation should call on_error with the final error before stopping. Additionally, join(timeout=2.0) at line 352 should be followed by an assertion that the thread actually terminated—if it times out, the subscription count assertion at line 354 may pass spuriously.

🧪 Proposed fix

 def test_retry_logic_on_retryable_error(mock_client, error):
     """Test that the query retries on retryable errors but stops after max_attempts."""
     query = TopicMessageQuery(topic_id="0.0.123").set_max_attempts(2).set_max_backoff(0.5)

     mock_client.mirror_stub.subscribeTopic.side_effect = [error, error]
-    handle = query.subscribe(mock_client, on_message=MagicMock(), on_error=MagicMock())
+    on_error = MagicMock()
+    handle = query.subscribe(mock_client, on_message=MagicMock(), on_error=on_error)

     handle._thread.join(timeout=2.0)
+    assert not handle._thread.is_alive(), "Subscription thread should have terminated after retries exhausted"

     assert mock_client.mirror_stub.subscribeTopic.call_count == 2
+    on_error.assert_called_once_with(error)

As per coding guidelines, "Tests must provide useful error messages when they fail for future debugging" and should verify error handler contracts.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Pro

Run ID: e10049ac-fd87-4829-a382-583a10b742f7

📥 Commits

Reviewing files that changed from the base of the PR and between 45c80fb and bb49b14.

📒 Files selected for processing (4)

• src/hiero_sdk_python/query/topic_message_query.py
• src/hiero_sdk_python/utils/subscription_handle.py
• tests/unit/subscription_handle_test.py
• tests/unit/topic_message_query_test.py

[coderabbitai]

Actionable comments posted: 2

♻️ Duplicate comments (3)

src/hiero_sdk_python/query/topic_message_query.py (1)

176-177: ⚠️ Potential issue | 🟠 Major

Don't let a resumed query fall back to unlimited streaming.

Proto field: ConsensusTopicQuery.limit (#4). Issue type: Wrong default. Once state.count reaches the caller's cap, the retry path sends limit = 0, but the mirror schema defines zero/unset as "return messages indefinitely", so a reconnect after the cap is reached can over-deliver instead of stopping. Stop locally once state.count >= self._limit, or avoid emitting a retry request with limit=0. (github.com)
As per coding guidelines, "Compare the SDK class against the proto schema."

Also applies to: 227-259

src/hiero_sdk_python/utils/subscription_handle.py (1)

20-26: ⚠️ Potential issue | 🟠 Major

Cancel the active call after releasing _lock.

Both _set_call() and cancel() invoke call.cancel() inside the critical section. That makes the lock cover an external call and can stall the other thread trying to register or cancel the stream. Capture the call reference while locked, then invoke cancel() once the mutex is released.

Also applies to: 28-34

tests/unit/topic_message_query_test.py (1)

397-405: ⚠️ Potential issue | 🟡 Minor

Make this cancellation test event-driven.

time.sleep(0.2) plus an immediate liveness check is scheduler-dependent and can flap in CI. Wait for the first streamed message or for subscribeTopic() to be entered before asserting the thread is alive and calling cancel(). As per coding guidelines, "No timing-dependent or unseeded random assertions."

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Pro

Run ID: eeaf40a6-b73f-4be7-8c19-c131fa2cfddf

📥 Commits

Reviewing files that changed from the base of the PR and between bb49b14 and 7b8a3bb.

📒 Files selected for processing (4)

• src/hiero_sdk_python/query/topic_message_query.py
• src/hiero_sdk_python/utils/subscription_handle.py
• tests/unit/subscription_handle_test.py
• tests/unit/topic_message_query_test.py

[aceppaluni]

LGTM, just make sure to resolve the open conversations

[github-actions]

👋 Hi @manishdait! Great work completing a Advanced issue! 🎉

Thanks for your contribution! 🚀

Here are some issues you might want to explore next:

• [Advanced]: Integrate Code Coverage for C++ SDK
• [Advanced]: Upgrade gRPC and protobuf dependencies

🌟 Stay connected:

• ⭐ Star this repository
• 👀 Watch for new issues
• 💬 Join us on Discord

Happy coding! 🚀
— Hiero Python SDK Team