Display connection rejection errors passed to client (#6101)

Co-authored-by: sp.wack <[email protected]>

frontend/__tests__/context/ws-client-provider.test.tsx

@@ -0,0 +1,30 @@
+import { describe, it, expect, vi } from "vitest";
+import { render, screen } from "@testing-library/react";
+import * as ChatSlice from "#/state/chat-slice";
+import {
+  updateStatusWhenErrorMessagePresent,
+} from "#/context/ws-client-provider";
+
+describe("Propagate error message", () => {
+  it("should do nothing when no message was passed from server", () => {
+    const addErrorMessageSpy = vi.spyOn(ChatSlice, "addErrorMessage");
+    updateStatusWhenErrorMessagePresent(null)
+    updateStatusWhenErrorMessagePresent(undefined)
+    updateStatusWhenErrorMessagePresent({})
+    updateStatusWhenErrorMessagePresent({message: null})
+
+    expect(addErrorMessageSpy).not.toHaveBeenCalled();
+  });
+
+  it("should display error to user when present", () => {
+    const message = "We have a problem!"
+    const addErrorMessageSpy = vi.spyOn(ChatSlice, "addErrorMessage")
+    updateStatusWhenErrorMessagePresent({message})
+
+    expect(addErrorMessageSpy).toHaveBeenCalledWith({
+      message,
+      status_update: true,
+      type: 'error'
+     });
+  });
+});

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 {
@@ -64,6 +67,21 @@ interface WsClientProviderProps {
   conversationId: string;
 }
 
+export function updateStatusWhenErrorMessagePresent(data: unknown) {
+  if (
+    data &&
+    typeof data === "object" &&
+    "message" in data &&
+    typeof data.message === "string"
+  ) {
+    handleStatusMessage({
+      type: "error",
+      message: data.message,
+      status_update: true,
+    });
+  }
+}
+
 export function WsClientProvider({
   conversationId,
   children,
@@ -101,19 +119,21 @@ export function WsClientProvider({
     handleAssistantMessage(event);
   }
 
-  function handleDisconnect() {
+  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(() => {

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));
   }

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 = [];
     },
   },

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;
 }