From 0657e1aaceb9d2ccc3d296793a98b4972a86b2ea Mon Sep 17 00:00:00 2001
From: Georgiy Tarasov <gtarasov.work@gmail.com>
Date: Tue, 29 Oct 2024 13:03:11 +0100
Subject: [PATCH] feat: handle network/parsing errors

---
 frontend/src/scenes/max/maxLogic.ts | 34 +++++++++++++++++++++++++----
 1 file changed, 30 insertions(+), 4 deletions(-)

diff --git a/frontend/src/scenes/max/maxLogic.ts b/frontend/src/scenes/max/maxLogic.ts
index 8a0f4e4df1119..c30550914a2e6 100644
--- a/frontend/src/scenes/max/maxLogic.ts
+++ b/frontend/src/scenes/max/maxLogic.ts
@@ -1,15 +1,17 @@
+import { captureException } from '@sentry/react'
 import { shuffle } from 'd3'
 import { createParser } from 'eventsource-parser'
 import { actions, kea, key, listeners, path, props, reducers, selectors } from 'kea'
 import { loaders } from 'kea-loaders'
 import api from 'lib/api'
-import { isHumanMessage } from 'scenes/max/utils'
+import { isHumanMessage, isVisualizationMessage } from 'scenes/max/utils'
 
 import {
     AssistantEventType,
     AssistantGenerationStatusEvent,
     AssistantGenerationStatusType,
     AssistantMessageType,
+    FailureMessage,
     NodeKind,
     RootAssistantMessage,
     SuggestedQuestionsQuery,
@@ -27,6 +29,11 @@ export type ThreadMessage = RootAssistantMessage & {
     status?: MessageStatus
 }
 
+const failureMessage: FailureMessage = {
+    type: AssistantMessageType.Failure,
+    content: 'Oops! It looks like I’m having trouble generating this trends insight. Could you please try again?',
+}
+
 export const maxLogic = kea<maxLogicType>([
     path(['scenes', 'max', 'maxLogic']),
     props({} as MaxLogicProps),
@@ -180,12 +187,31 @@ export const maxLogic = kea<maxLogicType>([
                     parser.feed(decoder.decode(value))
 
                     if (done) {
-                        actions.setMessageStatus(newIndex, 'completed')
+                        const generatedMessage = values.thread[newIndex]
+                        if (generatedMessage && isVisualizationMessage(generatedMessage) && generatedMessage.plan) {
+                            actions.setMessageStatus(newIndex, 'completed')
+                        } else if (generatedMessage) {
+                            actions.replaceMessage(newIndex, failureMessage)
+                        } else {
+                            actions.addMessage({
+                                ...failureMessage,
+                                status: 'completed',
+                            })
+                        }
                         break
                     }
                 }
-            } catch {
-                actions.setMessageStatus(values.thread.length - 1 === newIndex ? newIndex : newIndex - 1, 'error')
+            } catch (e) {
+                captureException(e)
+
+                if (values.thread[newIndex]) {
+                    actions.replaceMessage(newIndex, failureMessage)
+                } else {
+                    actions.addMessage({
+                        ...failureMessage,
+                        status: 'completed',
+                    })
+                }
             }
 
             actions.setThreadLoaded()