@marimo-team/frontend 0.25.2-dev66 → 0.25.2-dev89

package/src/components/editor/chrome/wrapper/footer-items/machine-stats.tsx

@@ -10,7 +10,7 @@ import { connectionAtom } from "@/core/network/connection";
 import { useRequestClient } from "@/core/network/requests";
 import type { UsageResponse } from "@/core/network/types";
 import { isWasm } from "@/core/wasm/utils";
-import { WebSocketState } from "@/core/websocket/types";
+import { ConnectionState } from "@/core/websocket/types";
 import { useAsyncData } from "@/hooks/useAsyncData";
 import { useInterval } from "@/hooks/useInterval";
 import { cn } from "@/utils/cn";
@@ -29,7 +29,7 @@ export const MachineStats: React.FC = () => {
     if (isWasm()) {
       return null;
     }
-    if (connection.state !== WebSocketState.OPEN) {
+    if (connection.state !== ConnectionState.OPEN) {
       return null;
     }
     return getUsageStats();

package/src/components/editor/controls/Controls.tsx

@@ -21,7 +21,7 @@ import {
   getConnectionTooltip,
   isAppInteractionDisabled,
 } from "@/core/websocket/connection-utils";
-import { WebSocketState } from "@/core/websocket/types";
+import { ConnectionState } from "@/core/websocket/types";
 import { cn } from "@/utils/cn";
 import { Functions } from "@/utils/functions";
 import {
@@ -43,7 +43,7 @@ interface ControlsProps {
   onTogglePresenting: () => void;
   onInterrupt: () => void;
   onRun: () => void;
-  connectionState: WebSocketState;
+  connectionState: ConnectionState;
   running: boolean;
   appConfig: AppConfig;
 }
@@ -60,7 +60,7 @@ export const Controls = ({
   const undoLabel = useAtomValue(undoLabelAtom);
   const needsRun = useAtomValue(needsRunAtom);
   const { undoDeleteCell } = useCellActions();
-  const closed = connectionState === WebSocketState.CLOSED;
+  const closed = connectionState === ConnectionState.CLOSED;
 
   let undoControl: JSX.Element | null = null;
   if (!closed && undoAvailable) {

package/src/components/editor/controls/__tests__/notebook-menu-dropdown.test.tsx

@@ -25,7 +25,7 @@ import { requestClientAtom } from "@/core/network/requests";
 import { filenameAtom } from "@/core/saving/file-state";
 import { store } from "@/core/state/jotai";
 import { isWasm } from "@/core/wasm/utils";
-import { WebSocketState } from "@/core/websocket/types";
+import { ConnectionState } from "@/core/websocket/types";
 import {
   DEFAULT_EXPORT_OPTIONS,
   exportOptionsAtom,
@@ -111,7 +111,7 @@ describe("NotebookMenuDropdown", () => {
     store.set(requestClientAtom, MockRequestClient.create());
     store.set(filenameAtom, "/project/notebook.py");
     store.set(viewStateAtom, { mode: "edit", cellAnchor: null });
-    store.set(connectionAtom, { state: WebSocketState.OPEN });
+    store.set(connectionAtom, { state: ConnectionState.OPEN });
     store.set(kioskModeAtom, false);
     store.set(layoutStateAtom, {
       selectedLayout: "vertical",

package/src/components/editor/header/__tests__/status.test.tsx

@@ -10,7 +10,7 @@ import { viewStateAtom } from "@/core/mode";
 import {
   type ConnectionStatus,
   WebSocketClosedReason,
-  WebSocketState,
+  ConnectionState,
 } from "@/core/websocket/types";
 import { StatusOverlay } from "../status";
 
@@ -40,7 +40,7 @@ describe("StatusOverlay disconnect indicator", () => {
     const onReconnect = vi.fn();
     const { getByTestId } = renderOverlay(
       {
-        state: WebSocketState.CLOSED,
+        state: ConnectionState.CLOSED,
         code: WebSocketClosedReason.KERNEL_DISCONNECTED,
         reason: "kernel not found",
       },
@@ -57,7 +57,7 @@ describe("StatusOverlay disconnect indicator", () => {
 
   it("renders a disabled button when no onReconnect is provided", () => {
     const { getByTestId } = renderOverlay({
-      state: WebSocketState.CLOSED,
+      state: ConnectionState.CLOSED,
       code: WebSocketClosedReason.KERNEL_DISCONNECTED,
       reason: "kernel not found",
     });
@@ -68,7 +68,7 @@ describe("StatusOverlay disconnect indicator", () => {
 
   it("leaves editor startup failures to the inline notice", () => {
     const { queryByTestId } = renderOverlay({
-      state: WebSocketState.CLOSED,
+      state: ConnectionState.CLOSED,
       code: WebSocketClosedReason.KERNEL_STARTUP_ERROR,
       reason: "Failed to start kernel sandbox",
     });

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

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

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

@@ -10,7 +10,7 @@ import { viewStateAtom } from "@/core/mode";
 import {
   type ConnectionStatus,
   WebSocketClosedReason,
-  WebSocketState,
+  ConnectionState,
 } from "@/core/websocket/types";
 import { cn } from "@/utils/cn";
 
@@ -21,10 +21,10 @@ export const StatusOverlay: React.FC<{
 }> = ({ connection, isRunning, onReconnect }) => {
   const { mode } = useAtomValue(viewStateAtom);
   const isClosed =
-    connection.state === WebSocketState.CLOSED &&
+    connection.state === ConnectionState.CLOSED &&
     (mode === "read" ||
       connection.code !== WebSocketClosedReason.KERNEL_STARTUP_ERROR);
-  const isOpen = connection.state === WebSocketState.OPEN;
+  const isOpen = connection.state === ConnectionState.OPEN;
   // Editor startup failures have their own recovery UI.
   const canReconnect =
     isClosed && connection.code === WebSocketClosedReason.KERNEL_DISCONNECTED;

package/src/components/ui/progress-stages.tsx

@@ -0,0 +1,114 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+import { AlertCircleIcon, CheckIcon, CircleIcon } from "lucide-react";
+import type { ReactNode } from "react";
+import { Spinner } from "@/components/icons/spinner";
+import { cn } from "@/utils/cn";
+
+export type ProgressState = "waiting" | "running" | "succeeded" | "failed";
+
+export interface ProgressStage {
+  id: string;
+  state: ProgressState;
+  title: string;
+  description: string;
+  details?: ReactNode;
+}
+
+export function ProgressStatusIcon({
+  state,
+  variant = "icon",
+}: {
+  state: ProgressState;
+  variant?: "icon" | "dot";
+}) {
+  if (variant === "dot") {
+    return (
+      <span
+        aria-hidden={true}
+        className={cn(
+          "size-1.5 shrink-0 rounded-full",
+          {
+            waiting: "bg-muted-foreground",
+            running: "bg-muted-foreground motion-safe:animate-pulse",
+            succeeded: "bg-(--grass-9)",
+            failed: "bg-(--red-9)",
+          }[state],
+        )}
+      />
+    );
+  }
+  const Icon = {
+    waiting: CircleIcon,
+    running: Spinner,
+    succeeded: CheckIcon,
+    failed: AlertCircleIcon,
+  }[state];
+  return (
+    <Icon
+      aria-hidden={true}
+      className={cn(
+        "size-4 shrink-0",
+        {
+          waiting: "text-muted-foreground/50",
+          running: "motion-reduce:animate-none",
+          succeeded: "text-(--grass-11)",
+          failed: "text-(--red-11)",
+        }[state],
+      )}
+    />
+  );
+}
+
+export function ProgressStages({
+  stages,
+  label,
+  className,
+}: {
+  stages: readonly ProgressStage[];
+  label: string;
+  className?: string;
+}) {
+  return (
+    <ol aria-label={label} className={cn("space-y-6", className)}>
+      {stages.map((stage, index) => (
+        <li
+          key={stage.id}
+          aria-current={
+            stage.state === "running" || stage.state === "failed"
+              ? "step"
+              : undefined
+          }
+          className="relative flex items-start gap-3"
+        >
+          {index < stages.length - 1 && (
+            <span
+              aria-hidden={true}
+              className="absolute left-2 top-4.5 -bottom-6.5 border-l border-border"
+            />
+          )}
+          <span className="mt-0.5 text-muted-foreground">
+            <ProgressStatusIcon state={stage.state} />
+          </span>
+          <div className="min-w-0 flex-1">
+            <div
+              className={cn(
+                "text-sm",
+                stage.state === "failed"
+                  ? "text-(--red-11)"
+                  : stage.state === "waiting"
+                    ? "text-muted-foreground"
+                    : "text-foreground",
+              )}
+            >
+              {stage.title}
+            </div>
+            <p className="mt-1 text-xs text-muted-foreground">
+              {stage.description}
+            </p>
+            {stage.details}
+          </div>
+        </li>
+      ))}
+    </ol>
+  );
+}

package/src/core/codemirror/lsp/__tests__/notebook-lsp.test.ts

@@ -359,11 +359,111 @@ describe("NotebookLanguageServerClient", () => {
       return undefined;
     });
 
-    notebookClient = new NotebookLanguageServerClient(mockClient, {}, () => ({
-      [Cells.cell1]: new EditorView({ doc: "# this is a comment" }),
-      [Cells.cell2]: new EditorView({ doc: "import math\nimport numpy" }),
-      [Cells.cell3]: new EditorView({ doc: "print(math.sqrt(4))" }),
-    }));
+    notebookClient = new NotebookLanguageServerClient(
+      mockClient,
+      {},
+      () => ({}),
+    );
+  });
+
+  describe("live editor snapshots", () => {
+    function createEditor(code: string) {
+      return new EditorView({
+        doc: code,
+        extensions: [
+          languageAdapterState.init(() => new PythonLanguageAdapter()),
+          languageMetadataField.init(() => ({})),
+        ],
+      });
+    }
+
+    it("sends edits and their removal without invalidating the notebook atom", async () => {
+      const cachedCodes = store.get(topologicalCodesAtom);
+      const editor = createEditor("print(math.sqrt(4))");
+      replaceNotebookEditors(notebookClient, () => ({ [Cells.cell3]: editor }));
+      const uri = CellDocumentUri.of(Cells.cell3);
+      try {
+        await notebookClient.textDocumentDidOpen({
+          textDocument: {
+            uri,
+            languageId: "python",
+            version: 1,
+            text: editor.state.doc.toString(),
+          },
+        });
+        for (const [index, code] of [
+          "math.missing()",
+          "print(math.sqrt(4))",
+        ].entries()) {
+          editor.dispatch({
+            changes: { from: 0, to: editor.state.doc.length, insert: code },
+          });
+          await notebookClient.textDocumentDidChange({
+            textDocument: { uri, version: index + 2 },
+            contentChanges: [{ text: code }],
+          });
+          expect(store.get(topologicalCodesAtom)).toEqual(cachedCodes);
+          expect(mockClient.textDocumentDidChange).toHaveBeenNthCalledWith(
+            index + 1,
+            {
+              textDocument: {
+                uri: notebookClient.documentUri,
+                version: index + 2,
+              },
+              contentChanges: [
+                {
+                  text: `# this is a comment\nimport math\nimport numpy\n${code}`,
+                },
+              ],
+            },
+          );
+        }
+      } finally {
+        editor.destroy();
+      }
+    });
+
+    it("syncs pending edits before completion and uses their updated line offsets", async () => {
+      const editor = createEditor("import math\nimport numpy");
+      replaceNotebookEditors(notebookClient, () => ({ [Cells.cell2]: editor }));
+      try {
+        await notebookClient.textDocumentDidOpen({
+          textDocument: {
+            uri: CellDocumentUri.of(Cells.cell2),
+            languageId: "python",
+            version: 1,
+            text: editor.state.doc.toString(),
+          },
+        });
+        editor.dispatch({
+          changes: { from: editor.state.doc.length, insert: "\nextra = 1" },
+        });
+        mockClient.textDocumentCompletion.mockImplementation(async () => {
+          expect(mockClient.textDocumentDidChange).toHaveBeenCalledWith({
+            textDocument: { uri: notebookClient.documentUri, version: 2 },
+            contentChanges: [
+              {
+                text: "# this is a comment\nimport math\nimport numpy\nextra = 1\nprint(math.sqrt(4))",
+              },
+            ],
+          });
+          return null;
+        });
+        const params: LSP.CompletionParams = {
+          textDocument: { uri: CellDocumentUri.of(Cells.cell3) },
+          position: { line: 0, character: 11 },
+        };
+        await notebookClient.textDocumentCompletion(params);
+        expect(mockClient.textDocumentCompletion).toHaveBeenCalledWith({
+          textDocument: { uri: notebookClient.documentUri },
+          position: { line: 4, character: 11 },
+        });
+        await notebookClient.textDocumentCompletion(params);
+        expect(mockClient.textDocumentDidChange).toHaveBeenCalledTimes(1);
+      } finally {
+        editor.destroy();
+      }
+    });
   });
 
   describe("document lifecycle", () => {
@@ -1893,6 +1993,7 @@ describe("NotebookLanguageServerClient", () => {
       const isolatedNotebookClient = new NotebookLanguageServerClient(
         isolatedClient,
         {},
+        () => ({}),
       );
       await isolatedNotebookClient.sync();
       isolatedNotebookClient.patchProcessNotification();

package/src/core/codemirror/lsp/notebook-lsp.ts

@@ -343,7 +343,21 @@ export class NotebookLanguageServerClient implements ILanguageServerClient {
   }
 
   private getNotebookCode() {
-    return store.get(topologicalCodesAtom);
+    const { cellIds, codes } = store.get(topologicalCodesAtom);
+    const editors = this.getNotebookEditors();
+    // EditorView edits do not invalidate the cached notebook atom.
+    return {
+      cellIds,
+      codes: Objects.fromEntries(
+        cellIds.map((id) => {
+          const editor = editors[id];
+          return [
+            id,
+            editor ? getEditorCodeAsPython(editor) : (codes[id] ?? ""),
+          ];
+        }),
+      ),
+    };
   }
 
   private assertCellDocumentUri(
@@ -870,7 +884,7 @@ export class NotebookLanguageServerClient implements ILanguageServerClient {
   }
 
   public async textDocumentCompletion(params: LSP.CompletionParams) {
-    const { lens } = this.snapshotter.getLatestSnapshot();
+    const { lens } = await this.sync();
 
     // Check if URI is valid
     if (!CellDocumentUri.is(params.textDocument.uri)) {

package/src/core/codemirror/readonly/__tests__/extension.test.ts

@@ -4,12 +4,12 @@ import { EditorState } from "@codemirror/state";
 import { EditorView } from "@codemirror/view";
 import { createStore } from "jotai";
 import { describe, expect, it } from "vitest";
-import { WebSocketState } from "@/core/websocket/types";
+import { ConnectionState } from "@/core/websocket/types";
 import { connectionAtom } from "../../../network/connection";
 import { dynamicReadonly, isEditorReadonly } from "../extension";
 
 function makeStoreWithConnection(
-  state: typeof WebSocketState.CONNECTING | typeof WebSocketState.OPEN,
+  state: typeof ConnectionState.CONNECTING | typeof ConnectionState.OPEN,
 ) {
   const store = createStore();
   store.set(connectionAtom, { state });
@@ -18,7 +18,7 @@ function makeStoreWithConnection(
 
 describe("dynamicReadonly", () => {
   it("should be readonly when connection is not OPEN", () => {
-    const store = makeStoreWithConnection(WebSocketState.CONNECTING);
+    const store = makeStoreWithConnection(ConnectionState.CONNECTING);
     const state = EditorState.create({
       doc: "test",
       extensions: [dynamicReadonly(store)],
@@ -27,7 +27,7 @@ describe("dynamicReadonly", () => {
   });
 
   it("should be editable when connection is OPEN", async () => {
-    const store = makeStoreWithConnection(WebSocketState.OPEN);
+    const store = makeStoreWithConnection(ConnectionState.OPEN);
     const state = EditorState.create({
       doc: "test",
       extensions: [dynamicReadonly(store)],
@@ -37,7 +37,7 @@ describe("dynamicReadonly", () => {
   });
 
   it("should toggle readonly when connection state changes", () => {
-    const store = makeStoreWithConnection(WebSocketState.CONNECTING);
+    const store = makeStoreWithConnection(ConnectionState.CONNECTING);
     const state = EditorState.create({
       doc: "test",
       extensions: [dynamicReadonly(store)],
@@ -45,7 +45,7 @@ describe("dynamicReadonly", () => {
     const view = new EditorView({ state });
     expect(isEditorReadonly(state)).toBe(true);
     // Simulate connection opening
-    store.set(connectionAtom, { state: WebSocketState.OPEN });
+    store.set(connectionAtom, { state: ConnectionState.OPEN });
     // The extension uses a plugin to update the state, so we need to create a view to trigger it
     view.requestMeasure();
     // The state should now be editable

package/src/core/codemirror/readonly/extension.ts

@@ -7,7 +7,7 @@ import {
   isAppConnecting,
 } from "@/core/websocket/connection-utils";
 import { connectionAtom } from "../../network/connection";
-import { WebSocketState } from "../../websocket/types";
+import { ConnectionState } from "../../websocket/types";
 
 /**
  * State effect for updating readonly status based on connection
@@ -111,7 +111,7 @@ export function dynamicReadonly(store: ReturnType<typeof createStore>) {
     // Subscribe to connectionAtom and manage listeners
     const unsubscribe = store.sub(connectionAtom, () => {
       const connection = store.get(connectionAtom);
-      if (connection.state === WebSocketState.CONNECTING) {
+      if (connection.state === ConnectionState.CONNECTING) {
         attachListeners();
       } else {
         detachListeners();

package/src/core/datasets/__tests__/data-source-discovery.test.ts

@@ -2,7 +2,7 @@
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
 import { connectionAtom } from "@/core/network/connection";
 import { store } from "@/core/state/jotai";
-import { WebSocketState } from "@/core/websocket/types";
+import { ConnectionState } from "@/core/websocket/types";
 import {
   fetchDataSourceDiscovery,
   invalidateDataSourceDiscovery,
@@ -30,7 +30,7 @@ describe("fetchDataSourceDiscovery", () => {
   beforeEach(() => {
     // Discovery waits for the session's websocket connection to be open
     // before requesting, so it doesn't race a still-connecting session.
-    store.set(connectionAtom, { state: WebSocketState.OPEN });
+    store.set(connectionAtom, { state: ConnectionState.OPEN });
   });
 
   afterEach(() => {

package/src/core/edit-app.tsx

@@ -197,7 +197,11 @@ export const EditApp: React.FC<AppProps> = ({
           )}
         </div>
 
-        <ConnectionNotice appConfig={appConfig} onRetry={reconnect} />
+        <ConnectionNotice
+          appConfig={appConfig}
+          onRetry={reconnect}
+          canOpenPackages={!hideControls}
+        />
 
         {/* Don't render until we have a single cell. NotStartedConnectionAlert
             still covers the "no remote runtime started" prompt. */}

package/src/core/export/hooks.ts

@@ -16,7 +16,7 @@ import { connectionAtom } from "../network/connection";
 import { useRequestClient } from "../network/requests";
 import type { UpdateCellOutputsRequest } from "../network/types";
 import { VirtualFileTracker } from "../static/virtual-file-tracker";
-import { WebSocketState } from "../websocket/types";
+import { ConnectionState } from "../websocket/types";
 
 const DELAY = 5000; // 5 seconds;
 
@@ -28,7 +28,7 @@ export function useAutoExport() {
   const htmlEnabled = appConfig.auto_download.includes("html");
   const ipynbEnabled = appConfig.auto_download.includes("ipynb");
 
-  const isConnected = state === WebSocketState.OPEN;
+  const isConnected = state === ConnectionState.OPEN;
   const markdownDisabled = !markdownEnabled || !isConnected;
   const htmlDisabled = !htmlEnabled || !isConnected;
   const ipynbDisabled = !ipynbEnabled || !isConnected;

package/src/core/network/__tests__/connection.test.ts

@@ -3,7 +3,7 @@
 import { afterEach, describe, expect, it } from "vitest";
 import { initialModeAtom } from "@/core/mode";
 import { store } from "@/core/state/jotai";
-import { WebSocketState } from "../../websocket/types";
+import { ConnectionState } from "../../websocket/types";
 import { connectionAtom, waitForConnectionOpenIfNotebook } from "../connection";
 
 const NEVER = Symbol("never");
@@ -20,7 +20,7 @@ function settledOrNever(promise: Promise<unknown>) {
 
 describe("waitForConnectionOpenIfNotebook", () => {
   afterEach(() => {
-    store.set(connectionAtom, { state: WebSocketState.NOT_STARTED });
+    store.set(connectionAtom, { state: ConnectionState.NOT_STARTED });
     store.set(initialModeAtom, undefined);
   });
 
@@ -39,7 +39,7 @@ describe("waitForConnectionOpenIfNotebook", () => {
     const promise = waitForConnectionOpenIfNotebook();
     await expect(settledOrNever(promise)).resolves.toBe(NEVER);
 
-    store.set(connectionAtom, { state: WebSocketState.OPEN });
+    store.set(connectionAtom, { state: ConnectionState.OPEN });
     await expect(settledOrNever(promise)).resolves.not.toBe(NEVER);
   });
 });

package/src/core/network/connection-notice.ts

@@ -5,7 +5,7 @@ import { sandboxAtom, sandboxSyncAtom } from "@/core/packages/sandbox-state";
 import {
   type ConnectionPhase,
   WebSocketClosedReason,
-  WebSocketState,
+  ConnectionState,
 } from "@/core/websocket/types";
 import { connectionAtom } from "./connection";
 
@@ -37,7 +37,7 @@ export const connectionNoticeAtom = atom<ConnectionNotice | null>((get) => {
   if (
     sandbox &&
     (syncing || syncError !== null) &&
-    connection.state === WebSocketState.OPEN
+    connection.state === ConnectionState.OPEN
   ) {
     return {
       kind: "sync",
@@ -56,7 +56,7 @@ export const connectionNoticeAtom = atom<ConnectionNotice | null>((get) => {
       error: syncError,
     };
   }
-  if (connection.state === WebSocketState.CONNECTING) {
+  if (connection.state === ConnectionState.CONNECTING) {
     const messages = {
       "preparing-environment": {
         title: "Preparing environment…",
@@ -88,11 +88,11 @@ export const connectionNoticeAtom = atom<ConnectionNotice | null>((get) => {
       error: null,
     };
   }
-  if (connection.state === WebSocketState.CLOSED) {
+  if (connection.state === ConnectionState.CLOSED) {
     const startupFailed =
       connection.code === WebSocketClosedReason.KERNEL_STARTUP_ERROR;
     const environmentFailed =
-      startupFailed && connection.phase === "preparing-environment";
+      sandbox && startupFailed && connection.phase === "preparing-environment";
     return {
       kind: startupFailed ? "startup" : "connection",
       phase: connection.phase,

package/src/core/network/connection.ts

@@ -4,14 +4,14 @@ import { atomWithReducer } from "jotai/utils";
 import type { NotificationMessageData } from "../kernel/messages";
 import { isNotebookPage } from "../mode";
 import { waitFor } from "../state/jotai";
-import { type ConnectionStatus, WebSocketState } from "../websocket/types";
+import { type ConnectionStatus, ConnectionState } from "../websocket/types";
 
 /**
  * Atom for storing the connection status.
  * Initialized to NOT_STARTED.
  */
 export const connectionAtom = atom<ConnectionStatus>({
-  state: WebSocketState.NOT_STARTED,
+  state: ConnectionState.NOT_STARTED,
 });
 
 type StartupProgress = NotificationMessageData<"startup-progress">;
@@ -33,7 +33,7 @@ export const startupProgressAtom = atomWithReducer<
 
 export function waitForConnectionOpen() {
   return waitFor(connectionAtom, (value) => {
-    return value.state === WebSocketState.OPEN;
+    return value.state === ConnectionState.OPEN;
   });
 }
 
@@ -53,28 +53,28 @@ export function waitForConnectionOpenIfNotebook() {
 
 export const isConnectingAtom = atom((get) => {
   const connection = get(connectionAtom);
-  return connection.state === WebSocketState.CONNECTING;
+  return connection.state === ConnectionState.CONNECTING;
 });
 
 export const isConnectedAtom = atom((get) => {
   const connection = get(connectionAtom);
-  return connection.state === WebSocketState.OPEN;
+  return connection.state === ConnectionState.OPEN;
 });
 
 export const canInteractWithAppAtom = atom((get) => {
   const connection = get(connectionAtom);
   return (
-    connection.state === WebSocketState.OPEN ||
-    connection.state === WebSocketState.NOT_STARTED
+    connection.state === ConnectionState.OPEN ||
+    connection.state === ConnectionState.NOT_STARTED
   );
 });
 
 export const isClosedAtom = atom((get) => {
   const connection = get(connectionAtom);
-  return connection.state === WebSocketState.CLOSED;
+  return connection.state === ConnectionState.CLOSED;
 });
 
 export const isNotStartedAtom = atom((get) => {
   const connection = get(connectionAtom);
-  return connection.state === WebSocketState.NOT_STARTED;
+  return connection.state === ConnectionState.NOT_STARTED;
 });