refactor(frontend): migration of browser-slice.ts to zustand (#11081)

frontend/__tests__/components/browser.test.tsx

@@ -13,7 +13,8 @@ vi.mock("react-router", async () => {
 
 vi.mock("#/context/conversation-context", () => ({
   useConversation: () => ({ conversationId: "test-conversation-id" }),
-  ConversationProvider: ({ children }: { children: React.ReactNode }) => children,
+  ConversationProvider: ({ children }: { children: React.ReactNode }) =>
+    children,
 }));
 
 vi.mock("react-i18next", async () => {
@@ -29,21 +30,18 @@ vi.mock("react-i18next", async () => {
   };
 });
 
-// Mock redux
-const mockDispatch = vi.fn();
+// Mock Zustand browser store
 let mockBrowserState = {
   url: "https://example.com",
   screenshotSrc: "",
+  setUrl: vi.fn(),
+  setScreenshotSrc: vi.fn(),
+  reset: vi.fn(),
 };
 
-vi.mock("react-redux", async () => {
-  const actual = await vi.importActual("react-redux");
-  return {
-    ...actual,
-    useDispatch: () => mockDispatch,
-    useSelector: () => mockBrowserState,
-  };
-});
+vi.mock("#/stores/browser-store", () => ({
+  useBrowserStore: () => mockBrowserState,
+}));
 
 // Import the component after all mocks are set up
 import { BrowserPanel } from "#/components/features/browser/browser";
@@ -55,6 +53,9 @@ describe("Browser", () => {
     mockBrowserState = {
       url: "https://example.com",
       screenshotSrc: "",
+      setUrl: vi.fn(),
+      setScreenshotSrc: vi.fn(),
+      reset: vi.fn(),
     };
   });
 
@@ -63,6 +64,9 @@ describe("Browser", () => {
     mockBrowserState = {
       url: "https://example.com",
       screenshotSrc: "",
+      setUrl: vi.fn(),
+      setScreenshotSrc: vi.fn(),
+      reset: vi.fn(),
     };
 
     render(<BrowserPanel />);
@@ -75,7 +79,11 @@ describe("Browser", () => {
     // Set the mock state for this test
     mockBrowserState = {
       url: "https://example.com",
-      screenshotSrc: "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mN0uGvyHwAFCAJS091fQwAAAABJRU5ErkJggg==",
+      screenshotSrc:
+        "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mN0uGvyHwAFCAJS091fQwAAAABJRU5ErkJggg==",
+      setUrl: vi.fn(),
+      setScreenshotSrc: vi.fn(),
+      reset: vi.fn(),
     };
 
     render(<BrowserPanel />);

frontend/src/components/features/browser/browser.tsx

@@ -1,26 +1,16 @@
 import { useEffect } from "react";
-import { useSelector, useDispatch } from "react-redux";
-import { RootState } from "#/store";
 import { BrowserSnapshot } from "./browser-snapshot";
 import { EmptyBrowserMessage } from "./empty-browser-message";
 import { useConversationId } from "#/hooks/use-conversation-id";
-import {
-  initialState as browserInitialState,
-  setUrl,
-  setScreenshotSrc,
-} from "#/state/browser-slice";
+import { useBrowserStore } from "#/stores/browser-store";
 
 export function BrowserPanel() {
-  const { url, screenshotSrc } = useSelector(
-    (state: RootState) => state.browser,
-  );
+  const { url, screenshotSrc, reset } = useBrowserStore();
   const { conversationId } = useConversationId();
-  const dispatch = useDispatch();
 
   useEffect(() => {
-    dispatch(setUrl(browserInitialState.url));
-    dispatch(setScreenshotSrc(browserInitialState.screenshotSrc));
-  }, [conversationId]);
+    reset();
+  }, [conversationId, reset]);
 
   const imgSrc =
     screenshotSrc && screenshotSrc.startsWith("data:image/png;base64,")

frontend/src/services/observations.ts

@@ -1,10 +1,10 @@
 import { setCurrentAgentState } from "#/state/agent-slice";
-import { setUrl, setScreenshotSrc } from "#/state/browser-slice";
 import store from "#/store";
 import { ObservationMessage } from "#/types/message";
 import { useCommandStore } from "#/state/command-store";
 import { appendJupyterOutput } from "#/state/jupyter-slice";
 import ObservationType from "#/types/observation-type";
+import { useBrowserStore } from "#/stores/browser-store";
 
 export function handleObservationMessage(message: ObservationMessage) {
   switch (message.observation) {
@@ -34,11 +34,14 @@ export function handleObservationMessage(message: ObservationMessage) {
       break;
     case ObservationType.BROWSE:
     case ObservationType.BROWSE_INTERACTIVE:
-      if (message.extras?.screenshot) {
-        store.dispatch(setScreenshotSrc(message.extras?.screenshot));
+      if (
+        message.extras?.screenshot &&
+        typeof message.extras.screenshot === "string"
+      ) {
+        useBrowserStore.getState().setScreenshotSrc(message.extras.screenshot);
       }
-      if (message.extras?.url) {
-        store.dispatch(setUrl(message.extras.url));
+      if (message.extras?.url && typeof message.extras.url === "string") {
+        useBrowserStore.getState().setUrl(message.extras.url);
       }
       break;
     case ObservationType.AGENT_STATE_CHANGED:
@@ -63,19 +66,29 @@ export function handleObservationMessage(message: ObservationMessage) {
 
     switch (observation) {
       case "browse":
-        if (message.extras?.screenshot) {
-          store.dispatch(setScreenshotSrc(message.extras.screenshot));
+        if (
+          message.extras?.screenshot &&
+          typeof message.extras.screenshot === "string"
+        ) {
+          useBrowserStore
+            .getState()
+            .setScreenshotSrc(message.extras.screenshot);
         }
-        if (message.extras?.url) {
-          store.dispatch(setUrl(message.extras.url));
+        if (message.extras?.url && typeof message.extras.url === "string") {
+          useBrowserStore.getState().setUrl(message.extras.url);
         }
         break;
       case "browse_interactive":
-        if (message.extras?.screenshot) {
-          store.dispatch(setScreenshotSrc(message.extras.screenshot));
+        if (
+          message.extras?.screenshot &&
+          typeof message.extras.screenshot === "string"
+        ) {
+          useBrowserStore
+            .getState()
+            .setScreenshotSrc(message.extras.screenshot);
         }
-        if (message.extras?.url) {
-          store.dispatch(setUrl(message.extras.url));
+        if (message.extras?.url && typeof message.extras.url === "string") {
+          useBrowserStore.getState().setUrl(message.extras.url);
         }
         break;
       default:

frontend/src/state/browser-slice.ts

@@ -1,25 +0,0 @@
-import { createSlice } from "@reduxjs/toolkit";
-
-export const initialState = {
-  // URL of browser window (placeholder for now, will be replaced with the actual URL later)
-  url: "https://github.com/All-Hands-AI/OpenHands",
-  // Base64-encoded screenshot of browser window (placeholder for now, will be replaced with the actual screenshot later)
-  screenshotSrc: "",
-};
-
-export const browserSlice = createSlice({
-  name: "browser",
-  initialState,
-  reducers: {
-    setUrl: (state, action) => {
-      state.url = action.payload;
-    },
-    setScreenshotSrc: (state, action) => {
-      state.screenshotSrc = action.payload;
-    },
-  },
-});
-
-export const { setUrl, setScreenshotSrc } = browserSlice.actions;
-
-export default browserSlice.reducer;

frontend/src/store.ts

@@ -1,12 +1,10 @@
 import { combineReducers, configureStore } from "@reduxjs/toolkit";
 import agentReducer from "./state/agent-slice";
-import browserReducer from "./state/browser-slice";
 import { jupyterReducer } from "./state/jupyter-slice";
 import microagentManagementReducer from "./state/microagent-management-slice";
 import eventMessageReducer from "./state/event-message-slice";
 
 export const rootReducer = combineReducers({
-  browser: browserReducer,
   agent: agentReducer,
   jupyter: jupyterReducer,
   microagentManagement: microagentManagementReducer,

frontend/src/stores/browser-store.ts

@@ -0,0 +1,26 @@
+import { create } from "zustand";
+
+interface BrowserState {
+  // URL of browser window (placeholder for now, will be replaced with the actual URL later)
+  url: string;
+  // Base64-encoded screenshot of browser window (placeholder for now, will be replaced with the actual screenshot later)
+  screenshotSrc: string;
+}
+
+interface BrowserStore extends BrowserState {
+  setUrl: (url: string) => void;
+  setScreenshotSrc: (screenshotSrc: string) => void;
+  reset: () => void;
+}
+
+const initialState: BrowserState = {
+  url: "https://github.com/All-Hands-AI/OpenHands",
+  screenshotSrc: "",
+};
+
+export const useBrowserStore = create<BrowserStore>((set) => ({
+  ...initialState,
+  setUrl: (url: string) => set({ url }),
+  setScreenshotSrc: (screenshotSrc: string) => set({ screenshotSrc }),
+  reset: () => set(initialState),
+}));