Migrate metrics slice from Redux to React Query and fix test setup

frontend/__tests__/services/actions.test.ts

@@ -1,9 +1,12 @@
 import { describe, it, expect, vi, beforeEach } from "vitest";
-import { handleStatusMessage, handleActionMessage } from "#/services/actions";
+import { handleActionMessage } from "#/services/actions";
+import { handleStatusMessage } from "#/services/status-service";
 import store from "#/store";
 import { trackError } from "#/utils/error-handler";
 import ActionType from "#/types/action-type";
 import { ActionMessage } from "#/types/message";
+import { queryClient } from "#/entry.client";
+import { statusKeys } from "#/hooks/query/use-status";
 
 // Mock dependencies
 vi.mock("#/utils/error-handler", () => ({
@@ -16,13 +19,19 @@ vi.mock("#/store", () => ({
   },
 }));
 
+vi.mock("#/entry.client", () => ({
+  queryClient: {
+    setQueryData: vi.fn(),
+  },
+}));
+
 describe("Actions Service", () => {
   beforeEach(() => {
     vi.clearAllMocks();
   });
 
   describe("handleStatusMessage", () => {
-    it("should dispatch info messages to status state", () => {
+    it("should update status message in React Query", () => {
       const message = {
         type: "info",
         message: "Runtime is not available",
@@ -32,9 +41,10 @@ describe("Actions Service", () => {
 
       handleStatusMessage(message);
 
-      expect(store.dispatch).toHaveBeenCalledWith(expect.objectContaining({
-        payload: message,
-      }));
+      expect(queryClient.setQueryData).toHaveBeenCalledWith(
+        statusKeys.current(),
+        message
+      );
     });
 
     it("should log error messages and display them in chat", () => {

frontend/__tests__/utils/error-handler.test.ts

@@ -5,8 +5,8 @@ import {
   showErrorToast,
   showChatError,
 } from "#/utils/error-handler";
-import * as Actions from "#/services/actions";
 import * as CustomToast from "#/utils/custom-toast-handlers";
+import * as StatusService from "#/services/status-service";
 
 vi.mock("posthog-js", () => ({
   default: {
@@ -14,7 +14,7 @@ vi.mock("posthog-js", () => ({
   },
 }));
 
-vi.mock("#/services/actions", () => ({
+vi.mock("#/services/status-service", () => ({
   handleStatusMessage: vi.fn(),
 }));
 
@@ -177,7 +177,7 @@ describe("Error Handler", () => {
       );
 
       // Verify error message was shown in chat
-      expect(Actions.handleStatusMessage).toHaveBeenCalledWith({
+      expect(StatusService.handleStatusMessage).toHaveBeenCalledWith({
         type: "error",
         message: "Chat error",
         id: "123",

frontend/src/components/features/controls/agent-status-bar.tsx

@@ -11,6 +11,7 @@ import {
 } from "#/context/ws-client-provider";
 import { useNotification } from "#/hooks/useNotification";
 import { browserTab } from "#/utils/browser-tab";
+import { useStatusMessage } from "#/hooks/query/use-status";
 
 const notificationStates = [
   AgentState.AWAITING_USER_INPUT,
@@ -21,7 +22,7 @@ const notificationStates = [
 export function AgentStatusBar() {
   const { t, i18n } = useTranslation();
   const { curAgentState } = useSelector((state: RootState) => state.agent);
-  const { curStatusMessage } = useSelector((state: RootState) => state.status);
+  const { statusMessage: curStatusMessage } = useStatusMessage();
   const { status } = useWsClient();
   const { notify } = useNotification();
 

frontend/src/components/features/conversation-panel/conversation-card.tsx

@@ -1,5 +1,4 @@
 import React from "react";
-import { useSelector } from "react-redux";
 import posthog from "posthog-js";
 import { formatTimeDelta } from "#/utils/format-time-delta";
 import { ConversationRepoLink } from "./conversation-repo-link";
@@ -11,7 +10,7 @@ import { EllipsisButton } from "./ellipsis-button";
 import { ConversationCardContextMenu } from "./conversation-card-context-menu";
 import { cn } from "#/utils/utils";
 import { BaseModal } from "../../shared/modals/base-modal/base-modal";
-import { RootState } from "#/store";
+import { useMetrics } from "#/hooks/query/use-metrics";
 
 interface ConversationCardProps {
   onClick?: () => void;
@@ -45,8 +44,8 @@ export function ConversationCard({
   const [metricsModalVisible, setMetricsModalVisible] = React.useState(false);
   const inputRef = React.useRef<HTMLInputElement>(null);
 
-  // Subscribe to metrics data from Redux store
-  const metrics = useSelector((state: RootState) => state.metrics);
+  // Get metrics data from React Query
+  const { metrics } = useMetrics();
 
   const handleBlur = () => {
     if (inputRef.current?.value) {

frontend/src/hooks/query/use-metrics.ts

@@ -0,0 +1,51 @@
+import { useQuery, useMutation, useQueryClient } from "@tanstack/react-query";
+
+interface MetricsState {
+  cost: number | null;
+  usage: {
+    prompt_tokens: number;
+    completion_tokens: number;
+    total_tokens: number;
+  } | null;
+}
+
+const initialState: MetricsState = {
+  cost: null,
+  usage: null,
+};
+
+// Define query keys
+export const metricsKeys = {
+  all: ["metrics"] as const,
+  current: () => [...metricsKeys.all, "current"] as const,
+};
+
+// Custom hook to get and update metrics
+export function useMetrics() {
+  const queryClient = useQueryClient();
+
+  // Query to get the current metrics
+  const query = useQuery({
+    queryKey: metricsKeys.current(),
+    queryFn: () =>
+      // Return the cached value or initial value
+      queryClient.getQueryData<MetricsState>(metricsKeys.current()) ||
+      initialState,
+    // Initialize with the default metrics
+    initialData: initialState,
+  });
+
+  // Mutation to update the metrics
+  const mutation = useMutation({
+    mutationFn: (newMetrics: MetricsState) => Promise.resolve(newMetrics),
+    onSuccess: (newMetrics) => {
+      queryClient.setQueryData(metricsKeys.current(), newMetrics);
+    },
+  });
+
+  return {
+    metrics: query.data,
+    setMetrics: mutation.mutate,
+    isLoading: query.isLoading,
+  };
+}

frontend/src/hooks/query/use-status.ts

@@ -0,0 +1,46 @@
+import { useQuery, useMutation, useQueryClient } from "@tanstack/react-query";
+import { StatusMessage } from "#/types/message";
+
+const initialStatusMessage: StatusMessage = {
+  status_update: true,
+  type: "info",
+  id: "",
+  message: "",
+};
+
+// Define query keys
+export const statusKeys = {
+  all: ["status"] as const,
+  current: () => [...statusKeys.all, "current"] as const,
+};
+
+// Custom hook to get and update status message
+export function useStatusMessage() {
+  const queryClient = useQueryClient();
+
+  // Query to get the current status message
+  const query = useQuery({
+    queryKey: statusKeys.current(),
+    queryFn: () =>
+      // Return the cached value or initial value
+      queryClient.getQueryData<StatusMessage>(statusKeys.current()) ||
+      initialStatusMessage,
+    // Initialize with the default status message
+    initialData: initialStatusMessage,
+  });
+
+  // Mutation to update the status message
+  const mutation = useMutation({
+    mutationFn: (newStatusMessage: StatusMessage) =>
+      Promise.resolve(newStatusMessage),
+    onSuccess: (newStatusMessage) => {
+      queryClient.setQueryData(statusKeys.current(), newStatusMessage);
+    },
+  });
+
+  return {
+    statusMessage: query.data,
+    setStatusMessage: mutation.mutate,
+    isLoading: query.isLoading,
+  };
+}

frontend/src/services/actions.ts

@@ -8,8 +8,6 @@ import { trackError } from "#/utils/error-handler";
 import { appendSecurityAnalyzerInput } from "#/state/security-analyzer-slice";
 import { setCode, setActiveFilepath } from "#/state/code-slice";
 import { appendJupyterInput } from "#/state/jupyter-slice";
-import { setCurStatusMessage } from "#/state/status-slice";
-import { setMetrics } from "#/state/metrics-slice";
 import store from "#/store";
 import ActionType from "#/types/action-type";
 import {
@@ -18,6 +16,8 @@ import {
   StatusMessage,
 } from "#/types/message";
 import { handleObservationMessage } from "./observations";
+import { handleStatusMessage } from "./status-service";
+import { updateMetrics } from "./metrics-service";
 import { appendInput } from "#/state/command-slice";
 
 const messageActions = {
@@ -95,7 +95,7 @@ export function handleActionMessage(message: ActionMessage) {
       cost: message.llm_metrics?.accumulated_cost ?? null,
       usage: message.tool_call_metadata?.model_response?.usage ?? null,
     };
-    store.dispatch(setMetrics(metrics));
+    updateMetrics(metrics);
   }
 
   if (message.action === ActionType.RUN) {
@@ -122,26 +122,7 @@ export function handleActionMessage(message: ActionMessage) {
   }
 }
 
-export function handleStatusMessage(message: StatusMessage) {
-  if (message.type === "info") {
-    store.dispatch(
-      setCurStatusMessage({
-        ...message,
-      }),
-    );
-  } else if (message.type === "error") {
-    trackError({
-      message: message.message,
-      source: "chat",
-      metadata: { msgId: message.id },
-    });
-    store.dispatch(
-      addErrorMessage({
-        ...message,
-      }),
-    );
-  }
-}
+// handleStatusMessage has been moved to status-service.ts
 
 export function handleAssistantMessage(message: Record<string, unknown>) {
   if (message.action) {

frontend/src/services/metrics-service.ts

@@ -0,0 +1,15 @@
+import { queryClient } from "#/entry.client";
+import { metricsKeys } from "#/hooks/query/use-metrics";
+
+interface MetricsState {
+  cost: number | null;
+  usage: {
+    prompt_tokens: number;
+    completion_tokens: number;
+    total_tokens: number;
+  } | null;
+}
+
+export function updateMetrics(metrics: MetricsState) {
+  queryClient.setQueryData(metricsKeys.current(), metrics);
+}

frontend/src/services/status-service.ts

@@ -0,0 +1,24 @@
+import { StatusMessage } from "#/types/message";
+import { trackError } from "#/utils/error-handler";
+import { queryClient } from "#/entry.client";
+import { statusKeys } from "#/hooks/query/use-status";
+import { addErrorMessage } from "#/state/chat-slice";
+import store from "#/store";
+
+export function handleStatusMessage(message: StatusMessage) {
+  if (message.type === "info") {
+    // Update the status message using React Query
+    queryClient.setQueryData(statusKeys.current(), message);
+  } else if (message.type === "error") {
+    trackError({
+      message: message.message,
+      source: "chat",
+      metadata: { msgId: message.id },
+    });
+    store.dispatch(
+      addErrorMessage({
+        ...message,
+      }),
+    );
+  }
+}

frontend/src/utils/error-handler.ts

@@ -1,5 +1,5 @@
 import posthog from "posthog-js";
-import { handleStatusMessage } from "#/services/actions";
+import { handleStatusMessage } from "#/services/status-service";
 import { displayErrorToast } from "./custom-toast-handlers";
 
 interface ErrorDetails {

frontend/test-utils.tsx

@@ -8,6 +8,7 @@ import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
 import { I18nextProvider, initReactI18next } from "react-i18next";
 import i18n from "i18next";
 import { vi } from "vitest";
+import { createMemoryRouter, RouterProvider } from "react-router";
 import { AppStore, RootState, rootReducer } from "./src/store";
 import { AuthProvider } from "#/context/auth-context";
 import { ConversationProvider } from "#/context/conversation-context";
@@ -62,6 +63,14 @@ export function renderWithProviders(
     ...renderOptions
   }: ExtendedRenderOptions = {},
 ) {
+  // Create a basic router for testing
+  const router = createMemoryRouter([
+    {
+      path: "/",
+      element: ui,
+    },
+  ]);
+
   function Wrapper({ children }: PropsWithChildren) {
     return (
       <Provider store={store}>
@@ -74,12 +83,43 @@ export function renderWithProviders(
             }
           >
             <ConversationProvider>
-              <I18nextProvider i18n={i18n}>{children}</I18nextProvider>
+              <I18nextProvider i18n={i18n}>
+                {children}
+              </I18nextProvider>
             </ConversationProvider>
           </QueryClientProvider>
         </AuthProvider>
       </Provider>
     );
   }
-  return { store, ...render(ui, { wrapper: Wrapper, ...renderOptions }) };
+  
+  // For components that need RouterProvider, use this approach
+  if (renderOptions.wrapper) {
+    return { store, ...render(ui, { wrapper: Wrapper, ...renderOptions }) };
+  }
+  
+  // For components that need RouterProvider
+  return { 
+    store, 
+    ...render(
+      <Provider store={store}>
+        <AuthProvider initialGithubTokenIsSet>
+          <QueryClientProvider
+            client={
+              new QueryClient({
+                defaultOptions: { queries: { retry: false } },
+              })
+            }
+          >
+            <ConversationProvider>
+              <I18nextProvider i18n={i18n}>
+                <RouterProvider router={router} />
+              </I18nextProvider>
+            </ConversationProvider>
+          </QueryClientProvider>
+        </AuthProvider>
+      </Provider>,
+      renderOptions
+    ) 
+  };
 }