Port the functionality in #907 to main and add test coverage.
Port the functionality in #907 to main and add test coverage.