From aa1bda2abcb6312b006bc5a0a97877f7a5fa122e Mon Sep 17 00:00:00 2001
From: Eric Bailey <git@esb.lol>
Date: Fri, 10 May 2024 10:08:28 -0500
Subject: [PATCH] Setup/teardown events

---
 src/state/messages/convo/agent.ts | 56 ++++++++++++++++++++-----------
 1 file changed, 36 insertions(+), 20 deletions(-)

diff --git a/src/state/messages/convo/agent.ts b/src/state/messages/convo/agent.ts
index 685cc58b10..b8e64b241e 100644
--- a/src/state/messages/convo/agent.ts
+++ b/src/state/messages/convo/agent.ts
@@ -107,26 +107,6 @@ export class Convo {
     } else {
       DEBUG_ACTIVE_CHAT = this.convoId
     }
-
-    this.events.on(
-      event => {
-        switch (event.type) {
-          case 'connect': {
-            this.onFirehoseConnect()
-            break
-          }
-          case 'error': {
-            this.onFirehoseError(event.error)
-            break
-          }
-          case 'logs': {
-            this.ingestFirehose(event.logs)
-            break
-          }
-        }
-      },
-      {convoId: this.convoId},
-    )
   }
 
   private commit() {
@@ -225,6 +205,7 @@ export class Convo {
           case ConvoDispatchEvent.Init: {
             this.status = ConvoStatus.Initializing
             this.setup()
+            this.setupFirehose()
             this.requestPollInterval(ACTIVE_POLL_INTERVAL)
             break
           }
@@ -246,12 +227,14 @@ export class Convo {
           }
           case ConvoDispatchEvent.Suspend: {
             this.status = ConvoStatus.Suspended
+            this.cleanupFirehoseConnection?.()
             this.withdrawRequestedPollInterval()
             break
           }
           case ConvoDispatchEvent.Error: {
             this.status = ConvoStatus.Error
             this.error = action.payload
+            this.cleanupFirehoseConnection?.()
             this.withdrawRequestedPollInterval()
             break
           }
@@ -272,12 +255,14 @@ export class Convo {
           }
           case ConvoDispatchEvent.Suspend: {
             this.status = ConvoStatus.Suspended
+            this.cleanupFirehoseConnection?.()
             this.withdrawRequestedPollInterval()
             break
           }
           case ConvoDispatchEvent.Error: {
             this.status = ConvoStatus.Error
             this.error = action.payload
+            this.cleanupFirehoseConnection?.()
             this.withdrawRequestedPollInterval()
             break
           }
@@ -300,12 +285,14 @@ export class Convo {
           }
           case ConvoDispatchEvent.Suspend: {
             this.status = ConvoStatus.Suspended
+            this.cleanupFirehoseConnection?.()
             this.withdrawRequestedPollInterval()
             break
           }
           case ConvoDispatchEvent.Error: {
             this.status = ConvoStatus.Error
             this.error = action.payload
+            this.cleanupFirehoseConnection?.()
             this.withdrawRequestedPollInterval()
             break
           }
@@ -615,6 +602,33 @@ export class Convo {
     }
   }
 
+  private cleanupFirehoseConnection: (() => void) | undefined
+  private setupFirehose() {
+    // remove old listeners, if exist
+    this.cleanupFirehoseConnection?.()
+
+    // reconnect
+    this.cleanupFirehoseConnection = this.events.on(
+      event => {
+        switch (event.type) {
+          case 'connect': {
+            this.onFirehoseConnect()
+            break
+          }
+          case 'error': {
+            this.onFirehoseError(event.error)
+            break
+          }
+          case 'logs': {
+            this.ingestFirehose(event.logs)
+            break
+          }
+        }
+      },
+      {convoId: this.convoId},
+    )
+  }
+
   onFirehoseConnect() {
     this.footerItems.delete(ConvoItemError.FirehoseFailed)
     this.commit()
@@ -727,6 +741,8 @@ export class Convo {
       id: tempId,
       message,
     })
+    // remove on each send, it might go through now without user having to click
+    this.footerItems.delete(ConvoItemError.PendingFailed)
     this.commit()
 
     if (!this.isProcessingPendingMessages) {