diff --git a/modules/wave_survival/state_machine.cpp b/modules/wave_survival/state_machine.cpp
index 393e5c52..fc9b5613 100644
--- a/modules/wave_survival/state_machine.cpp
+++ b/modules/wave_survival/state_machine.cpp
@@ -5,10 +5,12 @@ void StateMachine::_bind_methods() {}
 
 void StateMachine::switch_to_state(State *state) {
 	if (this->current_state != nullptr) {
+		this->current_state_is_active = false;
 		this->current_state->exit_state();
 	}
 	this->current_state = state;
 	if (this->current_state != nullptr) {
+		this->current_state_is_active = true;
 		this->current_state->enter_state();
 	} else {
 		print_error("StateMachine::switch_to_state: New state is nullptr, StateMachine will now stop working");
@@ -38,6 +40,20 @@ void StateMachine::_notification(int what) {
 		case NOTIFICATION_PROCESS:
 			process(get_process_delta_time());
 			return;
+		case NOTIFICATION_DISABLED:
+		case NOTIFICATION_EXIT_TREE:
+			if (this->current_state_is_active && this->current_state != nullptr) {
+				this->current_state_is_active = false;
+				this->current_state->exit_state();
+			}
+			return;
+		case NOTIFICATION_ENABLED:
+		case NOTIFICATION_ENTER_TREE:
+			if (!this->current_state_is_active && this->current_state != nullptr) {
+				this->current_state_is_active = true;
+				this->current_state->enter_state();
+			}
+			return;
 	}
 }
 
diff --git a/modules/wave_survival/state_machine.h b/modules/wave_survival/state_machine.h
index fc2ce297..3e8f7630 100644
--- a/modules/wave_survival/state_machine.h
+++ b/modules/wave_survival/state_machine.h
@@ -19,6 +19,7 @@ public:
 
 private:
 	State *current_state{ nullptr };
+	bool current_state_is_active{ false };
 	HashMap<StringName, State *> states{};
 };
 
diff --git a/project/objects/enemies/enemy_wretched.tscn b/project/objects/enemies/enemy_wretched.tscn
index ba33d5f5..ea34cc51 100644
--- a/project/objects/enemies/enemy_wretched.tscn
+++ b/project/objects/enemies/enemy_wretched.tscn
@@ -6,7 +6,8 @@
 script/source = "extends EnemyWretched
 
 func _on_health_status_death() -> void:
-	%StateMachine.set_process(false)
+	%StateMachine.process_mode = Node.PROCESS_MODE_DISABLED
+	%NavigationAgent3D.process_mode = Node.PROCESS_MODE_DISABLED
 	$wretched/AnimationPlayer.play(\"death\")
 	$CollisionShape3D.disabled = true
 	set_movement_direction(Vector2())