migration: Fix migration state update issue