fix!: cherry-pick 19038 and add tests

[julienrbrt]

Description

Cherry-pick 19038 and add more tests to harden the feature.

---

Author Checklist

All items are required. Please add a note to the item if the item is not applicable and
please add links to any relevant follow up issues.

I have...

• included the correct type prefix in the PR title
• added ! to the type prefix if API or client breaking change
• targeted the correct branch (see PR Targeting)
• provided a link to the relevant issue or specification
• followed the guidelines for building modules
• included the necessary unit and integration tests
• added a changelog entry to CHANGELOG.md
• included comments for documenting Go code
• updated the relevant documentation or specification
• reviewed "Files changed" and left comments if necessary
• run make lint and make test
• confirmed all CI checks have passed

Reviewers Checklist

All items are required. Please add a note if the item is not applicable and please add
your handle next to the items reviewed if you only reviewed selected items.

I have...

• confirmed the correct type prefix in the PR title
• confirmed ! in the type prefix if API or client breaking change
• confirmed all author checklist items have been addressed
• reviewed state machine logic
• reviewed API design and naming
• reviewed documentation is accurate
• reviewed tests and test coverage
• manually tested (if applicable)