Refactor

frontend/src/components/modals/security/invariant/invariant.tsx

@@ -5,7 +5,7 @@ import { useTranslation } from "react-i18next";
 import { Editor, Monaco } from "@monaco-editor/react";
 import { editor } from "monaco-editor";
 import { Button, Select, SelectItem } from "@nextui-org/react";
-import { useMutation, useQuery } from "@tanstack/react-query";
+import { useMutation } from "@tanstack/react-query";
 import { RootState } from "#/store";
 import {
   ActionSecurityRisk,
@@ -18,6 +18,9 @@ import InvariantLogoIcon from "./assets/logo";
 import { getFormattedDateTime } from "#/utils/gget-formatted-datetime";
 import { downloadJSON } from "#/utils/download-json";
 import InvariantService from "#/api/invariant-service";
+import { useGetPolicy } from "#/hooks/query/use-get-policy";
+import { useGetRiskSeverity } from "#/hooks/query/use-get-risk-severity";
+import { useGetTraces } from "#/hooks/query/use-get-traces";
 
 type SectionType = "logs" | "policy" | "settings";
 
@@ -33,46 +36,26 @@ function SecurityInvariant(): JSX.Element {
 
   const logsRef = React.useRef<HTMLDivElement>(null);
 
-  const { data: fetchedPolicy } = useQuery<string>({
-    queryKey: ["policy"],
-    queryFn: InvariantService.getPolicy,
-  });
-
-  const { data: riskSeverity } = useQuery<number>({
-    queryKey: ["risk_severity"],
-    queryFn: InvariantService.getRiskSeverity,
-  });
+  useGetPolicy({ onSuccess: setPolicy });
 
-  const { data: traces, refetch: exportTraces } = useQuery({
-    queryKey: ["traces"],
-    queryFn: InvariantService.getTraces,
-    enabled: false,
+  useGetRiskSeverity({
+    onSuccess: (riskSeverity) => {
+      setSelectedRisk(
+        riskSeverity === 0
+          ? ActionSecurityRisk.LOW
+          : riskSeverity || ActionSecurityRisk.MEDIUM,
+      );
+    },
   });
 
-  React.useEffect(() => {
-    if (fetchedPolicy) {
-      setPolicy(fetchedPolicy);
-    }
-  }, [fetchedPolicy]);
-
-  React.useEffect(() => {
-    if (traces) {
+  const { refetch: exportTraces } = useGetTraces({
+    onSuccess: (traces) => {
       toast.info(t(I18nKey.INVARIANT$TRACE_EXPORTED_MESSAGE));
 
       const filename = `openhands-trace-${getFormattedDateTime()}.json`;
       downloadJSON(traces, filename);
-    }
-  }, [traces]);
-
-  React.useEffect(() => {
-    if (riskSeverity) {
-      setSelectedRisk(
-        riskSeverity === 0
-          ? ActionSecurityRisk.LOW
-          : riskSeverity || ActionSecurityRisk.MEDIUM,
-      );
-    }
-  }, [riskSeverity]);
+    },
+  });
 
   const { mutate: updatePolicy } = useMutation({
     mutationFn: (variables: { policy: string }) =>

frontend/src/hooks/query/use-get-policy.ts

@@ -0,0 +1,26 @@
+import { useQuery } from "@tanstack/react-query";
+import React from "react";
+import InvariantService from "#/api/invariant-service";
+
+type ResponseData = string;
+
+interface UseGetPolicyConfig {
+  onSuccess: (data: ResponseData) => void;
+}
+
+export const useGetPolicy = (config?: UseGetPolicyConfig) => {
+  const data = useQuery<ResponseData>({
+    queryKey: ["policy"],
+    queryFn: InvariantService.getPolicy,
+  });
+
+  const { isFetching, isSuccess, data: policy } = data;
+
+  React.useEffect(() => {
+    if (!isFetching && isSuccess && policy) {
+      config?.onSuccess(policy);
+    }
+  }, [isFetching, isSuccess, policy, config?.onSuccess]);
+
+  return data;
+};

frontend/src/hooks/query/use-get-risk-severity.ts

@@ -0,0 +1,26 @@
+import { useQuery } from "@tanstack/react-query";
+import React from "react";
+import InvariantService from "#/api/invariant-service";
+
+type ResponseData = number;
+
+interface UseGetRiskSeverityConfig {
+  onSuccess: (data: ResponseData) => void;
+}
+
+export const useGetRiskSeverity = (config?: UseGetRiskSeverityConfig) => {
+  const data = useQuery<ResponseData>({
+    queryKey: ["risk_severity"],
+    queryFn: InvariantService.getRiskSeverity,
+  });
+
+  const { isFetching, isSuccess, data: riskSeverity } = data;
+
+  React.useEffect(() => {
+    if (!isFetching && isSuccess && riskSeverity) {
+      config?.onSuccess(riskSeverity);
+    }
+  }, [isFetching, isSuccess, riskSeverity, config?.onSuccess]);
+
+  return data;
+};

frontend/src/hooks/query/use-get-traces.ts

@@ -0,0 +1,27 @@
+import { useQuery } from "@tanstack/react-query";
+import React from "react";
+import InvariantService from "#/api/invariant-service";
+
+type ResponseData = object;
+
+interface UseGetTracesConfig {
+  onSuccess: (data: ResponseData) => void;
+}
+
+export const useGetTraces = (config?: UseGetTracesConfig) => {
+  const data = useQuery({
+    queryKey: ["traces"],
+    queryFn: InvariantService.getTraces,
+    enabled: false,
+  });
+
+  const { isFetching, isSuccess, data: traces } = data;
+
+  React.useEffect(() => {
+    if (!isFetching && isSuccess && traces) {
+      config?.onSuccess(traces);
+    }
+  }, [isFetching, isSuccess, traces, config?.onSuccess]);
+
+  return data;
+};

frontend/src/hooks/query/use-is-authed.ts

@@ -15,5 +15,6 @@ export const useIsAuthed = () => {
     queryFn: () => OpenHands.authenticate(appMode!),
     enabled: !!appMode,
     staleTime: 1000 * 60 * 5, // 5 minutes
+    retry: false,
   });
 };