Rewrite websocket reducer with RTK

ui/frontend/.eslintrc.js

@@ -68,6 +68,7 @@ module.exports = {
         'reducers/output/execute.ts',
         'reducers/output/format.ts',
         'reducers/output/gist.ts',
+        'reducers/websocket.ts',
         'websocketActions.ts',
         'websocketMiddleware.ts',
       ],

ui/frontend/.prettierignore

@@ -17,5 +17,6 @@ node_modules
 !reducers/output/execute.ts
 !reducers/output/format.ts
 !reducers/output/gist.ts
+!reducers/websocket.ts
 !websocketActions.ts
 !websocketMiddleware.ts

ui/frontend/actions.ts

@@ -1,6 +1,5 @@
 import fetch from 'isomorphic-fetch';
 import { ThunkAction as ReduxThunkAction, AnyAction } from '@reduxjs/toolkit';
-import { z } from 'zod';
 
 import {
   codeSelector,
@@ -120,18 +119,8 @@ export enum ActionType {
   NotificationSeen = 'NOTIFICATION_SEEN',
   BrowserWidthChanged = 'BROWSER_WIDTH_CHANGED',
   SplitRatioChanged = 'SPLIT_RATIO_CHANGED',
-  WebSocketError = 'WEBSOCKET_ERROR',
-  WebSocketConnected = 'WEBSOCKET_CONNECTED',
-  WebSocketDisconnected = 'WEBSOCKET_DISCONNECTED',
-  WebSocketFeatureFlagEnabled = 'WEBSOCKET_FEATURE_FLAG_ENABLED',
 }
 
-export const WebSocketError = z.object({
-  type: z.literal(ActionType.WebSocketError),
-  error: z.string(),
-});
-export type WebSocketError = z.infer<typeof WebSocketError>;
-
 export const initializeApplication = () => createAction(ActionType.InitializeApplication);
 
 export const disableSyncChangesToStorage = () => createAction(ActionType.DisableSyncChangesToStorage);
@@ -675,11 +664,6 @@ export const browserWidthChanged = (isSmall: boolean) =>
 export const splitRatioChanged = () =>
   createAction(ActionType.SplitRatioChanged);
 
-export const websocketError = (error: string): WebSocketError => createAction(ActionType.WebSocketError, { error });
-export const websocketConnected = () => createAction(ActionType.WebSocketConnected);
-export const websocketDisconnected = () => createAction(ActionType.WebSocketDisconnected);
-export const websocketFeatureFlagEnabled = () => createAction(ActionType.WebSocketFeatureFlagEnabled);
-
 function parseChannel(s?: string): Channel | null {
   switch (s) {
     case 'stable':
@@ -821,9 +805,5 @@ export type Action =
   | ReturnType<typeof notificationSeen>
   | ReturnType<typeof browserWidthChanged>
   | ReturnType<typeof splitRatioChanged>
-  | ReturnType<typeof websocketError>
-  | ReturnType<typeof websocketConnected>
-  | ReturnType<typeof websocketDisconnected>
-  | ReturnType<typeof websocketFeatureFlagEnabled>
   | ReturnType<typeof wsExecuteRequest>
   ;

ui/frontend/index.tsx

@@ -19,11 +19,11 @@ import {
   performVersionsLoad,
   reExecuteWithBacktrace,
   browserWidthChanged,
-  websocketFeatureFlagEnabled,
 } from './actions';
 import { configureRustErrors } from './highlighting';
 import PageSwitcher from './PageSwitcher';
 import playgroundApp from './reducers';
+import { websocketFeatureFlagEnabled } from './reducers/websocket';
 import Router from './Router';
 import configureStore from './configureStore';
 

ui/frontend/reducers/websocket.ts

@@ -1,31 +1,57 @@
-import { Action, ActionType } from '../actions';
+import { PayloadAction, createSlice } from '@reduxjs/toolkit';
+import z from 'zod';
+
+import { createWebsocketResponseSchema } from '../websocketActions';
 
 export type State = {
   connected: boolean;
   error?: string;
   featureFlagEnabled: boolean;
 };
 
-const DEFAULT: State = {
+const initialState: State = {
   connected: false,
   featureFlagEnabled: false,
 };
 
-export default function websocket(state = DEFAULT, action: Action): State {
-  switch (action.type) {
-    case ActionType.WebSocketConnected:
-      return { ...state, connected: true, error: undefined };
+const websocketErrorPayloadSchema = z.object({
+  error: z.string(),
+});
+type websocketErrorPayload = z.infer<typeof websocketErrorPayloadSchema>;
+
+const slice = createSlice({
+  name: 'websocket',
+  initialState,
+  reducers: {
+    connected: (state) => {
+      state.connected = true;
+      delete state.error;
+    },
+
+    disconnected: (state) => {
+      state.connected = false;
+    },
+
+    error: (state, action: PayloadAction<websocketErrorPayload>) => {
+      state.error = action.payload.error;
+    },
 
-    case ActionType.WebSocketDisconnected:
-      return { ...state, connected: false };
+    featureFlagEnabled: (state) => {
+      state.featureFlagEnabled = true;
+    },
+  },
+});
 
-    case ActionType.WebSocketError:
-      return { ...state, error: action.error };
+export const {
+  connected: websocketConnected,
+  disconnected: websocketDisconnected,
+  error: websocketError,
+  featureFlagEnabled: websocketFeatureFlagEnabled,
+} = slice.actions;
 
-    case ActionType.WebSocketFeatureFlagEnabled:
-      return { ...state, featureFlagEnabled: true };
+export const websocketErrorSchema = createWebsocketResponseSchema(
+  websocketError,
+  websocketErrorPayloadSchema,
+);
 
-    default:
-      return state;
-  }
-}
+export default slice.reducer;

ui/frontend/websocketMiddleware.ts

@@ -1,15 +1,18 @@
 import { AnyAction, Middleware } from '@reduxjs/toolkit';
 import { z } from 'zod';
 
+import { wsExecuteResponseSchema } from './reducers/output/execute';
 import {
-  WebSocketError,
   websocketConnected,
   websocketDisconnected,
   websocketError,
-} from './actions';
-import { wsExecuteResponseSchema } from './reducers/output/execute';
+  websocketErrorSchema,
+} from './reducers/websocket';
 
-const WSMessageResponse = z.discriminatedUnion('type', [WebSocketError, wsExecuteResponseSchema]);
+const WSMessageResponse = z.discriminatedUnion('type', [
+  websocketErrorSchema,
+  wsExecuteResponseSchema,
+]);
 
 const reportWebSocketError = async (error: string) => {
   try {
@@ -93,7 +96,7 @@ export const websocketMiddleware =
           // We cannot get detailed information about the failure
           // https://stackoverflow.com/a/31003057/155423
           const error = 'Generic WebSocket Error';
-          store.dispatch(websocketError(error));
+          store.dispatch(websocketError({ error }));
           reportWebSocketError(error);
         });
 

ui/src/server_axum/websocket.rs

@@ -63,8 +63,8 @@ impl TryFrom<ExecuteRequest> for sandbox::ExecuteRequest {
 #[derive(Debug, serde::Serialize)]
 #[serde(tag = "type")]
 enum MessageResponse {
-    #[serde(rename = "WEBSOCKET_ERROR")]
-    Error(WSError),
+    #[serde(rename = "websocket/error")]
+    Error { payload: WSError, meta: Meta },
 
     #[serde(rename = "output/execute/wsExecuteResponse")]
     ExecuteResponse {
@@ -174,7 +174,12 @@ pub async fn handle(mut socket: WebSocket) {
 
 fn error_to_response(error: Error) -> MessageResponse {
     let error = error.to_string();
-    MessageResponse::Error(WSError { error })
+    // TODO: thread through the Meta from the originating request
+    let meta = serde_json::json!({ "sequenceNumber": -1 });
+    MessageResponse::Error {
+        payload: WSError { error },
+        meta,
+    }
 }
 
 fn response_to_message(response: MessageResponse) -> Message {