diff --git a/lib/controller.ts b/lib/controller.ts
index 722ddb21a6..d3503dbb79 100644
--- a/lib/controller.ts
+++ b/lib/controller.ts
@@ -39,7 +39,8 @@ const AllExtensions = [
 ];
 
 type ExtensionArgs = [Zigbee, MQTT, State, PublishEntityState, EventBus,
-    (enable: boolean, name: string) => Promise<void>, () => void, (extension: Extension) => Promise<void>];
+    enableDisableExtension: (enable: boolean, name: string) => Promise<void>, restartCallback: () => Promise<void>,
+    addExtension: (extension: Extension) => Promise<void>];
 
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
 let sdNotify: any = null;
@@ -54,12 +55,12 @@ export class Controller {
     private zigbee: Zigbee;
     private state: State;
     private mqtt: MQTT;
-    private restartCallback: () => void;
-    private exitCallback: (code: number, restart: boolean) => void;
+    private restartCallback: () => Promise<void>;
+    private exitCallback: (code: number, restart: boolean) => Promise<void>;
     private extensions: Extension[];
     private extensionArgs: ExtensionArgs;
 
-    constructor(restartCallback: () => void, exitCallback: (code: number, restart: boolean) => void) {
+    constructor(restartCallback: () => Promise<void>, exitCallback: (code: number, restart: boolean) => Promise<void>) {
         logger.init();
         zhSetLogger(logger);
         zhcSetLogger(logger);
@@ -113,7 +114,7 @@ export class Controller {
             logger.error('Check https://www.zigbee2mqtt.io/guide/installation/20_zigbee2mqtt-fails-to-start.html for possible solutions'); /* eslint-disable-line max-len */
             logger.error('Exiting...');
             logger.error(error.stack);
-            await this.exit(1);
+            return this.exit(1);
         }
 
         // Disable some legacy options on new network creation
@@ -154,7 +155,7 @@ export class Controller {
         } catch (error) {
             logger.error(`MQTT failed to connect, exiting...`);
             await this.zigbee.stop();
-            await this.exit(1);
+            return this.exit(1);
         }
 
         // Call extensions
@@ -212,22 +213,23 @@ export class Controller {
         // Wrap-up
         this.state.stop();
         await this.mqtt.disconnect();
+        let code = 0;
 
         try {
             await this.zigbee.stop();
             logger.info('Stopped Zigbee2MQTT');
-            await this.exit(0, restart);
         } catch (error) {
             logger.error('Failed to stop Zigbee2MQTT');
-            await this.exit(1, restart);
+            code = 1;
         }
 
         sdNotify?.stopWatchdogMode();
+        return this.exit(code, restart);
     }
 
     async exit(code: number, restart = false): Promise<void> {
         await logger.end();
-        this.exitCallback(code, restart);
+        return this.exitCallback(code, restart);
     }
 
     @bind async onZigbeeAdapterDisconnected(): Promise<void> {
diff --git a/lib/extension/extension.ts b/lib/extension/extension.ts
index efad18e6d5..d5ae7a798d 100644
--- a/lib/extension/extension.ts
+++ b/lib/extension/extension.ts
@@ -5,7 +5,7 @@ abstract class Extension {
     protected publishEntityState: PublishEntityState;
     protected eventBus: EventBus;
     protected enableDisableExtension: (enable: boolean, name: string) => Promise<void>;
-    protected restartCallback: () => void;
+    protected restartCallback: () => Promise<void>;
     protected addExtension: (extension: Extension) => Promise<void>;
 
     /**
@@ -22,7 +22,7 @@ abstract class Extension {
      */
     constructor(zigbee: Zigbee, mqtt: MQTT, state: State, publishEntityState: PublishEntityState,
         eventBus: EventBus, enableDisableExtension: (enable: boolean, name: string) => Promise<void>,
-        restartCallback: () => void, addExtension: (extension: Extension) => Promise<void>) {
+        restartCallback: () => Promise<void>, addExtension: (extension: Extension) => Promise<void>) {
         this.zigbee = zigbee;
         this.mqtt = mqtt;
         this.state = state;
diff --git a/lib/extension/externalConverters.ts b/lib/extension/externalConverters.ts
index 138d1e063c..df13750610 100644
--- a/lib/extension/externalConverters.ts
+++ b/lib/extension/externalConverters.ts
@@ -7,7 +7,7 @@ import logger from '../util/logger';
 export default class ExternalConverters extends Extension {
     constructor(zigbee: Zigbee, mqtt: MQTT, state: State, publishEntityState: PublishEntityState,
         eventBus: EventBus, enableDisableExtension: (enable: boolean, name: string) => Promise<void>,
-        restartCallback: () => void, addExtension: (extension: Extension) => Promise<void>) {
+        restartCallback: () => Promise<void>, addExtension: (extension: Extension) => Promise<void>) {
         super(zigbee, mqtt, state, publishEntityState, eventBus, enableDisableExtension, restartCallback, addExtension);
 
         for (const file of settings.get().external_converters) {
diff --git a/lib/extension/frontend.ts b/lib/extension/frontend.ts
index babc4c59c3..a85ebb8a5d 100644
--- a/lib/extension/frontend.ts
+++ b/lib/extension/frontend.ts
@@ -32,7 +32,7 @@ export default class Frontend extends Extension {
 
     constructor(zigbee: Zigbee, mqtt: MQTT, state: State, publishEntityState: PublishEntityState,
         eventBus: EventBus, enableDisableExtension: (enable: boolean, name: string) => Promise<void>,
-        restartCallback: () => void, addExtension: (extension: Extension) => Promise<void>) {
+        restartCallback: () => Promise<void>, addExtension: (extension: Extension) => Promise<void>) {
         super(zigbee, mqtt, state, publishEntityState, eventBus, enableDisableExtension, restartCallback, addExtension);
         this.eventBus.onMQTTMessagePublished(this, this.onMQTTPublishMessage);
     }
diff --git a/lib/extension/homeassistant.ts b/lib/extension/homeassistant.ts
index da684e7c9b..5b900bb254 100644
--- a/lib/extension/homeassistant.ts
+++ b/lib/extension/homeassistant.ts
@@ -126,7 +126,7 @@ export default class HomeAssistant extends Extension {
 
     constructor(zigbee: Zigbee, mqtt: MQTT, state: State, publishEntityState: PublishEntityState,
         eventBus: EventBus, enableDisableExtension: (enable: boolean, name: string) => Promise<void>,
-        restartCallback: () => void, addExtension: (extension: Extension) => Promise<void>) {
+        restartCallback: () => Promise<void>, addExtension: (extension: Extension) => Promise<void>) {
         super(zigbee, mqtt, state, publishEntityState, eventBus, enableDisableExtension, restartCallback, addExtension);
         if (settings.get().advanced.output === 'attribute') {
             throw new Error('Home Assistant integration is not possible with attribute output!');