refactor(frontend): migration of jupyter-slice.ts to zustand (#11019)

frontend/__tests__/components/jupyter/jupyter.test.tsx

@@ -2,8 +2,8 @@ import { render, screen } from "@testing-library/react";
 import { Provider } from "react-redux";
 import { configureStore } from "@reduxjs/toolkit";
 import { JupyterEditor } from "#/components/features/jupyter/jupyter";
-import { jupyterReducer } from "#/state/jupyter-slice";
-import { vi, describe, it, expect } from "vitest";
+import { useJupyterStore } from "#/state/jupyter-store";
+import { vi, describe, it, expect, beforeEach } from "vitest";
 
 describe("JupyterEditor", () => {
   const mockStore = configureStore({
@@ -15,19 +15,20 @@ describe("JupyterEditor", () => {
       code: () => ({}),
       cmd: () => ({}),
       agent: () => ({}),
-      jupyter: jupyterReducer,
       securityAnalyzer: () => ({}),
       status: () => ({}),
     },
-    preloadedState: {
-      jupyter: {
-        cells: Array(20).fill({
-          content: "Test cell content",
-          type: "input",
-          output: "Test output",
-        }),
-      },
-    },
+  });
+
+  beforeEach(() => {
+    // Reset the Zustand store before each test
+    useJupyterStore.setState({
+      cells: Array(20).fill({
+        content: "Test cell content",
+        type: "input",
+        imageUrls: undefined,
+      }),
+    });
   });
 
   it("should have a scrollable container", () => {
@@ -36,7 +37,7 @@ describe("JupyterEditor", () => {
         <div style={{ height: "100vh" }}>
           <JupyterEditor maxWidth={800} />
         </div>
-      </Provider>
+      </Provider>,
     );
 
     const container = screen.getByTestId("jupyter-container");

frontend/__tests__/services/actions.test.tsx

@@ -21,8 +21,12 @@ vi.mock("#/state/command-store", () => ({
   },
 }));
 
-vi.mock("#/state/jupyter-slice", () => ({
-  appendJupyterInput: mockAppendJupyterInput,
+vi.mock("#/state/jupyter-store", () => ({
+  useJupyterStore: {
+    getState: () => ({
+      appendJupyterInput: mockAppendJupyterInput,
+    }),
+  },
 }));
 
 vi.mock("#/state/metrics-slice", () => ({
@@ -81,8 +85,8 @@ describe("handleActionMessage", () => {
     handleActionMessage(ipythonAction);
 
     // Check that appendJupyterInput was called with the code
-    expect(mockDispatch).toHaveBeenCalledWith(
-      mockAppendJupyterInput("print('Hello from Jupyter!')"),
+    expect(mockAppendJupyterInput).toHaveBeenCalledWith(
+      "print('Hello from Jupyter!')",
     );
     expect(mockAppendInput).not.toHaveBeenCalled();
   });

frontend/src/components/features/jupyter/jupyter-cell.tsx

@@ -1,5 +1,5 @@
 import React from "react";
-import { Cell } from "#/state/jupyter-slice";
+import { Cell } from "#/state/jupyter-store";
 import { JupyterLine, parseCellContent } from "#/utils/parse-cell-content";
 import { JupytrerCellInput } from "./jupyter-cell-input";
 import { JupyterCellOutput } from "./jupyter-cell-output";

frontend/src/components/features/jupyter/jupyter.tsx

@@ -9,13 +9,14 @@ import { RUNTIME_INACTIVE_STATES } from "#/types/agent-state";
 import { I18nKey } from "#/i18n/declaration";
 import JupyterLargeIcon from "#/icons/jupyter-large.svg?react";
 import { WaitingForRuntimeMessage } from "../chat/waiting-for-runtime-message";
+import { useJupyterStore } from "#/state/jupyter-store";
 
 interface JupyterEditorProps {
   maxWidth: number;
 }
 
 export function JupyterEditor({ maxWidth }: JupyterEditorProps) {
-  const cells = useSelector((state: RootState) => state.jupyter?.cells ?? []);
+  const cells = useJupyterStore((state) => state.cells);
   const { curAgentState } = useSelector((state: RootState) => state.agent);
 
   const jupyterRef = React.useRef<HTMLDivElement>(null);

frontend/src/routes/conversation.tsx

@@ -6,7 +6,7 @@ import { useQueryClient } from "@tanstack/react-query";
 import { useConversationId } from "#/hooks/use-conversation-id";
 import { useCommandStore } from "#/state/command-store";
 import { useEffectOnce } from "#/hooks/use-effect-once";
-import { clearJupyter } from "#/state/jupyter-slice";
+import { useJupyterStore } from "#/state/jupyter-store";
 import { useConversationStore } from "#/state/conversation-store";
 import { setCurrentAgentState } from "#/state/agent-slice";
 import { AgentState } from "#/types/agent-state";
@@ -42,6 +42,7 @@ function AppContent() {
   const dispatch = useDispatch();
   const navigate = useNavigate();
   const clearTerminal = useCommandStore((state) => state.clearTerminal);
+  const clearJupyter = useJupyterStore((state) => state.clearJupyter);
   const queryClient = useQueryClient();
 
   // Fetch batch feedback data when conversation is loaded
@@ -86,14 +87,14 @@ function AppContent() {
 
   React.useEffect(() => {
     clearTerminal();
-    dispatch(clearJupyter());
+    clearJupyter();
     resetConversationState();
     dispatch(setCurrentAgentState(AgentState.LOADING));
   }, [conversationId, clearTerminal, resetConversationState]);
 
   useEffectOnce(() => {
     clearTerminal();
-    dispatch(clearJupyter());
+    clearJupyter();
     resetConversationState();
     dispatch(setCurrentAgentState(AgentState.LOADING));
   });

frontend/src/services/actions.ts

@@ -1,7 +1,6 @@
 import { trackError } from "#/utils/error-handler";
 import useMetricsStore from "#/stores/metrics-store";
 import { useStatusStore } from "#/state/status-store";
-import store from "#/store";
 import ActionType from "#/types/action-type";
 import {
   ActionMessage,
@@ -9,8 +8,8 @@ import {
   StatusMessage,
 } from "#/types/message";
 import { handleObservationMessage } from "./observations";
+import { useJupyterStore } from "#/state/jupyter-store";
 import { useCommandStore } from "#/state/command-store";
-import { appendJupyterInput } from "#/state/jupyter-slice";
 import { queryClient } from "#/query-client-config";
 import {
   ActionSecurityRisk,
@@ -37,7 +36,7 @@ export function handleActionMessage(message: ActionMessage) {
   }
 
   if (message.action === ActionType.RUN_IPYTHON) {
-    store.dispatch(appendJupyterInput(message.args.code));
+    useJupyterStore.getState().appendJupyterInput(message.args.code);
   }
 
   if ("args" in message && "security_risk" in message.args) {

frontend/src/services/observations.ts

@@ -1,8 +1,8 @@
 import { setCurrentAgentState } from "#/state/agent-slice";
 import store from "#/store";
 import { ObservationMessage } from "#/types/message";
+import { useJupyterStore } from "#/state/jupyter-store";
 import { useCommandStore } from "#/state/command-store";
-import { appendJupyterOutput } from "#/state/jupyter-slice";
 import ObservationType from "#/types/observation-type";
 import { useBrowserStore } from "#/stores/browser-store";
 
@@ -23,14 +23,12 @@ export function handleObservationMessage(message: ObservationMessage) {
       break;
     }
     case ObservationType.RUN_IPYTHON:
-      store.dispatch(
-        appendJupyterOutput({
-          content: message.content,
-          imageUrls: Array.isArray(message.extras?.image_urls)
-            ? message.extras.image_urls
-            : undefined,
-        }),
-      );
+      useJupyterStore.getState().appendJupyterOutput({
+        content: message.content,
+        imageUrls: Array.isArray(message.extras?.image_urls)
+          ? message.extras.image_urls
+          : undefined,
+      });
       break;
     case ObservationType.BROWSE:
     case ObservationType.BROWSE_INTERACTIVE:

frontend/src/state/jupyter-slice.ts

@@ -1,37 +0,0 @@
-import { createSlice } from "@reduxjs/toolkit";
-
-export type Cell = {
-  content: string;
-  type: "input" | "output";
-  imageUrls?: string[];
-};
-
-const initialCells: Cell[] = [];
-
-export const jupyterSlice = createSlice({
-  name: "jupyter",
-  initialState: {
-    cells: initialCells,
-  },
-  reducers: {
-    appendJupyterInput: (state, action) => {
-      state.cells.push({ content: action.payload, type: "input" });
-    },
-    appendJupyterOutput: (state, action) => {
-      state.cells.push({
-        content: action.payload.content,
-        type: "output",
-        imageUrls: action.payload.imageUrls,
-      });
-    },
-    clearJupyter: (state) => {
-      state.cells = [];
-    },
-  },
-});
-
-export const { appendJupyterInput, appendJupyterOutput, clearJupyter } =
-  jupyterSlice.actions;
-
-export const jupyterReducer = jupyterSlice.reducer;
-export default jupyterReducer;

frontend/src/state/jupyter-store.ts

@@ -0,0 +1,40 @@
+import { create } from "zustand";
+
+export type Cell = {
+  content: string;
+  type: "input" | "output";
+  imageUrls?: string[];
+};
+
+interface JupyterState {
+  cells: Cell[];
+  appendJupyterInput: (content: string) => void;
+  appendJupyterOutput: (payload: {
+    content: string;
+    imageUrls?: string[];
+  }) => void;
+  clearJupyter: () => void;
+}
+
+export const useJupyterStore = create<JupyterState>((set) => ({
+  cells: [],
+  appendJupyterInput: (content: string) =>
+    set((state) => ({
+      cells: [...state.cells, { content, type: "input" }],
+    })),
+  appendJupyterOutput: (payload: { content: string; imageUrls?: string[] }) =>
+    set((state) => ({
+      cells: [
+        ...state.cells,
+        {
+          content: payload.content,
+          type: "output",
+          imageUrls: payload.imageUrls,
+        },
+      ],
+    })),
+  clearJupyter: () =>
+    set(() => ({
+      cells: [],
+    })),
+}));

frontend/src/store.ts

@@ -1,10 +1,8 @@
 import { combineReducers, configureStore } from "@reduxjs/toolkit";
 import agentReducer from "./state/agent-slice";
-import { jupyterReducer } from "./state/jupyter-slice";
 
 export const rootReducer = combineReducers({
   agent: agentReducer,
-  jupyter: jupyterReducer,
 });
 
 const store = configureStore({