Only completely prevent `CAP END` when we get `RESUME SUCCESS` (resume.py)

1 files changed, 10 insertions, 2 deletions

diff --git a/modules/resume.py b/modules/resume.py
index 90e516aa..c51d25c7 100644
--- a/modules/resume.py
+++ b/modules/resume.py
@@ -32,12 +32,17 @@ class Module(ModuleManager.BaseModule):
 
     @utils.hook("received.resume")
     def on_resume(self, event):
+        cap_done = False
+
         if event["args"][0] == "SUCCESS":
             resume_channels = event["server"].get_setting("resume-channels", [])
             self.log.info("Successfully resumed session", [])
+            event["server"].cap_started = False
+            cap_done = True
 
         elif event["args"][0] == "ERR":
             self.log.info("Failed to resume session: %s", [event["args"][1]])
+            cap_done = True
 
         elif event["args"][0] == "TOKEN":
             token = self._get_token(event["server"])
@@ -48,9 +53,12 @@ class Module(ModuleManager.BaseModule):
 
                 event["server"].send("RESUME %s%s" %
                     (token, " %s" % timestamp if timestamp else ""))
-                event["server"].cap_started = False
+            else:
+                cap_done = True
+
+        if cap_done:
+            event["server"].capability_done("resume")
 
-        event["server"].capability_done("resume")
 
     @utils.hook("received.001")
     def on_connect(self, event):