diff --git a/core/main.py b/core/main.py
index 9d67efb..7e717a1 100644
--- a/core/main.py
+++ b/core/main.py
@@ -71,7 +71,6 @@ def main():
 
     if args.mode == "source":
         from source.source import run_source_mode
-        state.reset_state(mode="source")
         try:
             run_source_mode()
         except KeyboardInterrupt:
@@ -122,7 +121,6 @@ def main():
             _handle_keyboard_interrupt()
         if choice == "1":
             from source.source import run_source_mode
-            state.reset_state(mode="source")
             try:
                 run_source_mode()
             except KeyboardInterrupt:
diff --git a/source/source.py b/source/source.py
index b938ce5..acdee08 100644
--- a/source/source.py
+++ b/source/source.py
@@ -22,7 +22,7 @@ _ARCHIVE_DIR = "/tmp/docker-migrate-archives"
 
 
 def run_source_mode():
-    state.set_stage("INIT", mode="source")
+    state.reset_state(mode="source")
     from core.fsm import FSM
     fsm = FSM(mode="source")
     try:
diff --git a/target/target.py b/target/target.py
index b5da6b4..08f972b 100644
--- a/target/target.py
+++ b/target/target.py
@@ -17,7 +17,7 @@ _RESTORE_DIR = "/opt/migrate-restore"
 
 
 def run_target_mode():
-    state.set_stage("INIT", mode="target")
+    state.reset_state(mode="target")
     from core.fsm import FSM
     fsm = FSM(mode="target")
     fsm.resume_from("INIT")