refactor(python): optimize consecutive Session.call() invocations#6563
refactor(python): optimize consecutive Session.call() invocations#6563romanz wants to merge 3 commits into
Session.call() invocations#6563Conversation
It allows scoping "interactive" context, where THP ACK piggybacking will be enabled. [no changelog]
WalkthroughA new generic InteractionContext context manager is introduced to represent an active client-session interaction and delegate enter/exit and call behavior to the underlying client. Session now takes explicit client and id parameters, exposes interact() returning an InteractionContext, and its call() uses the interaction context. TrezorClient gains an interact(session) factory producing InteractionContext instances. The THP layer adds ThpInteractionContext that acquires/releases piggyback_acks in its enter/exit, and THP client/session code was updated to use the new interact-based workflow instead of the previous _call piggybacking approach. Sequence Diagram(s)sequenceDiagram
participant User
participant Session
participant InteractionContext
participant TrezorClient
User->>Session: interact()
Session->>TrezorClient: request InteractionContext
TrezorClient-->>InteractionContext: create bound context
InteractionContext-->>User: return context
User->>InteractionContext: __enter__()
InteractionContext->>TrezorClient: __enter__()
TrezorClient-->>InteractionContext: ready
InteractionContext-->>User: context ready
User->>InteractionContext: call(method, ...)
InteractionContext->>TrezorClient: _call(session, method, ...)
TrezorClient-->>InteractionContext: result
InteractionContext-->>User: return result
User->>InteractionContext: __exit__()
InteractionContext->>TrezorClient: __exit__()
TrezorClient-->>InteractionContext: cleanup
InteractionContext-->>User: closed
sequenceDiagram
participant User
participant ThpSession
participant ThpInteractionContext
participant ThpClient
participant piggyback_acks
User->>ThpSession: interact()
ThpSession->>ThpClient: request ThpInteractionContext
ThpClient-->>ThpInteractionContext: create bound context
ThpInteractionContext-->>User: return context
User->>ThpInteractionContext: __enter__()
ThpInteractionContext->>piggyback_acks: acquire ownership
ThpInteractionContext->>ThpClient: __enter__()
ThpClient-->>ThpInteractionContext: ready
ThpInteractionContext-->>User: context ready
User->>ThpInteractionContext: call(method, ...)
ThpInteractionContext->>ThpClient: _call(session, method, ...)
ThpClient-->>ThpInteractionContext: result
ThpInteractionContext-->>User: return result
User->>ThpInteractionContext: __exit__()
ThpInteractionContext->>ThpClient: __exit__()
ThpClient-->>ThpInteractionContext: cleanup
ThpInteractionContext->>piggyback_acks: release ownership
ThpInteractionContext-->>User: closed
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes 🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches
🧪 Generate unit tests (beta)
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. Comment |
|
| model | device_test | click_test | persistence_test |
|---|---|---|---|
| T2T1 |  test(all)  main(all)  |
 test(all)  main(all)  |
 test(all)  main(all)  |
| T3B1 |  test(all)  main(all)  |
 test(all)  main(all)  |
 test(all)  main(all)  |
| T3T1 |  test(all)  main(all)  |
 test(all)  main(all)  |
 test(all)  main(all)  |
| T3W1 |  test(all)  main(all)  |
 test(all)  main(all)  |
 test(all)  main(all)  |
Latest CI run: 22776484968
|
that's what |
I am not sure... IIUC,
Using this feature, the host will not send THP ACKs when Current THP implementation may require host-side retransmission in the following case. |
|
Maybe I should rename |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@python/src/trezorlib/client.py`:
- Around line 137-149: Session.call currently always opens a new context with
self.interact(), which breaks THP single-workflow constraints; change
Session.call so it first checks for and reuses an already-active interaction
context (e.g. an attribute like self._interaction_ctx or whatever the client
stores for an active InteractionContext) and if present call ctx.call(msg,
expect=expect, timeout=timeout) on that context, otherwise fall back to the
existing with self.interact() as ctx: return ctx.call(...) flow; keep existing
error handling (is_invalid / exceptions.InvalidSessionError) and use the
InteractionContext.call method when reusing the context.
In `@python/src/trezorlib/thp/client.py`:
- Around line 41-55: ThpInteractionContext.__enter__ currently calls
super().__enter__() then piggyback_acks.__enter__(), which leaks the transport
if piggyback setup raises; similarly __exit__ may skip super().__exit__ if
piggyback_acks.__exit__ raises. Wrap the piggyback_acks setup in a
try/except/finally in __enter__ so that if self.piggyback_acks.__enter__() fails
you call super().__exit__() (or otherwise close/unwind the client) before
re-raising, and in __exit__ call piggyback_acks.__exit__ inside a try/except (or
try/finally) that ensures super().__exit__(...) is always invoked afterward;
reference ThpInteractionContext, __enter__, __exit__, self.piggyback_acks,
client.channel.piggyback_acks, and super().__enter__/super().__exit__ when
making the changes.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: c328cea1-ff36-401b-be4c-4bd874d0e813
📒 Files selected for processing (2)
python/src/trezorlib/client.pypython/src/trezorlib/thp/client.py
| """ | ||
| Call a method on this session, process and return the response. | ||
|
|
||
| Use `self.interact()` for consecutive calls (to allow THP ACK piggybacking). | ||
| """ | ||
| with self.interact() as ctx: | ||
| return ctx.call(msg, expect=expect, timeout=timeout) | ||
|
|
||
| def interact(self) -> InteractionContext: | ||
| """Use the returned context manager to call methods on this session.""" | ||
| if self.is_invalid: | ||
| raise exceptions.InvalidSessionError(self.id) | ||
| with self: | ||
| return self.client._call(self, msg, expect=expect, timeout=timeout) | ||
| return self.client.interact(self) |
There was a problem hiding this comment.
Make Session.call() reuse an active interaction
call() always opens a fresh self.interact(). On THP, python/src/trezorlib/thp/channel.py:445-447 asserts that only one workflow is active, so even the straightforward pattern with session.interact(): session.call(...) fails on the first call. That also prevents helpers that still use session.call() from participating in a caller-managed interactive scope. Please reuse the current interaction here instead of unconditionally nesting a new one.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@python/src/trezorlib/client.py` around lines 137 - 149, Session.call
currently always opens a new context with self.interact(), which breaks THP
single-workflow constraints; change Session.call so it first checks for and
reuses an already-active interaction context (e.g. an attribute like
self._interaction_ctx or whatever the client stores for an active
InteractionContext) and if present call ctx.call(msg, expect=expect,
timeout=timeout) on that context, otherwise fall back to the existing with
self.interact() as ctx: return ctx.call(...) flow; keep existing error handling
(is_invalid / exceptions.InvalidSessionError) and use the
InteractionContext.call method when reusing the context.
| class ThpInteractionContext(client.InteractionContext["TrezorClientThp", "ThpSession"]): | ||
| def __enter__(self) -> tx.Self: | ||
| super().__enter__() | ||
| self.piggyback_acks = self.client.channel.piggyback_acks(self) | ||
| self.piggyback_acks.__enter__() | ||
| return self | ||
|
|
||
| def __exit__( | ||
| self, | ||
| exc_type: type[BaseException] | None, | ||
| exc_value: BaseException | None, | ||
| traceback: t.Any, | ||
| ) -> None: | ||
| self.piggyback_acks.__exit__(exc_type, exc_value, traceback) | ||
| return super().__exit__(exc_type, exc_value, traceback) |
There was a problem hiding this comment.
Make THP interaction cleanup exception-safe
super().__enter__() opens the transport before channel.piggyback_acks(self).__enter__() runs. If piggyback setup fails — for example, a nested workflow trips the assert in python/src/trezorlib/thp/channel.py:445-447 — this context never reaches __exit__() and the transport stays open. The same leak exists on the way out: if piggyback_acks.__exit__() raises while sending the final ACK (python/src/trezorlib/thp/channel.py:454-460), super().__exit__() is skipped. Please unwind the client context in an except/finally path here.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@python/src/trezorlib/thp/client.py` around lines 41 - 55,
ThpInteractionContext.__enter__ currently calls super().__enter__() then
piggyback_acks.__enter__(), which leaks the transport if piggyback setup raises;
similarly __exit__ may skip super().__exit__ if piggyback_acks.__exit__ raises.
Wrap the piggyback_acks setup in a try/except/finally in __enter__ so that if
self.piggyback_acks.__enter__() fails you call super().__exit__() (or otherwise
close/unwind the client) before re-raising, and in __exit__ call
piggyback_acks.__exit__ inside a try/except (or try/finally) that ensures
super().__exit__(...) is always invoked afterward; reference
ThpInteractionContext, __enter__, __exit__, self.piggyback_acks,
client.channel.piggyback_acks, and super().__enter__/super().__exit__ when
making the changes.
`session.call_raw()` doesn't send a THP ACK when it is done, so the device will re-transmit the error message over and over. [no changelog]
Also rename ACK-sending methods, since "flushing" ACKs should not be skipped (otherwise the device may get stuck waiting for the ACK, while the host is waiting on DebugLink). [no changelog]
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@python/src/trezorlib/thp/channel.py`:
- Around line 455-463: The ACK is being sent on every piggyback_acks() exit;
modify the guard so the final ACK is only sent when leaving the outermost
workflow: after restoring self._active_workflow = previous, only call
self._send_ack_internal(...) if previous is None (i.e. there is no enclosing
workflow) and self.is_ack_piggybacking_allowed is true. Update the condition
that currently checks is_ack_piggybacking_allowed to also require previous is
None (or equivalently self._active_workflow is None) before sending the ACK,
referencing the existing variables/methods _active_workflow, previous,
is_ack_piggybacking_allowed, and _send_ack_internal.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 04aa7a7b-3fb7-49d8-916e-a1d54e2d8385
📒 Files selected for processing (2)
python/src/trezorlib/thp/channel.pytests/device_tests/test_msg_backup_device.py
| active = self._active_workflow | ||
| self._active_workflow = None | ||
| self._active_workflow = previous | ||
| # Verify proper nesting of workflows | ||
| assert active is marker | ||
| if self.is_ack_piggybacking_allowed: | ||
| # Explicitly ACK the latest received message. The device may restart | ||
| # the event loop, so the next request will be sent in a separate message. | ||
| self._send_ack(None) | ||
| # Explicitly ACK the last message received. | ||
| # Otherwise, the next message sent may be lost due to an event loop restart. | ||
| ack = control_byte.make_ack(not self.sync_bit_receive) | ||
| self._send_ack_internal(ctrl_byte=ack, cid=self.channel_id) |
There was a problem hiding this comment.
Only flush the final ACK on the outermost workflow exit.
Line 459 currently runs for both inner and outer piggyback_acks() exits. In a nested interaction that sends an intermediate ACK while an outer workflow is still active, which breaks the “single final ACK” contract and reduces piggybacking across consecutive calls.
Suggested fix
finally:
active = self._active_workflow
self._active_workflow = previous
# Verify proper nesting of workflows
assert active is marker
- if self.is_ack_piggybacking_allowed:
+ if previous is None and self.is_ack_piggybacking_allowed:
# Explicitly ACK the last message received.
# Otherwise, the next message sent may be lost due to an event loop restart.
ack = control_byte.make_ack(not self.sync_bit_receive)
self._send_ack_internal(ctrl_byte=ack, cid=self.channel_id)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@python/src/trezorlib/thp/channel.py` around lines 455 - 463, The ACK is being
sent on every piggyback_acks() exit; modify the guard so the final ACK is only
sent when leaving the outermost workflow: after restoring self._active_workflow
= previous, only call self._send_ack_internal(...) if previous is None (i.e.
there is no enclosing workflow) and self.is_ack_piggybacking_allowed is true.
Update the condition that currently checks is_ack_piggybacking_allowed to also
require previous is None (or equivalently self._active_workflow is None) before
sending the ACK, referencing the existing variables/methods _active_workflow,
previous, is_ack_piggybacking_allowed, and _send_ack_internal.
that's how it's used, for sure, but the intention of (you'll notice that |
|
the intention is like so:
(note: the ACK context should not be scoped to a session because you may want to bundle multiple sessions together) aside: when processing ButtonRequests, the host sends ButtonAck immediately, ACKing the receipt of ButtonRequest, but how does Trezor ACK the ButtonAck? the next message won't arrive until the user clicks a button. does Trezor send an explicit ACK here? |
Yes - the device sends an explicit THP ACKs. Currently, send-side THP ACK piggybacking is not enabled on the device - so we'll need to handle those cases when it will be enabled. |
|
Closing in favor of #6594. |
Following #6563 (comment). [no changelog]
Following #6563 (comment). [no changelog]
Following #6563 (comment). [no changelog]
Following #6563 (comment). [no changelog]
Following #6563 (comment). [no changelog]
Following #6563 (comment). [no changelog]
Following #6563 (comment). [no changelog]
Following #6563 (comment). [no changelog]
Following #6563 (comment). [no changelog]
Following #6563 (comment). [no changelog]
Following #6563 (comment). [no changelog]
Following #6563 (comment). [no changelog]
2 participants
It allows scoping "interactive" context, where THP ACK piggybacking will be enabled (#6557).