Very cool, I was looking for something like this. Instead of (or in addition to) having get create a SessionState, could you consider allowing the user to create a Dataclass as state so the code complete in the editor is more helpful? From what I can tell SessionState is just a regular object so it should work.
Edit: Also once this is implemented it would be great if buttons can specify a callback. I know it might not fit 100% with the rest of streamlit’s top down declarative style but right now buttons are not very useful and this would enable a lot of things that are hard right now.