Change loki connection management

Author: ch1bo

ui/src/components/Sim/hooks/useLokiWebSocket.ts

@@ -10,6 +10,7 @@ import {
   ITransactionReceived,
 } from "@/components/Sim/types";
 import { useRef } from "react";
+import { EConnectionState } from "@/contexts/SimContext/types";
 
 // FIXME: latency in topology is wrong
 
@@ -387,11 +388,18 @@ function connectLokiWebSocket(lokiHost: string, dispatch: any): () => void {
     '{service="cardano-node"} |~ "BlockFetchServer|MsgBlock|CompletedBlockFetch|MsgLeiosBlock|MsgLeiosBlockTxs"';
   const wsUrl = `ws://${lokiHost}/loki/api/v1/tail?query=${encodeURIComponent(query)}&limit=5000`;
   console.log("Connecting to Loki:", wsUrl);
+  dispatch({
+    type: "SET_LOKI_CONNECTION_STATE",
+    payload: EConnectionState.Connecting,
+  });
 
   const ws = new WebSocket(wsUrl);
 
   ws.onopen = () => {
-    dispatch({ type: "SET_LOKI_CONNECTED", payload: true });
+    dispatch({
+      type: "SET_LOKI_CONNECTION_STATE",
+      payload: EConnectionState.Connected,
+    });
   };
 
   let count = 0;
@@ -438,25 +446,32 @@ function connectLokiWebSocket(lokiHost: string, dispatch: any): () => void {
 
   ws.onerror = (error) => {
     console.error("WebSocket error:", error);
-    dispatch({ type: "SET_LOKI_CONNECTED", payload: false });
+    dispatch({
+      type: "SET_LOKI_CONNECTION_STATE",
+      payload: EConnectionState.NotConnected,
+    });
   };
 
   ws.onclose = () => {
-    dispatch({ type: "SET_LOKI_CONNECTED", payload: false });
+    dispatch({
+      type: "SET_LOKI_CONNECTION_STATE",
+      payload: EConnectionState.NotConnected,
+    });
   };
 
   return () => ws.close();
 }
 
 export const useLokiWebSocket = () => {
   const {
-    state: { lokiHost, lokiConnected },
+    state: { lokiHost, lokiConnectionState },
     dispatch,
   } = useSimContext();
+
   const cleanupRef = useRef<(() => void) | null>(null);
 
   const connect = () => {
-    if (!lokiHost || lokiConnected) return;
+    if (!lokiHost || lokiConnectionState === EConnectionState.Connected) return;
 
     dispatch({ type: "RESET_TIMELINE" });
 
@@ -466,7 +481,10 @@ export const useLokiWebSocket = () => {
   const disconnect = () => {
     cleanupRef.current?.();
     cleanupRef.current = null;
-    dispatch({ type: "SET_LOKI_CONNECTED", payload: false });
+    dispatch({
+      type: "SET_LOKI_CONNECTION_STATE",
+      payload: EConnectionState.NotConnected,
+    });
   };
 
   return { connect, disconnect };

ui/src/components/Sim/modules/Scenario.tsx

@@ -2,19 +2,27 @@ import {
   useSimContext,
   defaultAggregatedData,
 } from "@/contexts/SimContext/context";
-import { IScenario } from "@/contexts/SimContext/types";
+import { EConnectionState, IScenario } from "@/contexts/SimContext/types";
 import { ChangeEvent, FC, useCallback, useEffect, useState } from "react";
 import { useStreamMessagesHandler } from "../hooks/useStreamMessagesHandler";
 import { useLokiWebSocket } from "../hooks/useLokiWebSocket";
 import { Button } from "@/components/Button";
 
 export const Scenario: FC = () => {
   const {
-    state: { allScenarios, activeScenario, events, lokiHost, lokiConnected, autoStart },
+    state: {
+      allScenarios,
+      activeScenario,
+      events,
+      lokiHost,
+      lokiConnectionState,
+      autoStart,
+    },
     dispatch,
   } = useSimContext();
   const { startStream, streaming, stopStream } = useStreamMessagesHandler();
-  const { connect: connectLoki, disconnect: disconnectLoki, connecting: lokiConnecting } = useLokiWebSocket();
+  const { connect: connectLoki, disconnect: disconnectLoki } =
+    useLokiWebSocket();
   const [includeTransactions, setIncludeTransactions] = useState(true);
 
   const isLokiMode = !!lokiHost;
@@ -29,7 +37,9 @@ export const Scenario: FC = () => {
           scenario.topology,
           window.location.toString(),
         ).toString(),
-        trace: scenario.trace ? new URL(scenario.trace, window.location.toString()).toString() : undefined,
+        trace: scenario.trace
+          ? new URL(scenario.trace, window.location.toString()).toString()
+          : undefined,
       }));
       dispatch({ type: "SET_SCENARIOS", payload: scenarios });
 
@@ -40,7 +50,11 @@ export const Scenario: FC = () => {
         const scenarioIndex = parseInt(scenarioParam, 10);
         if (scenarioIndex >= 0 && scenarioIndex < scenarios.length) {
           const targetScenario = scenarios[scenarioIndex];
-          dispatch({ type: "SET_SCENARIO", payload: targetScenario.name, autoStart: true });
+          dispatch({
+            type: "SET_SCENARIO",
+            payload: targetScenario.name,
+            autoStart: true,
+          });
         }
       }
     })();
@@ -55,9 +69,20 @@ export const Scenario: FC = () => {
         startStream(true); // Include transactions by default for auto-load
       }
       // Reset autoStart flag after triggering
-      dispatch({ type: "SET_SCENARIO", payload: activeScenario, autoStart: false });
+      dispatch({
+        type: "SET_SCENARIO",
+        payload: activeScenario,
+        autoStart: false,
+      });
     }
-  }, [autoStart, activeScenario, isLokiMode, connectLoki, startStream, dispatch]);
+  }, [
+    autoStart,
+    activeScenario,
+    isLokiMode,
+    connectLoki,
+    startStream,
+    dispatch,
+  ]);
 
   const chooseScenario = useCallback(
     (event: ChangeEvent<HTMLSelectElement>) => {
@@ -89,8 +114,9 @@ export const Scenario: FC = () => {
     }
   }, [isLokiMode, connectLoki, startStream, includeTransactions]);
 
-  const isLoaded = events.length > 0 || streaming || lokiConnected;
-  const isConnecting = streaming || lokiConnecting;
+  const isConnected = lokiConnectionState == EConnectionState.Connected;
+  const isConnecting = lokiConnectionState == EConnectionState.Connecting;
+  const isLoaded = events.length > 0 || streaming || isConnected;
 
   return (
     <div className="flex items-center justify-start gap-4 border-2 rounded-md p-4 border-gray-200 bg-white/80 backdrop-blur-xs">
@@ -130,10 +156,17 @@ export const Scenario: FC = () => {
           onClick={handleStartStream}
           disabled={isConnecting || isLoaded}
         >
-          {isLokiMode 
-            ? (lokiConnecting ? "Connecting..." : lokiConnected ? "Connected" : "Connect") 
-            : (streaming ? "Loading..." : isLoaded ? "Loaded" : "Load Scenario")
-          }
+          {isLokiMode
+            ? isConnecting
+              ? "Connecting..."
+              : isConnected
+                ? "Connected"
+                : "Connect"
+            : streaming
+              ? "Loading..."
+              : isLoaded
+                ? "Loaded"
+                : "Load"}
         </Button>
         {isLoaded && (
           <Button

ui/src/contexts/SimContext/context.ts

@@ -4,6 +4,7 @@ import {
   ISimContext,
   ISimContextState,
   ISimulationAggregatedDataState,
+  EConnectionState,
 } from "./types";
 
 export const defaultAggregatedData: ISimulationAggregatedDataState = {
@@ -34,7 +35,7 @@ export const defaultState: ISimContextState = {
   aggregatedData: defaultAggregatedData,
   tracePath: "",
   lokiHost: undefined,
-  lokiConnected: false,
+  lokiConnectionState: EConnectionState.NotConnected,
   topography: { links: new Map(), nodes: new Map() },
   topologyPath: "",
   topologyLoaded: false,

ui/src/contexts/SimContext/reducer.ts

@@ -1,5 +1,9 @@
 import { defaultAggregatedData } from "./context";
-import { ISimContextState, TSimContextActions } from "./types";
+import {
+  ISimContextState,
+  TSimContextActions,
+  EConnectionState,
+} from "./types";
 import {
   computeAggregatedDataAtTime,
   clearLatencyCache,
@@ -37,7 +41,7 @@ export const reducer = (
         autoStart: action.autoStart || false,
         tracePath: scenario.trace || "",
         lokiHost: scenario.loki,
-        lokiConnected: false,
+        lokiConnectionState: EConnectionState.NotConnected,
         topologyPath: scenario.topology,
         topologyLoaded:
           state.topologyLoaded && scenario.topology === state.topologyPath,
@@ -199,10 +203,10 @@ export const reducer = (
         speedMultiplier: 1,
       };
 
-    case "SET_LOKI_CONNECTED":
+    case "SET_LOKI_CONNECTION_STATE":
       return {
         ...state,
-        lokiConnected: action.payload,
+        lokiConnectionState: action.payload,
       };
 
     default:

ui/src/contexts/SimContext/types.ts

@@ -63,6 +63,12 @@ export interface IGraphContextState {
   currentNode?: string;
 }
 
+export enum EConnectionState {
+  NotConnected = "NotConnected",
+  Connecting = "Connecting",
+  Connected = "Connected",
+}
+
 export interface IScenario {
   name: string;
   topology: string;
@@ -79,7 +85,7 @@ export interface ISimContextState {
   aggregatedData: ISimulationAggregatedDataState;
   tracePath: string;
   lokiHost?: string;
-  lokiConnected: boolean;
+  lokiConnectionState: EConnectionState;
   topography: ITransformedNodeMap;
   topologyPath: string;
   topologyLoaded: boolean;
@@ -114,7 +120,7 @@ export type TSimContextActions =
   | { type: "SET_TIMELINE_PLAYING"; payload: boolean }
   | { type: "SET_TIMELINE_SPEED"; payload: number }
   | { type: "RESET_TIMELINE" }
-  | { type: "SET_LOKI_CONNECTED"; payload: boolean };
+  | { type: "SET_LOKI_CONNECTION_STATE"; payload: EConnectionState };
 
 export interface ISimContext {
   state: ISimContextState;