From 4c0e093c77eee3219e34bb2ee4b1f7731cc19e54 Mon Sep 17 00:00:00 2001
From: gnuxie <Gnuxie@protonmail.com>
Date: Thu, 19 Sep 2024 20:12:07 +0100
Subject: [PATCH] Fix draupnir manager for safe mode.

---
 .../StandardDraupnirManager.ts                | 19 +++++++++++++++----
 1 file changed, 15 insertions(+), 4 deletions(-)

diff --git a/src/draupnirfactory/StandardDraupnirManager.ts b/src/draupnirfactory/StandardDraupnirManager.ts
index 254ece75..dac0b5cb 100644
--- a/src/draupnirfactory/StandardDraupnirManager.ts
+++ b/src/draupnirfactory/StandardDraupnirManager.ts
@@ -68,7 +68,7 @@ export class StandardDraupnirManager {
       config,
       this.makeSafeModeToggle(clientUserID, managementRoom, config)
     );
-    if (this.isDraupnirAvailable(clientUserID)) {
+    if (this.isNormalDraupnir(clientUserID)) {
       return ActionError.Result(
         `There is a draupnir for ${clientUserID} already running`
       );
@@ -83,6 +83,7 @@ export class StandardDraupnirManager {
     }
     this.draupnir.set(clientUserID, draupnir.ok);
     this.failedDraupnir.delete(clientUserID);
+    this.safeModeDraupnir.delete(clientUserID);
     draupnir.ok.start();
     return draupnir;
   }
@@ -93,7 +94,7 @@ export class StandardDraupnirManager {
     config: IConfig,
     cause: SafeModeCause
   ): Promise<ActionResult<SafeModeDraupnir>> {
-    if (this.isDraupnirAvailable(clientUserID)) {
+    if (this.isSafeModeDraupnir(clientUserID)) {
       return ActionError.Result(
         `There is a draupnir for ${clientUserID} already running`
       );
@@ -115,16 +116,26 @@ export class StandardDraupnirManager {
     }
     safeModeDraupnir.ok.start();
     this.safeModeDraupnir.set(clientUserID, safeModeDraupnir.ok);
+    this.draupnir.delete(clientUserID);
+    this.failedDraupnir.delete(clientUserID);
     return safeModeDraupnir;
   }
 
+  private isNormalDraupnir(drapunirClientID: StringUserID): boolean {
+    return this.draupnir.has(drapunirClientID);
+  }
+
+  private isSafeModeDraupnir(drapunirClientID: StringUserID): boolean {
+    return this.safeModeDraupnir.has(drapunirClientID);
+  }
+
   /**
    * Whether the draupnir is available to the user, either normally or via safe mode.
    */
   public isDraupnirAvailable(draupnirClientID: StringUserID): boolean {
     return (
-      this.draupnir.has(draupnirClientID) ||
-      this.safeModeDraupnir.has(draupnirClientID)
+      this.isNormalDraupnir(draupnirClientID) ||
+      this.isSafeModeDraupnir(draupnirClientID)
     );
   }