Display connection rejection errors passed to client

All-Hands-AI · Jan 7, 2025 · c26e9fd · c26e9fd
1 parent 0973446
commit c26e9fd
Show file tree

Hide file tree

Showing 4 changed files with 30 additions and 10 deletions.
diff --git a/frontend/src/context/ws-client-provider.tsx b/frontend/src/context/ws-client-provider.tsx
@@ -2,7 +2,10 @@ import posthog from "posthog-js";
 import React from "react";
 import { io, Socket } from "socket.io-client";
 import EventLogger from "#/utils/event-logger";
-import { handleAssistantMessage } from "#/services/actions";
+import {
+  handleAssistantMessage,
+  handleStatusMessage,
+} from "#/services/actions";
 import { useRate } from "#/hooks/use-rate";
 import { OpenHandsParsedEvent } from "#/types/core";
 import {
@@ -104,19 +107,32 @@ export function WsClientProvider({
     handleAssistantMessage(event);
   }
 
-  function handleDisconnect() {
+  function updateStatusWhenErrorMessagePresent(data: unknown) {
+    if (data && typeof data === 'object' && 'message' in data
+      && typeof data.message === "string") {
+      handleStatusMessage({
+        type: "error",
+        message: data.message as string,
+        status_update: true,
+      });
+    }
+  }
+
+  function handleDisconnect(data: unknown) {
     setStatus(WsClientProviderStatus.DISCONNECTED);
     const sio = sioRef.current;
     if (!sio) {
       return;
     }
     sio.io.opts.query = sio.io.opts.query || {};
     sio.io.opts.query.latest_event_id = lastEventRef.current?.id;
+    updateStatusWhenErrorMessagePresent(data);
   }
 
-  function handleError() {
-    posthog.capture("socket_error");
+  function handleError(data: unknown) {
     setStatus(WsClientProviderStatus.DISCONNECTED);
+    updateStatusWhenErrorMessagePresent(data);
+    posthog.capture("socket_error");
   }
 
   React.useEffect(() => {

diff --git a/frontend/src/services/actions.ts b/frontend/src/services/actions.ts
@@ -75,6 +75,7 @@ export function handleActionMessage(message: ActionMessage) {
     if (message.args && message.args.thought) {
       store.dispatch(addAssistantMessage(message.args.thought));
     }
+    // Need to convert ActionMessage to RejectAction
     // @ts-expect-error TODO: fix
     store.dispatch(addAssistantAction(message));
   }

diff --git a/frontend/src/state/chat-slice.ts b/frontend/src/state/chat-slice.ts
@@ -73,7 +73,7 @@ export const chatSlice = createSlice({
       state.messages.push(message);
     },
 
-    addAssistantMessage(state, action: PayloadAction<string>) {
+    addAssistantMessage(state: SliceState, action: PayloadAction<string>) {
       const message: Message = {
         type: "thought",
         sender: "assistant",
@@ -85,7 +85,10 @@ export const chatSlice = createSlice({
       state.messages.push(message);
     },
 
-    addAssistantAction(state, action: PayloadAction<OpenHandsAction>) {
+    addAssistantAction(
+      state: SliceState,
+      action: PayloadAction<OpenHandsAction>,
+    ) {
       const actionID = action.payload.action;
       if (!HANDLED_ACTIONS.includes(actionID)) {
         return;
@@ -125,7 +128,7 @@ export const chatSlice = createSlice({
     },
 
     addAssistantObservation(
-      state,
+      state: SliceState,
       observation: PayloadAction<OpenHandsObservation>,
     ) {
       const observationID = observation.payload.observation;
@@ -179,7 +182,7 @@ export const chatSlice = createSlice({
     },
 
     addErrorMessage(
-      state,
+      state: SliceState,
       action: PayloadAction<{ id?: string; message: string }>,
     ) {
       const { id, message } = action.payload;
@@ -192,7 +195,7 @@ export const chatSlice = createSlice({
       });
     },
 
-    clearMessages(state) {
+    clearMessages(state: SliceState) {
       state.messages = [];
     },
   },

diff --git a/frontend/src/types/message.tsx b/frontend/src/types/message.tsx
@@ -43,6 +43,6 @@ export interface ObservationMessage {
 export interface StatusMessage {
   status_update: true;
   type: string;
-  id: string;
+  id?: string;
   message: string;
 }