Migrate

frontend/__tests__/components/browser.test.tsx

@@ -1,4 +1,7 @@
 import { describe, it, expect, afterEach, vi } from "vitest";
+import { screen } from "@testing-library/react";
+import { renderWithProviders } from "../../test-utils";
+import { BrowserPanel } from "#/components/features/browser/browser";
 
 // Mock useParams before importing components
 vi.mock("react-router", async () => {
@@ -23,38 +26,37 @@ vi.mock("react-i18next", async () => {
   };
 });
 
-import { screen } from "@testing-library/react";
-import { renderWithProviders } from "../../test-utils";
-import { BrowserPanel } from "#/components/features/browser/browser";
+const { useBrowserMock } = vi.hoisted(() => ({
+  useBrowserMock: vi.fn(),
+}));
+
+vi.mock("#/hooks/state/use-browser", async () => ({
+  useBrowser: useBrowserMock,
+}));
 
 describe("Browser", () => {
   afterEach(() => {
     vi.clearAllMocks();
   });
+
   it("renders a message if no screenshotSrc is provided", () => {
-    renderWithProviders(<BrowserPanel />, {
-      preloadedState: {
-        browser: {
-          url: "https://example.com",
-          screenshotSrc: "",
-        },
-      },
+    useBrowserMock.mockReturnValue({
+      url: "https://example.com",
+      screenshotSrc: "",
     });
+    renderWithProviders(<BrowserPanel />);
 
     // i18n empty message key
     expect(screen.getByText("BROWSER$NO_PAGE_LOADED")).toBeInTheDocument();
   });
 
   it("renders the url and a screenshot", () => {
-    renderWithProviders(<BrowserPanel />, {
-      preloadedState: {
-        browser: {
-          url: "https://example.com",
-          screenshotSrc:
-            "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mN0uGvyHwAFCAJS091fQwAAAABJRU5ErkJggg==",
-        },
-      },
+    useBrowserMock.mockReturnValue({
+      url: "https://example.com",
+      screenshotSrc:
+        "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mN0uGvyHwAFCAJS091fQwAAAABJRU5ErkJggg==",
     });
+    renderWithProviders(<BrowserPanel />);
 
     expect(screen.getByText("https://example.com")).toBeInTheDocument();
     expect(screen.getByAltText(/browser screenshot/i)).toBeInTheDocument();

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

@@ -1,12 +1,9 @@
-import { useSelector } from "react-redux";
-import { RootState } from "#/store";
 import { BrowserSnapshot } from "./browser-snapshot";
 import { EmptyBrowserMessage } from "./empty-browser-message";
+import { useBrowser } from "#/hooks/state/use-browser";
 
 export function BrowserPanel() {
-  const { url, screenshotSrc } = useSelector(
-    (state: RootState) => state.browser,
-  );
+  const { url, screenshotSrc } = useBrowser();
 
   const imgSrc =
     screenshotSrc && screenshotSrc.startsWith("data:image/png;base64,")

frontend/src/context/ws-client-provider.tsx

@@ -9,6 +9,11 @@ import {
   AssistantMessageAction,
   UserMessageAction,
 } from "#/types/core/actions";
+import {
+  BrowseInteractiveObservation,
+  BrowseObservation,
+} from "#/types/core/observations";
+import { useBrowser } from "#/hooks/state/use-browser";
 
 const isOpenHandsEvent = (event: unknown): event is OpenHandsParsedEvent =>
   typeof event === "object" &&
@@ -39,6 +44,14 @@ const isMessageAction = (
 ): event is UserMessageAction | AssistantMessageAction =>
   isUserMessage(event) || isAssistantMessage(event);
 
+const isBrowserObservation = (
+  event: OpenHandsParsedEvent,
+): event is BrowseObservation | BrowseInteractiveObservation =>
+  "source" in event &&
+  "observation" in event &&
+  (event.observation === "browse" ||
+    event.observation === "browse_interactive");
+
 export enum WsClientProviderStatus {
   CONNECTED,
   DISCONNECTED,
@@ -104,6 +117,7 @@ export function WsClientProvider({
   conversationId,
   children,
 }: React.PropsWithChildren<WsClientProviderProps>) {
+  const { setScreenshotSrc, setUrl } = useBrowser();
   const sioRef = React.useRef<Socket | null>(null);
   const [status, setStatus] = React.useState(
     WsClientProviderStatus.DISCONNECTED,
@@ -134,6 +148,11 @@ export function WsClientProvider({
       lastEventRef.current = event;
     }
 
+    if (isOpenHandsEvent(event) && isBrowserObservation(event)) {
+      if (event.extras.url) setUrl(event.extras.url);
+      if (event.extras.screenshot) setScreenshotSrc(event.extras.screenshot);
+    }
+
     handleAssistantMessage(event);
   }
 

frontend/src/hooks/state/use-browser.ts

@@ -0,0 +1,48 @@
+import { useQueryClient } from "@tanstack/react-query";
+import React from "react";
+
+const BROWSER_KEY = ["_STATE", "browser"];
+
+interface BrowserData {
+  url: string;
+  screenshotSrc: string;
+}
+
+export const DEFAULT_BROWSER_DATA: BrowserData = {
+  url: "https://github.com/All-Hands-AI/OpenHands",
+  screenshotSrc: "",
+};
+
+export const useBrowser = () => {
+  const queryClient = useQueryClient();
+
+  const setUrl = React.useCallback(
+    (url: string) => {
+      queryClient.setQueryData<BrowserData>(BROWSER_KEY, (old) => ({
+        ...(old || DEFAULT_BROWSER_DATA),
+        url,
+      }));
+    },
+    [queryClient],
+  );
+
+  const setScreenshotSrc = React.useCallback(
+    (screenshotSrc: string) => {
+      queryClient.setQueryData<BrowserData>(BROWSER_KEY, (old) => ({
+        ...(old || DEFAULT_BROWSER_DATA),
+        screenshotSrc,
+      }));
+    },
+    [queryClient],
+  );
+
+  const { url, screenshotSrc } =
+    queryClient.getQueryData<BrowserData>(BROWSER_KEY) || DEFAULT_BROWSER_DATA;
+
+  return {
+    url,
+    screenshotSrc,
+    setUrl,
+    setScreenshotSrc,
+  };
+};

frontend/src/hooks/use-end-session.ts

@@ -1,13 +1,10 @@
 import { useDispatch } from "react-redux";
 import { useNavigate } from "react-router";
-import {
-  initialState as browserInitialState,
-  setScreenshotSrc,
-  setUrl,
-} from "#/state/browser-slice";
 import { clearSelectedRepository } from "#/state/initial-query-slice";
+import { DEFAULT_BROWSER_DATA, useBrowser } from "./state/use-browser";
 
 export const useEndSession = () => {
+  const { setScreenshotSrc, setUrl } = useBrowser();
   const navigate = useNavigate();
   const dispatch = useDispatch();
 
@@ -18,8 +15,8 @@ export const useEndSession = () => {
     dispatch(clearSelectedRepository());
 
     // Reset browser state to initial values
-    dispatch(setUrl(browserInitialState.url));
-    dispatch(setScreenshotSrc(browserInitialState.screenshotSrc));
+    setUrl(DEFAULT_BROWSER_DATA.url);
+    setScreenshotSrc(DEFAULT_BROWSER_DATA.screenshotSrc);
 
     navigate("/");
   };

frontend/src/services/observations.ts

@@ -1,5 +1,4 @@
 import { setCurrentAgentState } from "#/state/agent-slice";
-import { setUrl, setScreenshotSrc } from "#/state/browser-slice";
 import store from "#/store";
 import { ObservationMessage } from "#/types/message";
 import { AgentState } from "#/types/agent-state";
@@ -29,15 +28,6 @@ export function handleObservationMessage(message: ObservationMessage) {
       // FIXME: render this as markdown
       store.dispatch(appendJupyterOutput(message.content));
       break;
-    case ObservationType.BROWSE:
-    case ObservationType.BROWSE_INTERACTIVE:
-      if (message.extras?.screenshot) {
-        store.dispatch(setScreenshotSrc(message.extras?.screenshot));
-      }
-      if (message.extras?.url) {
-        store.dispatch(setUrl(message.extras.url));
-      }
-      break;
     case ObservationType.AGENT_STATE_CHANGED:
       store.dispatch(setCurrentAgentState(message.extras.agent_state));
       break;

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,6 +1,5 @@
 import { combineReducers, configureStore } from "@reduxjs/toolkit";
 import agentReducer from "./state/agent-slice";
-import browserReducer from "./state/browser-slice";
 import chatReducer from "./state/chat-slice";
 import codeReducer from "./state/code-slice";
 import fileStateReducer from "./state/file-state-slice";
@@ -14,7 +13,6 @@ import metricsReducer from "./state/metrics-slice";
 export const rootReducer = combineReducers({
   fileState: fileStateReducer,
   initialQuery: initialQueryReducer,
-  browser: browserReducer,
   chat: chatReducer,
   code: codeReducer,
   cmd: commandReducer,