diff --git a/frontend/src/components/modals/security/invariant/invariant.tsx b/frontend/src/components/modals/security/invariant/invariant.tsx index 5813a7e3ff9d..dd53e61311d2 100644 --- a/frontend/src/components/modals/security/invariant/invariant.tsx +++ b/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(null); - const { data: fetchedPolicy } = useQuery({ - queryKey: ["policy"], - queryFn: InvariantService.getPolicy, - }); - - const { data: riskSeverity } = useQuery({ - 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 }) => diff --git a/frontend/src/hooks/query/use-get-policy.ts b/frontend/src/hooks/query/use-get-policy.ts new file mode 100644 index 000000000000..b8c450c2a303 --- /dev/null +++ b/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({ + 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; +}; diff --git a/frontend/src/hooks/query/use-get-risk-severity.ts b/frontend/src/hooks/query/use-get-risk-severity.ts new file mode 100644 index 000000000000..a72e4b30d500 --- /dev/null +++ b/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({ + 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; +}; diff --git a/frontend/src/hooks/query/use-get-traces.ts b/frontend/src/hooks/query/use-get-traces.ts new file mode 100644 index 000000000000..7eaa1fed7917 --- /dev/null +++ b/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; +}; diff --git a/frontend/src/hooks/query/use-is-authed.ts b/frontend/src/hooks/query/use-is-authed.ts index 66a2c9a657b4..45c51b70bf15 100644 --- a/frontend/src/hooks/query/use-is-authed.ts +++ b/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, }); };