refactor: move apolloClient out of Zustand store (#507)

Author: lihebi

apps/ui/src/components/Canvas.tsx

@@ -48,6 +48,7 @@ import FloatingEdge from "./nodes/FloatingEdge";
 import CustomConnectionLine from "./nodes/CustomConnectionLine";
 import HelperLines from "./HelperLines";
 import { getAbsPos, newNodeShapeConfig } from "../lib/store/canvasSlice";
+import { useApolloClient } from "@apollo/client";
 
 const nodeTypes = { SCOPE: ScopeNode, CODE: CodeNode, RICH: RichNode };
 const edgeTypes = {
@@ -149,6 +150,7 @@ function useJump() {
   const selectPod = useStore(store, (state) => state.selectPod);
 
   const yjsRun = useStore(store, (state) => state.yjsRun);
+  const apolloClient = useApolloClient();
 
   const setCenterSelection = useStore(
     store,
@@ -238,15 +240,15 @@ function useJump() {
         if (pod.type == "CODE") {
           if (event.shiftKey) {
             // Hitting "SHIFT"+"Enter" will run the code pod
-            yjsRun(id);
+            yjsRun(id, apolloClient);
           } else {
             // Hitting "Enter" on a Code pod will go to "Edit" mode.
             setFocusedEditor(id);
           }
         } else if (pod.type === "SCOPE") {
           if (event.shiftKey) {
             // Hitting "SHIFT"+"Enter" on a Scope will run the scope.
-            yjsRun(id);
+            yjsRun(id, apolloClient);
           }
         } else if (pod.type === "RICH") {
           // Hitting "Enter" on a Rich pod will go to "Edit" mode.

apps/ui/src/components/MyMonaco.tsx

@@ -10,6 +10,7 @@ import { RepoContext } from "../lib/store";
 import { MonacoBinding } from "y-monaco";
 import { useReactFlow } from "reactflow";
 import { Annotation } from "../lib/parser";
+import { useApolloClient } from "@apollo/client";
 
 const theme: monaco.editor.IStandaloneThemeData = {
   base: "vs",
@@ -395,6 +396,7 @@ export const MyMonaco = memo<MyMonacoProps>(function MyMonaco({
   const readOnly = useStore(store, (state) => state.role === "GUEST");
   const showLineNumbers = useStore(store, (state) => state.showLineNumbers);
   const yjsRun = useStore(store, (state) => state.yjsRun);
+  const apolloClient = useApolloClient();
   const focusedEditor = useStore(store, (state) => state.focusedEditor);
   const setFocusedEditor = useStore(store, (state) => state.setFocusedEditor);
   const annotations = useStore(
@@ -483,7 +485,7 @@ export const MyMonaco = memo<MyMonacoProps>(function MyMonaco({
       label: "Run",
       keybindings: [monaco.KeyMod.Shift | monaco.KeyCode.Enter],
       run: () => {
-        yjsRun(id);
+        yjsRun(id, apolloClient);
       },
     });
     editor.addAction({

apps/ui/src/components/nodes/Code.tsx

@@ -58,6 +58,7 @@ import { timeDifference } from "../../lib/utils/utils";
 import { ButtonGroup } from "@mui/material";
 
 import { ConfirmDeleteButton } from "./utils";
+import { useApolloClient } from "@apollo/client";
 
 function Timer({ lastExecutedAt }) {
   const [counter, setCounter] = useState(0);
@@ -388,6 +389,7 @@ function MyFloatingToolbar({ id, layout, setLayout }) {
   // const pod = useStore(store, (state) => state.pods[id]);
   const yjsRun = useStore(store, (state) => state.yjsRun);
   const yjsRunChain = useStore(store, (state) => state.yjsRunChain);
+  const apolloClient = useApolloClient();
   // right, bottom
   const isGuest = useStore(store, (state) => state.role === "GUEST");
 
@@ -413,7 +415,7 @@ function MyFloatingToolbar({ id, layout, setLayout }) {
         <Tooltip title="Run (shift-enter)">
           <IconButton
             onClick={() => {
-              yjsRun(id);
+              yjsRun(id, apolloClient);
             }}
           >
             <PlayCircleOutlineIcon style={{ fontSize: iconFontSize }} />
@@ -424,7 +426,7 @@ function MyFloatingToolbar({ id, layout, setLayout }) {
         <Tooltip title="Run chain">
           <IconButton
             onClick={() => {
-              yjsRunChain(id);
+              yjsRunChain(id, apolloClient);
             }}
           >
             <KeyboardDoubleArrowRightIcon style={{ fontSize: iconFontSize }} />

apps/ui/src/components/nodes/Scope.tsx

@@ -50,13 +50,15 @@ import {
   level2fontsize,
 } from "./utils";
 import { CopyToClipboard } from "react-copy-to-clipboard";
+import { useApolloClient } from "@apollo/client";
 
 function MyFloatingToolbar({ id }: { id: string }) {
   const store = useContext(RepoContext);
   if (!store) throw new Error("Missing BearContext.Provider in the tree");
   const reactFlowInstance = useReactFlow();
   const isGuest = useStore(store, (state) => state.role === "GUEST");
   const yjsRun = useStore(store, (state) => state.yjsRun);
+  const apolloClient = useApolloClient();
 
   const autoLayout = useStore(store, (state) => state.autoLayout);
 
@@ -85,7 +87,7 @@ function MyFloatingToolbar({ id }: { id: string }) {
         <Tooltip title="Run (shift-enter)">
           <IconButton
             onClick={() => {
-              yjsRun(id);
+              yjsRun(id, apolloClient);
             }}
           >
             <PlayCircleOutlineIcon style={{ fontSize: iconFontSize }} />

apps/ui/src/lib/store/runtimeSlice.ts

@@ -172,15 +172,12 @@ type PodResult = {
 };
 
 export interface RuntimeSlice {
-  apolloClient?: ApolloClient<any>;
-  setApolloClient: (client: ApolloClient<any>) => void;
-
   parsePod: (id: string) => void;
   parseAllPods: () => void;
   resolvePod: (id) => void;
   resolveAllPods: () => void;
-  yjsRun: (id: string) => void;
-  yjsRunChain: (id: string) => void;
+  yjsRun: (id: string, apolloClient: ApolloClient<any>) => void;
+  yjsRunChain: (id: string, apolloClient: ApolloClient<any>) => void;
   clearResults: (id) => void;
   setRunning: (id) => void;
   parseResult: Record<
@@ -196,16 +193,14 @@ export interface RuntimeSlice {
   getResultMap(): Y.Map<PodResult>;
   activeRuntime?: string;
   setActiveRuntime(id?: string): void;
-  yjsSendRun(ids: string[]): void;
+  yjsSendRun(ids: string[], apolloClient: ApolloClient<any>): void;
   isRuntimeReady(): boolean;
 }
 
 export const createRuntimeSlice: StateCreator<MyState, [], [], RuntimeSlice> = (
   set,
   get
 ) => ({
-  apolloClient: undefined,
-  setApolloClient: (client) => set({ apolloClient: client }),
   parseResult: {},
 
   // new yjs-based runtime
@@ -306,7 +301,7 @@ export const createRuntimeSlice: StateCreator<MyState, [], [], RuntimeSlice> = (
     }
     return true;
   },
-  yjsSendRun(ids) {
+  yjsSendRun(ids, apolloClient) {
     const activeRuntime = get().activeRuntime!;
     let specs = ids.map((id) => {
       // Actually send the run request.
@@ -337,7 +332,7 @@ export const createRuntimeSlice: StateCreator<MyState, [], [], RuntimeSlice> = (
       return false;
     });
     if (specs) {
-      get().apolloClient?.mutate({
+      apolloClient.mutate({
         mutation: gql`
           mutation RunChain($specs: [RunSpecInput], $runtimeId: String) {
             runChain(specs: $specs, runtimeId: $runtimeId)
@@ -353,21 +348,21 @@ export const createRuntimeSlice: StateCreator<MyState, [], [], RuntimeSlice> = (
       });
     }
   },
-  yjsRun: (id) => {
+  yjsRun: (id, apolloClient) => {
     if (!get().isRuntimeReady()) return;
     const nodesMap = get().getNodesMap();
     const nodes = Array.from<Node>(nodesMap.values());
     const node = nodesMap.get(id);
     if (!node) return;
     const chain = getDescendants(node, nodes);
-    get().yjsSendRun(chain);
+    get().yjsSendRun(chain, apolloClient);
   },
   /**
    * Add the pod and all its downstream pods (defined by edges) to the chain and run the chain.
    * @param id the id of the pod to start the chain
    * @returns
    */
-  yjsRunChain: async (id) => {
+  yjsRunChain: async (id, apolloClient) => {
     if (!get().isRuntimeReady()) return;
     // Get the chain: get the edges, and then get the pods
     const edgesMap = get().getEdgesMap();
@@ -387,7 +382,7 @@ export const createRuntimeSlice: StateCreator<MyState, [], [], RuntimeSlice> = (
       chain.push(nodeid);
       nodeid = node2target[nodeid];
     }
-    get().yjsSendRun(chain);
+    get().yjsSendRun(chain, apolloClient);
   },
   clearResults: (id) => {
     const resultMap = get().getResultMap();

apps/ui/src/pages/repo.tsx

@@ -392,14 +392,10 @@ function RepoImpl() {
     }
   }, [addClient, deleteClient, provider]);
 
-  const apolloClient = useApolloClient();
-  const setApolloClient = useStore(store, (state) => state.setApolloClient);
-
   useEffect(() => {
     if (hasToken()) {
       if (!loading && me) {
         setUser(me);
-        setApolloClient(apolloClient);
       }
     }
   }, [loading, me]);