refactor: lazy-load top-level pages, reduce initial page load (#1295)

* refactor: lazy-load top-level pages, reduce initial page load

* styles

frontend/src/components/editor/actions/useNotebookActions.tsx

@@ -22,13 +22,12 @@ import {
   Undo2Icon,
 } from "lucide-react";
 import { commandPaletteAtom } from "../controls/command-palette";
+import { useCellActions, useNotebook } from "@/core/cells/cells";
 import {
   canUndoDeletes,
   disabledCellIds,
   enabledCellIds,
-  useCellActions,
-  useNotebook,
-} from "@/core/cells/cells";
+} from "@/core/cells/utils";
 import { readCode, saveCellConfig } from "@/core/network/requests";
 import { Objects } from "@/utils/objects";
 import { ActionButton } from "./types";

frontend/src/components/editor/app-container.tsx

@@ -0,0 +1,37 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import { WebSocketState } from "@/core/websocket/types";
+import { cn } from "@/utils/cn";
+import React, { PropsWithChildren } from "react";
+import { StatusOverlay } from "./header/status";
+import { AppConfig } from "@/core/config/config-schema";
+
+interface Props {
+  connectionState: WebSocketState;
+  isRunning: boolean;
+  width: AppConfig["width"];
+}
+
+export const AppContainer: React.FC<PropsWithChildren<Props>> = ({
+  width,
+  connectionState,
+  isRunning,
+  children,
+}) => {
+  return (
+    <>
+      <StatusOverlay state={connectionState} isRunning={isRunning} />
+      <div
+        id="App"
+        className={cn(
+          connectionState === WebSocketState.CLOSED && "disconnected",
+          "bg-background w-full h-full text-textColor",
+          "flex flex-col overflow-y-auto overflow-x-hidden",
+          width === "full" && "config-width-full",
+        )}
+      >
+        {children}
+      </div>
+    </>
+  );
+};

frontend/src/components/editor/header/app-header.tsx

@@ -0,0 +1,24 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+import { ConnectionStatus, WebSocketState } from "@/core/websocket/types";
+import React, { PropsWithChildren } from "react";
+import { Disconnected } from "../Disconnected";
+
+interface Props {
+  className?: string;
+  connection: ConnectionStatus;
+}
+
+export const AppHeader: React.FC<PropsWithChildren<Props>> = ({
+  connection,
+  className,
+  children,
+}) => {
+  return (
+    <div className={className}>
+      {children}
+      {connection.state === WebSocketState.CLOSED && (
+        <Disconnected reason={connection.reason} />
+      )}
+    </div>
+  );
+};

frontend/src/components/editor/header/status.tsx

@@ -0,0 +1,48 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+import { Tooltip } from "@/components/ui/tooltip";
+import { notebookScrollToRunning } from "@/core/cells/actions";
+import { WebSocketState } from "@/core/websocket/types";
+import { UnlinkIcon, HourglassIcon } from "lucide-react";
+import React from "react";
+
+export const StatusOverlay: React.FC<{
+  state: WebSocketState;
+  isRunning: boolean;
+}> = ({ state, isRunning }) => {
+  return (
+    <>
+      {state === WebSocketState.OPEN && isRunning && <RunningIcon />}
+      {state === WebSocketState.CLOSED && <NoiseBackground />}
+      {state === WebSocketState.CLOSED && <DisconnectedIcon />}
+    </>
+  );
+};
+
+const topLeftStatus =
+  "absolute top-3 left-4 m-0 flex items-center space-x-3 min-h-[28px] no-print pointer-events-auto z-30";
+
+const DisconnectedIcon = () => (
+  <Tooltip content="App disconnected">
+    <div className={topLeftStatus}>
+      <UnlinkIcon className="closed-app-icon" />
+    </div>
+  </Tooltip>
+);
+
+const RunningIcon = () => (
+  <div
+    className={topLeftStatus}
+    data-testid="loading-indicator"
+    title={"Marimo is busy computing. Hang tight!"}
+    onClick={notebookScrollToRunning}
+  >
+    <HourglassIcon className="running-app-icon" size={30} strokeWidth={1} />
+  </div>
+);
+
+const NoiseBackground = () => (
+  <>
+    <div className="noise" />
+    <div className="disconnected-gradient" />
+  </>
+);

frontend/src/components/pages/edit-page.tsx

@@ -0,0 +1,21 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+import { EditApp } from "@/core/edit-app";
+import { AppChrome } from "../editor/chrome/wrapper/app-chrome";
+import { CommandPalette } from "../editor/controls/command-palette";
+import { AppConfig, UserConfig } from "@/core/config/config-schema";
+
+interface Props {
+  userConfig: UserConfig;
+  appConfig: AppConfig;
+}
+
+const EditPage = (props: Props) => {
+  return (
+    <AppChrome>
+      <EditApp {...props} />
+      <CommandPalette />
+    </AppChrome>
+  );
+};
+
+export default EditPage;

frontend/src/components/home/home-page.tsx → frontend/src/components/pages/home-page.tsx

@@ -32,7 +32,7 @@ function tabTarget(path: string) {
   return `${getSessionId()}-${encodeURIComponent(path)}`;
 }
 
-export const HomePage: React.FC = () => {
+const HomePage: React.FC = () => {
   const fileResponse = useAsyncData(async () => {
     const [workspace, recents] = await Promise.all([
       getWorkspaceFiles(),
@@ -276,3 +276,5 @@ const CreateNewNotebook: React.FC = () => {
 //     </div>
 //   );
 // };
+
+export default HomePage;

frontend/src/components/pages/run-page.tsx

@@ -0,0 +1,19 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+import { AppConfig } from "@/core/config/config-schema";
+import { RunApp } from "@/core/run-app";
+import { StaticBanner } from "../static-html/static-banner";
+
+interface Props {
+  appConfig: AppConfig;
+}
+
+const RunPage = (props: Props) => {
+  return (
+    <>
+      <StaticBanner />
+      <RunApp appConfig={props.appConfig} />
+    </>
+  );
+};
+
+export default RunPage;

frontend/src/core/MarimoApp.tsx

@@ -1,85 +1,90 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 import "../css/index.css";
+import "../css/app/App.css";
 import "iconify-icon";
 
-import { PropsWithChildren, memo, useEffect } from "react";
+import React, { PropsWithChildren, Suspense, memo } from "react";
 import { ErrorBoundary } from "../components/editor/boundary/ErrorBoundary";
-import { initializePlugins } from "../plugins/plugins";
-import { App } from "./App";
 import { TooltipProvider } from "../components/ui/tooltip";
 import { Toaster } from "../components/ui/toaster";
 import { ModalProvider } from "../components/modal/ImperativeModal";
 import { DayPickerProvider } from "react-day-picker";
-import { CommandPalette } from "../components/editor/controls/command-palette";
 import { useAppConfig, useUserConfig } from "@/core/config/config";
 import { initialMode } from "./mode";
-import { AppChrome } from "../components/editor/chrome/wrapper/app-chrome";
-import { StaticBanner } from "../components/static-html/static-banner";
 import { CssVariables } from "@/theme/ThemeProvider";
 import { useAsyncData } from "@/hooks/useAsyncData";
 import { isPyodide } from "./pyodide/utils";
 import { PyodideBridge } from "./pyodide/bridge";
 import { LargeSpinner } from "@/components/icons/large-spinner";
-import { TailwindIndicator } from "@/components/indicator";
-import { HomePage } from "@/components/home/home-page";
+import { TailwindIndicator } from "@/components/debug/indicator";
+
+// Lazy imports
+const LazyHomePage = React.lazy(() => import("@/components/pages/home-page"));
+const LazyRunPage = React.lazy(() => import("@/components/pages/run-page"));
+const LazyEditPage = React.lazy(() => import("@/components/pages/edit-page"));
 
 /**
  * The root component of the Marimo app.
  */
 export const MarimoApp: React.FC = memo(() => {
   const [userConfig] = useUserConfig();
   const [appConfig] = useAppConfig();
+  const editorFontSize = toRem(userConfig.display.code_editor_font_size);
 
-  useEffect(() => {
-    initializePlugins();
-  }, []);
+  const renderBody = () => {
+    if (initialMode === "home") {
+      return <LazyHomePage />;
+    } else if (initialMode === "read") {
+      return <LazyRunPage appConfig={appConfig} />;
+    } else {
+      return <LazyEditPage userConfig={userConfig} appConfig={appConfig} />;
+    }
+  };
 
-  const body =
-    initialMode === "home" ? (
-      <HomePage />
-    ) : initialMode === "read" ? (
-      <>
-        <StaticBanner />
-        <App userConfig={userConfig} appConfig={appConfig} />
-      </>
-    ) : (
-      <AppChrome>
-        <App userConfig={userConfig} appConfig={appConfig} />
-        <CommandPalette />
-      </AppChrome>
-    );
+  return (
+    <Providers>
+      <CssVariables
+        variables={{ "--marimo-code-editor-font-size": editorFontSize }}
+      >
+        {renderBody()}
+      </CssVariables>
+    </Providers>
+  );
+});
+MarimoApp.displayName = "MarimoApp";
 
+/**
+ * The root with all the providers.
+ */
+const Providers = memo(({ children }: PropsWithChildren) => {
   return (
     <ErrorBoundary>
-      <TooltipProvider>
-        <DayPickerProvider initialProps={{}}>
-          <PyodideLoader>
-            <ModalProvider>
-              <CssVariables
-                variables={{
-                  "--marimo-code-editor-font-size": toRem(
-                    userConfig.display.code_editor_font_size,
-                  ),
-                }}
-              >
-                {body}
+      <Suspense>
+        <TooltipProvider>
+          <DayPickerProvider initialProps={{}}>
+            <PyodideLoader>
+              <ModalProvider>
+                {children}
                 <Toaster />
-              </CssVariables>
-              <TailwindIndicator />
-            </ModalProvider>
-          </PyodideLoader>
-        </DayPickerProvider>
-      </TooltipProvider>
+                <TailwindIndicator />
+              </ModalProvider>
+            </PyodideLoader>
+          </DayPickerProvider>
+        </TooltipProvider>
+      </Suspense>
     </ErrorBoundary>
   );
 });
-MarimoApp.displayName = "MarimoApp";
+Providers.displayName = "Providers";
 
 function toRem(px: number) {
   return `${px / 16}rem`;
 }
 
-export const PyodideLoader: React.FC<PropsWithChildren> = ({ children }) => {
+/**
+ * HOC to load Pyodide before rendering children, if necessary.
+ */
+const PyodideLoader: React.FC<PropsWithChildren> = ({ children }) => {
   if (!isPyodide()) {
     return children;
   }

frontend/src/core/cells/__tests__/cells.test.ts

@@ -4,8 +4,8 @@ import {
   NotebookState,
   exportedForTesting,
   flattenNotebookCells,
-  notebookCells,
 } from "../cells";
+import { notebookCells } from "../utils";
 import { CellId } from "@/core/cells/ids";
 import { OutputMessage } from "@/core/kernel/messages";
 import { Seconds } from "@/utils/time";

frontend/src/core/cells/actions.ts

@@ -0,0 +1,25 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import { store } from "../state/jotai";
+import { Objects } from "@/utils/objects";
+import { EditorView } from "@codemirror/view";
+import { getCellEditorView, notebookAtom } from "./cells";
+
+/**
+ * Scroll to the first cell that is currently in "running" state.
+ */
+export function notebookScrollToRunning() {
+  // find cell that is currently in "running" state
+  const { cellRuntime } = store.get(notebookAtom);
+  const cell = Objects.entries(cellRuntime).find(
+    ([cellid, runtimestate]) => runtimestate.status === "running",
+  );
+  if (!cell) {
+    return;
+  }
+  const view = getCellEditorView(cell[0]);
+  view?.dispatch({
+    selection: { anchor: 0 },
+    effects: [EditorView.scrollIntoView(0, { y: "center" })],
+  });
+}