@sp33dy I've done this two ways, both for shape-oriented UIs. One way is to version each screen after a save and then allow you to revert to a previous saved screen. It's like a perpetual set of backup copies. This is similar to the Local History that you see in some apps.
I've also done something like this which grew out of an audit requirement. Each operation is supplemented with a data structure "Action" which is then recorded. The Undo feature would pair a rollback operation (an ADD paired with a DELETE, a DELETE paired with an ADD, and UPDATE paired with an offsetting UPDATE). Here's how that might look in TornadoFX.
https://courses.bekwam.net/public_tutorials/bkcourse_tornadofx_journalingapp.html