Fix exception after return of SM wrapper
Correct type hint on state container