refactor(frontend): migration of command-slice.ts to zustand (#11003)

frontend/__tests__/components/terminal/terminal.test.tsx

@@ -1,17 +1,14 @@
 import { act, screen } from "@testing-library/react";
 import { renderWithProviders } from "test-utils";
 import { vi, describe, afterEach, it, expect } from "vitest";
-import { Command, appendInput, appendOutput } from "#/state/command-slice";
+import { Command, useCommandStore } from "#/state/command-store";
 import Terminal from "#/components/features/terminal/terminal";
 
-const renderTerminal = (commands: Command[] = []) =>
-  renderWithProviders(<Terminal />, {
-    preloadedState: {
-      cmd: {
-        commands,
-      },
-    },
-  });
+const renderTerminal = (commands: Command[] = []) => {
+  // Set initial commands in Zustand store
+  useCommandStore.setState({ commands });
+  return renderWithProviders(<Terminal />);
+};
 
 describe.skip("Terminal", () => {
   global.ResizeObserver = vi.fn().mockImplementation(() => ({
@@ -58,25 +55,25 @@ describe.skip("Terminal", () => {
   });
 
   it("should write commands to the terminal", () => {
-    const { store } = renderTerminal();
+    renderTerminal();
 
     act(() => {
-      store.dispatch(appendInput("echo Hello"));
-      store.dispatch(appendOutput("Hello"));
+      useCommandStore.getState().appendInput("echo Hello");
+      useCommandStore.getState().appendOutput("Hello");
     });
 
     expect(mockTerminal.writeln).toHaveBeenNthCalledWith(1, "echo Hello");
     expect(mockTerminal.writeln).toHaveBeenNthCalledWith(2, "Hello");
 
     act(() => {
-      store.dispatch(appendInput("echo World"));
+      useCommandStore.getState().appendInput("echo World");
     });
 
     expect(mockTerminal.writeln).toHaveBeenNthCalledWith(3, "echo World");
   });
 
   it("should load and write commands to the terminal", () => {
-    const { store } = renderTerminal([
+    renderTerminal([
       { type: "input", content: "echo Hello" },
       { type: "output", content: "Hello" },
     ]);
@@ -85,17 +82,17 @@ describe.skip("Terminal", () => {
     expect(mockTerminal.writeln).toHaveBeenNthCalledWith(2, "Hello");
 
     act(() => {
-      store.dispatch(appendInput("echo Hello"));
+      useCommandStore.getState().appendInput("echo Hello");
     });
 
     expect(mockTerminal.writeln).toHaveBeenNthCalledWith(3, "echo Hello");
   });
 
   it("should end the line with a dollar sign after writing a command", () => {
-    const { store } = renderTerminal();
+    renderTerminal();
 
     act(() => {
-      store.dispatch(appendInput("echo Hello"));
+      useCommandStore.getState().appendInput("echo Hello");
     });
 
     expect(mockTerminal.writeln).toHaveBeenCalledWith("echo Hello");

frontend/__tests__/hooks/use-terminal.test.tsx

@@ -1,7 +1,7 @@
 import { beforeAll, describe, expect, it, vi } from "vitest";
 import { afterEach } from "node:test";
 import { useTerminal } from "#/hooks/use-terminal";
-import { Command } from "#/state/command-slice";
+import { Command, useCommandStore } from "#/state/command-store";
 import { AgentState } from "#/types/agent-state";
 import { renderWithProviders } from "../../test-utils";
 
@@ -19,10 +19,10 @@ interface TestTerminalComponentProps {
   commands: Command[];
 }
 
-function TestTerminalComponent({
-  commands,
-}: TestTerminalComponentProps) {
-  const ref = useTerminal({ commands });
+function TestTerminalComponent({ commands }: TestTerminalComponentProps) {
+  // Set commands in Zustand store
+  useCommandStore.setState({ commands });
+  const ref = useTerminal();
   return <div ref={ref} />;
 }
 
@@ -60,7 +60,6 @@ describe("useTerminal", () => {
     renderWithProviders(<TestTerminalComponent commands={[]} />, {
       preloadedState: {
         agent: { curAgentState: AgentState.RUNNING },
-        cmd: { commands: [] },
       },
     });
   });
@@ -74,7 +73,6 @@ describe("useTerminal", () => {
     renderWithProviders(<TestTerminalComponent commands={commands} />, {
       preloadedState: {
         agent: { curAgentState: AgentState.RUNNING },
-        cmd: { commands },
       },
     });
 
@@ -94,17 +92,11 @@ describe("useTerminal", () => {
       { content: secret, type: "output" },
     ];
 
-    renderWithProviders(
-      <TestTerminalComponent
-        commands={commands}
-      />,
-      {
-        preloadedState: {
-          agent: { curAgentState: AgentState.RUNNING },
-          cmd: { commands },
-        },
+    renderWithProviders(<TestTerminalComponent commands={commands} />, {
+      preloadedState: {
+        agent: { curAgentState: AgentState.RUNNING },
       },
-    );
+    });
 
     // This test is no longer relevant as secrets filtering has been removed
   });

frontend/__tests__/services/actions.test.tsx

@@ -1,6 +1,8 @@
 import { describe, it, expect, vi, beforeEach } from "vitest";
 import ActionType from "#/types/action-type";
 import { ActionMessage } from "#/types/message";
+import { setMetrics } from "#/state/metrics-slice";
+import { appendSecurityAnalyzerInput } from "#/state/security-analyzer-slice";
 
 // Mock the store and actions
 const mockDispatch = vi.fn();
@@ -13,14 +15,26 @@ vi.mock("#/store", () => ({
   },
 }));
 
-vi.mock("#/state/command-slice", () => ({
-  appendInput: mockAppendInput,
+vi.mock("#/state/command-store", () => ({
+  useCommandStore: {
+    getState: () => ({
+      appendInput: mockAppendInput,
+    }),
+  },
 }));
 
 vi.mock("#/state/jupyter-slice", () => ({
   appendJupyterInput: mockAppendJupyterInput,
 }));
 
+vi.mock("#/state/metrics-slice", () => ({
+  setMetrics: vi.fn(),
+}));
+
+vi.mock("#/state/security-analyzer-slice", () => ({
+  appendSecurityAnalyzerInput: vi.fn(),
+}));
+
 describe("handleActionMessage", () => {
   beforeEach(() => {
     // Clear all mocks before each test
@@ -45,7 +59,8 @@ describe("handleActionMessage", () => {
     handleActionMessage(runAction);
 
     // Check that appendInput was called with the command
-    expect(mockDispatch).toHaveBeenCalledWith(mockAppendInput("ls -la"));
+    expect(mockAppendInput).toHaveBeenCalledWith("ls -la");
+    expect(mockDispatch).not.toHaveBeenCalled();
     expect(mockAppendJupyterInput).not.toHaveBeenCalled();
   });
 
@@ -59,7 +74,8 @@ describe("handleActionMessage", () => {
       args: {
         code: "print('Hello from Jupyter!')",
       },
-      message: "Running Python code interactively: print('Hello from Jupyter!')",
+      message:
+        "Running Python code interactively: print('Hello from Jupyter!')",
       timestamp: "2023-01-01T00:00:00Z",
     };
 
@@ -67,7 +83,9 @@ describe("handleActionMessage", () => {
     handleActionMessage(ipythonAction);
 
     // Check that appendJupyterInput was called with the code
-    expect(mockDispatch).toHaveBeenCalledWith(mockAppendJupyterInput("print('Hello from Jupyter!')"));
+    expect(mockDispatch).toHaveBeenCalledWith(
+      mockAppendJupyterInput("print('Hello from Jupyter!')"),
+    );
     expect(mockAppendInput).not.toHaveBeenCalled();
   });
 
@@ -89,7 +107,9 @@ describe("handleActionMessage", () => {
     // Handle the action
     handleActionMessage(hiddenAction);
 
-    // Check that nothing was dispatched
+    // Check that nothing was dispatched or called
     expect(mockDispatch).not.toHaveBeenCalled();
+    expect(mockAppendInput).not.toHaveBeenCalled();
+    expect(mockAppendJupyterInput).not.toHaveBeenCalled();
   });
 });

frontend/src/components/features/terminal/terminal.tsx

@@ -7,14 +7,11 @@ import { cn } from "#/utils/utils";
 import { WaitingForRuntimeMessage } from "../chat/waiting-for-runtime-message";
 
 function Terminal() {
-  const { commands } = useSelector((state: RootState) => state.cmd);
   const { curAgentState } = useSelector((state: RootState) => state.agent);
 
   const isRuntimeInactive = RUNTIME_INACTIVE_STATES.includes(curAgentState);
 
-  const ref = useTerminal({
-    commands,
-  });
+  const ref = useTerminal();
 
   return (
     <div className="h-full flex flex-col rounded-xl">

frontend/src/hooks/use-terminal.ts

@@ -2,26 +2,18 @@ import { FitAddon } from "@xterm/addon-fit";
 import { Terminal } from "@xterm/xterm";
 import React from "react";
 import { useSelector } from "react-redux";
-import { Command } from "#/state/command-slice";
-import { RootState } from "#/store";
+import { Command, useCommandStore } from "#/state/command-store";
 import { RUNTIME_INACTIVE_STATES } from "#/types/agent-state";
 import { useWsClient } from "#/context/ws-client-provider";
 import { getTerminalCommand } from "#/services/terminal-service";
 import { parseTerminalOutput } from "#/utils/parse-terminal-output";
+import { RootState } from "#/store";
 
 /*
   NOTE: Tests for this hook are indirectly covered by the tests for the XTermTerminal component.
   The reason for this is that the hook exposes a ref that requires a DOM element to be rendered.
 */
 
-interface UseTerminalConfig {
-  commands: Command[];
-}
-
-const DEFAULT_TERMINAL_CONFIG: UseTerminalConfig = {
-  commands: [],
-};
-
 const renderCommand = (
   command: Command,
   terminal: Terminal,
@@ -44,11 +36,10 @@ const renderCommand = (
 // This ensures terminal history is preserved when navigating away and back
 const persistentLastCommandIndex = { current: 0 };
 
-export const useTerminal = ({
-  commands,
-}: UseTerminalConfig = DEFAULT_TERMINAL_CONFIG) => {
+export const useTerminal = () => {
   const { send } = useWsClient();
   const { curAgentState } = useSelector((state: RootState) => state.agent);
+  const commands = useCommandStore((state) => state.commands);
   const terminal = React.useRef<Terminal | null>(null);
   const fitAddon = React.useRef<FitAddon | null>(null);
   const ref = React.useRef<HTMLDivElement>(null);

frontend/src/routes/conversation.tsx

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

frontend/src/services/__tests__/actions.test.ts

@@ -27,10 +27,6 @@ vi.mock("#/state/status-store", () => ({
   },
 }));
 
-vi.mock("#/state/chat-slice", () => ({
-  addErrorMessage: vi.fn(),
-}));
-
 vi.mock("#/utils/error-handler", () => ({
   trackError: vi.fn(),
 }));
@@ -94,7 +90,7 @@ describe("handleStatusMessage", () => {
     expect(queryClient.invalidateQueries).not.toHaveBeenCalled();
   });
 
-  it("should dispatch addErrorMessage for error messages", () => {
+  it("should call trackError for error messages", () => {
     // Create an error status message
     const statusMessage: StatusMessage = {
       status_update: true,
@@ -113,6 +109,9 @@ describe("handleStatusMessage", () => {
       metadata: { msgId: "ERROR_ID" },
     });
 
+    // Verify that store.dispatch was not called
+    expect(store.dispatch).not.toHaveBeenCalled();
+
     // Verify that queryClient.invalidateQueries was not called
     expect(queryClient.invalidateQueries).not.toHaveBeenCalled();
   });

frontend/src/services/actions.ts

@@ -10,7 +10,7 @@ import {
   StatusMessage,
 } from "#/types/message";
 import { handleObservationMessage } from "./observations";
-import { appendInput } from "#/state/command-slice";
+import { useCommandStore } from "#/state/command-store";
 import { appendJupyterInput } from "#/state/jupyter-slice";
 import { queryClient } from "#/query-client-config";
 
@@ -30,7 +30,7 @@ export function handleActionMessage(message: ActionMessage) {
   }
 
   if (message.action === ActionType.RUN) {
-    store.dispatch(appendInput(message.args.command));
+    useCommandStore.getState().appendInput(message.args.command);
   }
 
   if (message.action === ActionType.RUN_IPYTHON) {

frontend/src/services/observations.ts

@@ -2,7 +2,7 @@ import { setCurrentAgentState } from "#/state/agent-slice";
 import { setUrl, setScreenshotSrc } from "#/state/browser-slice";
 import store from "#/store";
 import { ObservationMessage } from "#/types/message";
-import { appendOutput } from "#/state/command-slice";
+import { useCommandStore } from "#/state/command-store";
 import { appendJupyterOutput } from "#/state/jupyter-slice";
 import ObservationType from "#/types/observation-type";
 
@@ -19,7 +19,7 @@ export function handleObservationMessage(message: ObservationMessage) {
         content = `${head}\r\n\n... (truncated ${message.content.length - 5000} characters) ...\r\n\n${tail}`;
       }
 
-      store.dispatch(appendOutput(content));
+      useCommandStore.getState().appendOutput(content);
       break;
     }
     case ObservationType.RUN_IPYTHON:

frontend/src/state/command-slice.ts

@@ -1,31 +0,0 @@
-import { createSlice } from "@reduxjs/toolkit";
-
-export type Command = {
-  content: string;
-  type: "input" | "output";
-};
-
-const initialCommands: Command[] = [];
-
-export const commandSlice = createSlice({
-  name: "command",
-  initialState: {
-    commands: initialCommands,
-  },
-  reducers: {
-    appendInput: (state, action) => {
-      state.commands.push({ content: action.payload, type: "input" });
-    },
-    appendOutput: (state, action) => {
-      state.commands.push({ content: action.payload, type: "output" });
-    },
-    clearTerminal: (state) => {
-      state.commands = [];
-    },
-  },
-});
-
-export const { appendInput, appendOutput, clearTerminal } =
-  commandSlice.actions;
-
-export default commandSlice.reducer;

frontend/src/state/command-store.ts

@@ -0,0 +1,26 @@
+import { create } from "zustand";
+
+export type Command = {
+  content: string;
+  type: "input" | "output";
+};
+
+interface CommandState {
+  commands: Command[];
+  appendInput: (content: string) => void;
+  appendOutput: (content: string) => void;
+  clearTerminal: () => void;
+}
+
+export const useCommandStore = create<CommandState>((set) => ({
+  commands: [],
+  appendInput: (content: string) =>
+    set((state) => ({
+      commands: [...state.commands, { content, type: "input" }],
+    })),
+  appendOutput: (content: string) =>
+    set((state) => ({
+      commands: [...state.commands, { content, type: "output" }],
+    })),
+  clearTerminal: () => set({ commands: [] }),
+}));

frontend/src/store.ts

@@ -2,7 +2,6 @@ import { combineReducers, configureStore } from "@reduxjs/toolkit";
 import agentReducer from "./state/agent-slice";
 import browserReducer from "./state/browser-slice";
 import fileStateReducer from "./state/file-state-slice";
-import commandReducer from "./state/command-slice";
 import { jupyterReducer } from "./state/jupyter-slice";
 import securityAnalyzerReducer from "./state/security-analyzer-slice";
 import metricsReducer from "./state/metrics-slice";
@@ -13,7 +12,6 @@ import eventMessageReducer from "./state/event-message-slice";
 export const rootReducer = combineReducers({
   fileState: fileStateReducer,
   browser: browserReducer,
-  cmd: commandReducer,
   agent: agentReducer,
   jupyter: jupyterReducer,
   securityAnalyzer: securityAnalyzerReducer,