Refactor socket to handle both action and observation types (#250)

frontend/src/components/Terminal.tsx

@@ -2,7 +2,7 @@ import React, { useEffect, useRef } from "react";
 import { IDisposable, Terminal as XtermTerminal } from "@xterm/xterm";
 import { FitAddon } from "xterm-addon-fit";
 import "@xterm/xterm/css/xterm.css";
-import socket from "../state/socket";
+import socket from "../socket/socket";
 
 class JsonWebsocketAddon {
   _socket: WebSocket;

frontend/src/services/chatService.ts

@@ -1,5 +1,5 @@
 import { appendUserMessage } from "../state/chatSlice";
-import socket from "../state/socket";
+import socket from "../socket/socket";
 import store from "../store";
 
 export function sendChatMessage(message: string): void {

frontend/src/socket/actions.ts

@@ -0,0 +1,39 @@
+import store from "../store";
+import { ActionMessage } from "./types/Message";
+import { setScreenshotSrc, setUrl } from "../state/browserSlice";
+import { appendAssistantMessage } from "../state/chatSlice";
+import { setCode } from "../state/codeSlice";
+import { setInitialized } from "../state/taskSlice";
+
+const messageActions = {
+  initialize: () => {
+    store.dispatch(setInitialized(true));
+    store.dispatch(
+      appendAssistantMessage(
+        "Hello, I am OpenDevin, an AI Software Engineer. What would you like me to build you today?",
+      ),
+    );
+  },
+  browse: (message: ActionMessage) => {
+    const { url, screenshotSrc } = message.args;
+    store.dispatch(setUrl(url));
+    store.dispatch(setScreenshotSrc(screenshotSrc));
+  },
+  write: (message: ActionMessage) => {
+    store.dispatch(setCode(message.args.contents));
+  },
+  think: (message: ActionMessage) => {
+    store.dispatch(appendAssistantMessage(message.args.thought));
+  },
+  finish: (message: ActionMessage) => {
+    store.dispatch(appendAssistantMessage(message.message));
+  },
+};
+
+export function handleActionMessage(message: ActionMessage) {
+  if (message.action in messageActions) {
+    const actionFn =
+      messageActions[message.action as keyof typeof messageActions];
+    actionFn(message);
+  }
+}

frontend/src/socket/observations.ts

@@ -0,0 +1,7 @@
+import { appendAssistantMessage } from "../state/chatSlice";
+import store from "../store";
+import { ObservationMessage } from "./types/Message";
+
+export function handleObservationMessage(message: ObservationMessage) {
+  store.dispatch(appendAssistantMessage(message.message));
+}

frontend/src/socket/socket.ts

@@ -0,0 +1,34 @@
+import store from "../store";
+import { ActionMessage, ObservationMessage } from "./types/Message";
+import { appendError } from "../state/errorsSlice";
+import { handleActionMessage } from "./actions";
+import { handleObservationMessage } from "./observations";
+
+type SocketMessage = ActionMessage | ObservationMessage;
+
+const WS_URL = import.meta.env.VITE_TERMINAL_WS_URL;
+if (!WS_URL) {
+  throw new Error(
+    "The environment variable VITE_TERMINAL_WS_URL is not set. Please set it to the WebSocket URL of the terminal server.",
+  );
+}
+
+const socket = new WebSocket(WS_URL);
+
+socket.addEventListener("message", (event) => {
+  const socketMessage = JSON.parse(event.data) as SocketMessage;
+  if ("action" in socketMessage) {
+    handleActionMessage(socketMessage);
+  } else {
+    handleObservationMessage(socketMessage);
+  }
+});
+socket.addEventListener("error", () => {
+  store.dispatch(
+    appendError(
+      `Failed connection to server. Please ensure the server is reachable at ${WS_URL}.`,
+    ),
+  );
+});
+
+export default socket;

frontend/src/state/socket.ts

@@ -1,70 +0,0 @@
-import store from "../store";
-import { setScreenshotSrc, setUrl } from "./browserSlice";
-import { appendAssistantMessage } from "./chatSlice";
-import { setCode } from "./codeSlice";
-import { appendError } from "./errorsSlice";
-import { setInitialized } from "./taskSlice";
-
-const MESSAGE_ACTIONS = ["terminal", "planner", "code", "browser"] as const;
-type MessageAction = (typeof MESSAGE_ACTIONS)[number];
-
-type SocketMessage = {
-  action: MessageAction;
-  message: string;
-  args: Record<string, unknown>;
-};
-
-const messageActions = {
-  initialize: () => {
-    store.dispatch(setInitialized(true));
-    store.dispatch(
-      appendAssistantMessage(
-        "Hello, I am OpenDevin, an AI Software Engineer. What would you like me to build you today?",
-      ),
-    );
-  },
-  browse: (message: SocketMessage) => {
-    const { url, screenshotSrc } = message.args;
-    store.dispatch(setUrl(url));
-    store.dispatch(setScreenshotSrc(screenshotSrc));
-  },
-  write: (message: SocketMessage) => {
-    store.dispatch(setCode(message.args.contents));
-  },
-  think: (message: SocketMessage) => {
-    store.dispatch(appendAssistantMessage(message.args.thought));
-  },
-  finish: (message: SocketMessage) => {
-    store.dispatch(appendAssistantMessage(message.message));
-  },
-};
-
-const WS_URL = import.meta.env.VITE_TERMINAL_WS_URL;
-if (!WS_URL) {
-  throw new Error(
-    "The environment variable VITE_TERMINAL_WS_URL is not set. Please set it to the WebSocket URL of the terminal server.",
-  );
-}
-
-const socket = new WebSocket(WS_URL);
-
-socket.addEventListener("message", (event) => {
-  const socketMessage = JSON.parse(event.data) as SocketMessage;
-
-  if (socketMessage.action in messageActions) {
-    const actionFn =
-      messageActions[socketMessage.action as keyof typeof messageActions];
-    actionFn(socketMessage);
-  } else if (!socketMessage.action) {
-    store.dispatch(appendAssistantMessage(socketMessage.message));
-  }
-});
-socket.addEventListener("error", () => {
-  store.dispatch(
-    appendError(
-      `Failed connection to server. Please ensure the server is reachable at ${WS_URL}.`,
-    ),
-  );
-});
-
-export default socket;