refactor: Session log modal transition handling

react/src/components/ComputeSessionList.tsx

@@ -1,63 +1,9 @@
-import BAIModal from './BAIModal';
-import ContainerLogModalWithLazyQueryLoader from './ComputeSessionNodeItems/ContainerLogModalWithLazyQueryLoader';
-import SessionDetailDrawer from './SessionDetailDrawer';
-import { Skeleton } from 'antd';
-import { Suspense, useEffect, useState } from 'react';
-import { StringParam, useQueryParam } from 'use-query-params';
+import SessionDetailAndContainerLogOpenerLegacy from './SessionDetailAndContainerLogOpenerLegacy';
 
 const ComputeSessionList = () => {
-  const [sessionId, setSessionId] = useQueryParam('sessionDetail', StringParam);
-
-  const [containerLogModalSessionId, setContainerLogModalSessionId] =
-    useState<string>();
-  useEffect(() => {
-    const handler = (e: any) => {
-      setContainerLogModalSessionId(e.detail);
-    };
-    document.addEventListener('bai-open-session-log', handler);
-    return () => {
-      document.removeEventListener('bai-open-session-log', handler);
-    };
-  }, []);
   return (
     <>
-      <SessionDetailDrawer
-        open={!sessionId}
-        sessionId={sessionId || undefined}
-        onClose={() => {
-          setSessionId(null, 'replaceIn');
-        }}
-      />
-      {containerLogModalSessionId && (
-        <Suspense
-          fallback={
-            <BAIModal
-              open
-              // loading
-              width={'100%'}
-              styles={{
-                header: {
-                  width: '100%',
-                },
-                body: {
-                  height: 'calc(100vh - 100px)',
-                  maxHeight: 'calc(100vh - 100px)',
-                },
-              }}
-              footer={null}
-            >
-              <Skeleton active />
-            </BAIModal>
-          }
-        >
-          <ContainerLogModalWithLazyQueryLoader
-            sessionId={containerLogModalSessionId}
-            afterClose={() => {
-              setContainerLogModalSessionId(undefined);
-            }}
-          />
-        </Suspense>
-      )}
+      <SessionDetailAndContainerLogOpenerLegacy />
     </>
   );
 };

react/src/components/ComputeSessionNodeItems/ContainerLogModal.tsx

@@ -4,28 +4,21 @@ import { useCurrentUserRole } from '../../hooks/backendai';
 import { useTanQuery } from '../../hooks/reactQueryAlias';
 import { useMemoWithPrevious } from '../../hooks/useMemoWithPrevious';
 import BAIModal, { BAIModalProps } from '../BAIModal';
+import BAISelect from '../BAISelect';
 import Flex from '../Flex';
 import { ContainerLogModalFragment$key } from './__generated__/ContainerLogModalFragment.graphql';
 import { ReloadOutlined } from '@ant-design/icons';
 import { LazyLog, ScrollFollow } from '@melloware/react-logviewer';
-import {
-  Button,
-  Divider,
-  Grid,
-  Select,
-  theme,
-  Tooltip,
-  Typography,
-} from 'antd';
+import { Button, Divider, Grid, theme, Tooltip, Typography } from 'antd';
 import graphql from 'babel-plugin-relay/macro';
 import _ from 'lodash';
 import { DownloadIcon } from 'lucide-react';
-import React, { useState } from 'react';
+import React, { useEffect, useState } from 'react';
 import { useTranslation } from 'react-i18next';
 import { useFragment } from 'react-relay';
 
 interface ContainerLogModalProps extends BAIModalProps {
-  sessionFrgmt: ContainerLogModalFragment$key;
+  sessionFrgmt: ContainerLogModalFragment$key | null;
 }
 
 const ContainerLogModal: React.FC<ContainerLogModalProps> = ({
@@ -66,9 +59,14 @@ const ContainerLogModal: React.FC<ContainerLogModalProps> = ({
       kernelNodes[0]?.row_id,
   );
 
+  useEffect(() => {
+    if (modalProps.open === false) {
+      setSelectedKernelId(undefined);
+    }
+  }, [modalProps.open]);
+
   // Currently we can only fetch full logs
   // const [logSize, setLogSize] = useState<100|'full'>('full');
-
   const {
     data: logs,
     refetch,
@@ -129,20 +127,24 @@ const ContainerLogModal: React.FC<ContainerLogModalProps> = ({
           <Typography.Title level={4} style={{ margin: 0, flexShrink: 0 }}>
             Logs
           </Typography.Title>
-          <Typography.Text style={{ fontWeight: 'normal' }} ellipsis>
-            {session?.name}
-          </Typography.Text>
-          <Typography.Text
-            style={{ fontWeight: 'normal', fontFamily: 'monospace' }}
-            copyable={{
-              text: session?.row_id,
-              tooltips: t('button.CopySomething', {
-                name: t('session.SessionId'),
-              }),
-            }}
-          >
-            ({md ? session?.row_id : session?.row_id.split('-')?.[0]})
-          </Typography.Text>
+          {session ? (
+            <>
+              <Typography.Text style={{ fontWeight: 'normal' }} ellipsis>
+                {session?.name}
+              </Typography.Text>
+              <Typography.Text
+                style={{ fontWeight: 'normal', fontFamily: 'monospace' }}
+                copyable={{
+                  text: session?.row_id,
+                  tooltips: t('button.CopySomething', {
+                    name: t('session.SessionId'),
+                  }),
+                }}
+              >
+                ({md ? session?.row_id : session?.row_id.split('-')?.[0]})
+              </Typography.Text>
+            </>
+          ) : null}
         </Flex>
       }
       width={'100%'}
@@ -157,6 +159,7 @@ const ContainerLogModal: React.FC<ContainerLogModalProps> = ({
       }}
       {...modalProps}
       footer={null}
+      destroyOnClose
     >
       <Flex
         direction="column"
@@ -166,12 +169,13 @@ const ContainerLogModal: React.FC<ContainerLogModalProps> = ({
       >
         <Flex gap="sm" wrap="wrap">
           Kernel Role
-          <Select
+          <BAISelect
             value={selectedKernelId}
             onChange={(value) => {
               setSelectedKernelId(value);
               resetPreviousLineNumber();
             }}
+            autoSelectOption
             options={_.chain(session?.kernel_nodes?.edges)
               .sortBy((e) => `${e?.node?.cluster_role} ${e?.node?.cluster_idx}`)
               .map((e) => {

react/src/components/ComputeSessionNodeItems/ContainerLogModalWithLazyQueryLoader.tsx

@@ -2,15 +2,15 @@ import { useCurrentProjectValue } from '../../hooks/useCurrentProject';
 import ContainerLogModal from './ContainerLogModal';
 import { ContainerLogModalWithLazyQueryLoaderQuery } from './__generated__/ContainerLogModalWithLazyQueryLoaderQuery.graphql';
 import graphql from 'babel-plugin-relay/macro';
-import { useState } from 'react';
 import { useLazyLoadQuery } from 'react-relay';
 
 const ContainerLogModalWithLazyQueryLoader: React.FC<{
-  sessionId: string;
-  afterClose?: () => void;
-}> = ({ sessionId, afterClose }) => {
+  sessionId?: string;
+  open: boolean;
+  loading: boolean;
+  onRequestClose?: () => void;
+}> = ({ sessionId, open, loading, onRequestClose }) => {
   const currentProject = useCurrentProjectValue();
-  const [open, setOpen] = useState(true);
   const { compute_session_node } =
     useLazyLoadQuery<ContainerLogModalWithLazyQueryLoaderQuery>(
       graphql`
@@ -27,23 +27,20 @@ const ContainerLogModalWithLazyQueryLoader: React.FC<{
         sessionId,
         project_id: currentProject.id,
       },
+      {
+        fetchPolicy: sessionId ? 'network-only' : 'store-only',
+      },
     );
 
   return (
-    compute_session_node && (
-      <ContainerLogModal
-        maskTransitionName={open ? '' : undefined}
-        transitionName={open ? '' : undefined}
-        sessionFrgmt={compute_session_node}
-        open={open}
-        onCancel={() => {
-          setOpen(false);
-        }}
-        afterClose={() => {
-          afterClose?.();
-        }}
-      />
-    )
+    <ContainerLogModal
+      sessionFrgmt={compute_session_node || null}
+      open={open}
+      loading={loading}
+      onCancel={() => {
+        onRequestClose && onRequestClose();
+      }}
+    />
   );
 };
 

react/src/components/SessionDetailAndContainerLogOpenerLegacy.tsx

@@ -0,0 +1,44 @@
+import ContainerLogModalWithLazyQueryLoader from './ComputeSessionNodeItems/ContainerLogModalWithLazyQueryLoader';
+import SessionDetailDrawer from './SessionDetailDrawer';
+import { useState, useEffect, useTransition } from 'react';
+import { useQueryParam, StringParam } from 'use-query-params';
+
+const SessionDetailAndContainerLogOpenerLegacy = () => {
+  const [sessionId, setSessionId] = useQueryParam('sessionDetail', StringParam);
+  const [containerLogModalSessionId, setContainerLogModalSessionId] =
+    useState<string>();
+  const [isPendingLogModalOpen, startLogModalOpenTransition] = useTransition();
+
+  useEffect(() => {
+    const handler = (e: any) => {
+      startLogModalOpenTransition(() => {
+        setContainerLogModalSessionId(e.detail);
+      });
+    };
+    document.addEventListener('bai-open-session-log', handler);
+    return () => {
+      document.removeEventListener('bai-open-session-log', handler);
+    };
+  }, [startLogModalOpenTransition, setContainerLogModalSessionId]);
+
+  return (
+    <>
+      <SessionDetailDrawer
+        open={!sessionId}
+        sessionId={sessionId || undefined}
+        onClose={() => {
+          setSessionId(null, 'replaceIn');
+        }}
+      />
+      <ContainerLogModalWithLazyQueryLoader
+        open={!!containerLogModalSessionId || isPendingLogModalOpen}
+        loading={isPendingLogModalOpen}
+        sessionId={containerLogModalSessionId}
+        onRequestClose={() => {
+          setContainerLogModalSessionId(undefined);
+        }}
+      />
+    </>
+  );
+};
+export default SessionDetailAndContainerLogOpenerLegacy;