From d2cb543c76fdb2fc96ff269af913c2257800ffbc Mon Sep 17 00:00:00 2001 From: mbecker20 Date: Sat, 15 Feb 2025 13:58:13 -0800 Subject: [PATCH] use komodo_client.subscribe_to_update_websocket, and click indicator to reconnect --- frontend/src/lib/hooks.ts | 19 ++-- frontend/src/lib/socket.tsx | 205 +++++++++++++++--------------------- frontend/src/main.tsx | 2 +- 3 files changed, 98 insertions(+), 128 deletions(-) diff --git a/frontend/src/lib/hooks.ts b/frontend/src/lib/hooks.ts index a193cf4a2..83ebb5797 100644 --- a/frontend/src/lib/hooks.ts +++ b/frontend/src/lib/hooks.ts @@ -1,5 +1,5 @@ import { AUTH_TOKEN_STORAGE_KEY, KOMODO_BASE_URL } from "@main"; -import { KomodoClient as Client, Types } from "komodo_client"; +import { KomodoClient, Types } from "komodo_client"; import { AuthResponses, ExecuteResponses, @@ -27,19 +27,20 @@ import { RESOURCE_TARGETS } from "./utils"; const token = () => ({ jwt: localStorage.getItem(AUTH_TOKEN_STORAGE_KEY) ?? "", }); -const client = () => Client(KOMODO_BASE_URL, { type: "jwt", params: token() }); +export const komodo_client = () => + KomodoClient(KOMODO_BASE_URL, { type: "jwt", params: token() }); export const useLoginOptions = () => useQuery({ queryKey: ["GetLoginOptions"], - queryFn: () => client().auth("GetLoginOptions", {}), + queryFn: () => komodo_client().auth("GetLoginOptions", {}), }); export const useUser = () => { const userInvalidate = useUserInvalidate(); const query = useQuery({ queryKey: ["GetUser"], - queryFn: () => client().auth("GetUser", {}), + queryFn: () => komodo_client().auth("GetUser", {}), refetchInterval: 30_000, }); useEffect(() => { @@ -77,7 +78,7 @@ export const useRead = < ) => useQuery({ queryKey: [type, params], - queryFn: () => client().read(type, params), + queryFn: () => komodo_client().read(type, params), ...config, }); @@ -106,7 +107,7 @@ export const useManageUser = < const { toast } = useToast(); return useMutation({ mutationKey: [type], - mutationFn: (params: P) => client().user(type, params), + mutationFn: (params: P) => komodo_client().user(type, params), onError: (e: { result: { error?: string } }, v, c) => { console.log("Auth error:", e); const msg = e.result.error ?? "Unknown error. See console."; @@ -140,7 +141,7 @@ export const useWrite = < const { toast } = useToast(); return useMutation({ mutationKey: [type], - mutationFn: (params: P) => client().write(type, params), + mutationFn: (params: P) => komodo_client().write(type, params), onError: (e: { result: { error?: string } }, v, c) => { console.log("Write error:", e); const msg = e.result.error ?? "Unknown error. See console."; @@ -174,7 +175,7 @@ export const useExecute = < const { toast } = useToast(); return useMutation({ mutationKey: [type], - mutationFn: (params: P) => client().execute(type, params), + mutationFn: (params: P) => komodo_client().execute(type, params), onError: (e: { result: { error?: string } }, v, c) => { console.log("Execute error:", e); const msg = e.result.error ?? "Unknown error. See console."; @@ -208,7 +209,7 @@ export const useAuth = < const { toast } = useToast(); return useMutation({ mutationKey: [type], - mutationFn: (params: P) => client().auth(type, params), + mutationFn: (params: P) => komodo_client().auth(type, params), onError: (e: { result: { error?: string } }, v, c) => { console.log("Auth error:", e); const msg = e.result.error ?? "Unknown error. See console."; diff --git a/frontend/src/lib/socket.tsx b/frontend/src/lib/socket.tsx index f80b0fa79..710d9d97c 100644 --- a/frontend/src/lib/socket.tsx +++ b/frontend/src/lib/socket.tsx @@ -1,22 +1,29 @@ -import { useInvalidate, useUser } from "@lib/hooks"; -import { Types } from "komodo_client"; +import { komodo_client, useInvalidate, useUser } from "@lib/hooks"; +import { CancelToken, Types } from "komodo_client"; import { Button } from "@ui/button"; import { toast } from "@ui/use-toast"; import { atom, useAtom } from "jotai"; -import { Circle } from "lucide-react"; +import { Circle, Loader2 } from "lucide-react"; import { ReactNode, useCallback, useEffect, useState } from "react"; import { cn } from "@lib/utils"; -import { AUTH_TOKEN_STORAGE_KEY } from "@main"; import { ResourceComponents } from "@components/resources"; import { UsableResource } from "@types"; import { ResourceName } from "@components/resources/common"; -const rws_atom = atom(null); -const useWebsocket = () => useAtom(rws_atom); - const ws_connected = atom(false); export const useWebsocketConnected = () => useAtom(ws_connected); +const ws_cancel = atom(new CancelToken()); +const useWebsocketCancel = () => useAtom(ws_cancel)[0]; + +const useWebsocketReconnect = () => { + const [cancel, set] = useAtom(ws_cancel); + return () => { + cancel.cancel(); + set(new CancelToken()); + }; +}; + const onMessageHandlers: { [key: string]: (update: Types.UpdateListItem) => void; } = {}; @@ -34,13 +41,81 @@ export const useWebsocketMessages = ( }, []); }; -const on_message = ( - { data }: MessageEvent, +let count = 0; + +export const WebsocketProvider = ({ children }: { children: ReactNode }) => { + const user = useUser().data; + const invalidate = useInvalidate(); + const cancel = useWebsocketCancel(); + const [connected, setConnected] = useWebsocketConnected(); + + const on_update_fn = useCallback( + (update: Types.UpdateListItem) => on_update(update, invalidate), + [invalidate] + ); + + useEffect(() => { + if (user && !connected) { + count = count + 1; + const _count = count; + komodo_client().subscribe_to_update_websocket({ + on_login: () => { + setConnected(true); + console.info(_count + " | Logged into Update websocket"); + }, + on_update: on_update_fn, + on_close: () => { + setConnected(false); + console.info(_count + " | Update websocket connection closed"); + }, + cancel, + }); + } + }, [user, cancel, connected]); + + return <>{children}; +}; + +export const WsStatusIndicator = () => { + const [refreshing, setRefreshing] = useState(false); + const [connected] = useWebsocketConnected(); + const reconnect = useWebsocketReconnect(); + const onclick = () => { + setRefreshing(true); + setTimeout(() => setRefreshing(false), 500); + reconnect(); + toast({ + title: connected + ? "Triggered websocket reconnect" + : "Triggered websocket connect", + }); + }; + + return ( + + ); +}; + +const on_update = ( + update: Types.UpdateListItem, invalidate: ReturnType ) => { - if (data == "LOGGED_IN") return console.info("logged in to ws"); - const update = JSON.parse(data) as Types.UpdateListItem; - const Components = ResourceComponents[update.target.type as UsableResource]; const title = Components ? (
@@ -245,109 +320,3 @@ const on_message = ( // Run any attached handlers Object.values(onMessageHandlers).forEach((handler) => handler(update)); }; - -export const WebsocketProvider = ({ - url, - children, -}: { - url: string; - children: ReactNode; -}) => { - const user = useUser().data; - const invalidate = useInvalidate(); - const [ws, set] = useWebsocket(); - const [connected, setConnected] = useWebsocketConnected(); - // don't care about value, just use to make sure value changes - // to trigger connection useEffect - const [reconnect, setReconnect] = useState(false); - - const on_message_fn = useCallback( - (e: MessageEvent) => on_message(e, invalidate), - [invalidate] - ); - - // Connection useEffect - useEffect(() => { - if (user && !connected) { - const ws = make_websocket({ - url, - on_open: () => setConnected(true), - on_message: on_message_fn, - on_close: () => { - setConnected(false); - }, - }); - set(ws); - } - }, [set, url, user, connected, reconnect]); - - useEffect(() => { - // poll for CLOSED state. - // trigger reconnect after stale page - const interval = setInterval(() => { - if (ws?.readyState === WebSocket.CLOSED) { - setConnected(false); - // toggle to make sure connection useEffect runs. - // which could happen if connected is stuck in false state, - // so setConnected(false) doesn't trigger reconnect - setReconnect(!reconnect); - } - }, 3_000); - - return () => clearInterval(interval); - }, []); - - return <>{children}; -}; - -export const WsStatusIndicator = () => { - const [connected] = useWebsocketConnected(); - const onclick = () => - toast({ - title: connected ? "Websocket connected" : "Websocket disconnected", - }); - - return ( - - ); -}; - -const make_websocket = ({ - url, - on_open, - on_message, - on_close, -}: { - url: string; - on_open: () => void; - on_message: (e: MessageEvent) => void; - on_close: () => void; -}) => { - const ws = new WebSocket(url); - - const _on_open = () => { - const jwt = localStorage.getItem(AUTH_TOKEN_STORAGE_KEY); - if (!ws || !jwt) return; - const msg: Types.WsLoginMessage = { type: "Jwt", params: { jwt } }; - if (jwt && ws) ws.send(JSON.stringify(msg)); - on_open(); - }; - - ws.addEventListener("open", _on_open); - ws.addEventListener("message", on_message); - ws.addEventListener("close", on_close); - - return ws; -}; diff --git a/frontend/src/main.tsx b/frontend/src/main.tsx index f9ca3b84b..e091adbd5 100644 --- a/frontend/src/main.tsx +++ b/frontend/src/main.tsx @@ -33,7 +33,7 @@ init_monaco().then(() => ReactDOM.createRoot(document.getElementById("root")!).render( // - +