@marimo-team/frontend 0.23.7-dev13 → 0.23.7-dev15
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- package/dist/assets/__vite-browser-external-DuZehUbK.js +1 -0
- package/dist/assets/__vite-browser-external-Jpm67kL1.js +1 -0
- package/dist/assets/{cell-editor-DOImf428.js → cell-editor-BN3D1_w1.js} +1 -1
- package/dist/assets/{command-palette-CSa-ZC8t.js → command-palette-DuPUk7of.js} +1 -1
- package/dist/assets/{edit-page-CMzqgK6j.js → edit-page-hJAbnd__.js} +6 -6
- package/dist/assets/{hooks-BpoZRNxw.js → hooks-DbPDczTe.js} +1 -1
- package/dist/assets/{index-fXenSTRv.js → index-CH2ZE_q-.js} +3 -3
- package/dist/assets/{layout-D5myrs5s.js → layout-C1dvirgi.js} +2 -2
- package/dist/assets/panels-Z5fVmDRY.js +1 -0
- package/dist/assets/{reveal-component-CJYRD8Yf.js → reveal-component-DLRiYfBj.js} +1 -1
- package/dist/assets/run-page-DDMPBh9O.js +1 -0
- package/dist/assets/{save-worker-CvbUHJh7.js → save-worker-D2iQi-UK.js} +4 -4
- package/dist/assets/{scratchpad-panel-SgUSM06x.js → scratchpad-panel-CUMSD7sH.js} +1 -1
- package/dist/assets/{state-B_kYpAQh.js → state-xEqF8Q3P.js} +3 -3
- package/dist/assets/{useNotebookActions-CrtC0jVZ.js → useNotebookActions-CZ7s2FNR.js} +1 -1
- package/dist/assets/{worker-CF8V9c2V.js → worker-B38WhSlZ.js} +4 -4
- package/dist/assets/ws-BV7dcs53.js +22 -0
- package/dist/index.html +2 -2
- package/package.json +2 -2
- package/src/components/editor/app-container.tsx +7 -1
- package/src/components/editor/header/__tests__/status.test.tsx +108 -0
- package/src/components/editor/header/status.tsx +44 -10
- package/src/core/edit-app.tsx +2 -1
- package/src/core/islands/worker/worker.tsx +3 -2
- package/src/core/run-app.tsx +2 -1
- package/src/core/wasm/__tests__/utils.test.ts +34 -0
- package/src/core/wasm/utils.ts +14 -0
- package/src/core/wasm/worker/bootstrap.ts +3 -2
- package/src/core/wasm/worker/worker.ts +3 -2
- package/src/core/websocket/__tests__/useMarimoKernelConnection.hook.test.tsx +155 -0
- package/src/core/websocket/__tests__/useMarimoKernelConnection.test.ts +137 -0
- package/src/core/websocket/transports/basic.ts +2 -0
- package/src/core/websocket/transports/transport.ts +1 -0
- package/src/core/websocket/useMarimoKernelConnection.tsx +130 -55
- package/src/core/websocket/useWebSocket.tsx +5 -2
- package/dist/assets/__vite-browser-external-D0cSGXjR.js +0 -1
- package/dist/assets/__vite-browser-external-DQc2JVNq.js +0 -1
- package/dist/assets/panels-COAp5_gz.js +0 -1
- package/dist/assets/run-page-CTQJzZGR.js +0 -1
- package/dist/assets/ws-BZQmQxZ7.js +0 -22
|
@@ -0,0 +1,108 @@
|
|
|
1
|
+
/* Copyright 2026 Marimo. All rights reserved. */
|
|
2
|
+
// @vitest-environment jsdom
|
|
3
|
+
|
|
4
|
+
import { fireEvent, render } from "@testing-library/react";
|
|
5
|
+
import { createStore, Provider as JotaiProvider } from "jotai";
|
|
6
|
+
import type React from "react";
|
|
7
|
+
import { describe, expect, it, vi } from "vitest";
|
|
8
|
+
import { TooltipProvider } from "@/components/ui/tooltip";
|
|
9
|
+
import { viewStateAtom } from "@/core/mode";
|
|
10
|
+
import {
|
|
11
|
+
type ConnectionStatus,
|
|
12
|
+
WebSocketClosedReason,
|
|
13
|
+
WebSocketState,
|
|
14
|
+
} from "@/core/websocket/types";
|
|
15
|
+
import { StatusOverlay } from "../status";
|
|
16
|
+
|
|
17
|
+
function renderOverlay(
|
|
18
|
+
connection: ConnectionStatus,
|
|
19
|
+
onReconnect?: () => void,
|
|
20
|
+
): ReturnType<typeof render> {
|
|
21
|
+
const store = createStore();
|
|
22
|
+
store.set(viewStateAtom, { mode: "edit", cellAnchor: null });
|
|
23
|
+
const wrapper: React.FC<React.PropsWithChildren> = ({ children }) => (
|
|
24
|
+
<JotaiProvider store={store}>
|
|
25
|
+
<TooltipProvider>{children}</TooltipProvider>
|
|
26
|
+
</JotaiProvider>
|
|
27
|
+
);
|
|
28
|
+
return render(
|
|
29
|
+
<StatusOverlay
|
|
30
|
+
connection={connection}
|
|
31
|
+
isRunning={false}
|
|
32
|
+
onReconnect={onReconnect}
|
|
33
|
+
/>,
|
|
34
|
+
{ wrapper },
|
|
35
|
+
);
|
|
36
|
+
}
|
|
37
|
+
|
|
38
|
+
describe("StatusOverlay disconnect indicator", () => {
|
|
39
|
+
it("invokes onReconnect when the disconnect icon is clicked", () => {
|
|
40
|
+
const onReconnect = vi.fn();
|
|
41
|
+
const { getByTestId } = renderOverlay(
|
|
42
|
+
{
|
|
43
|
+
state: WebSocketState.CLOSED,
|
|
44
|
+
code: WebSocketClosedReason.KERNEL_DISCONNECTED,
|
|
45
|
+
reason: "kernel not found",
|
|
46
|
+
},
|
|
47
|
+
onReconnect,
|
|
48
|
+
);
|
|
49
|
+
|
|
50
|
+
const icon = getByTestId("disconnected-indicator") as HTMLButtonElement;
|
|
51
|
+
expect(icon.tagName).toBe("BUTTON");
|
|
52
|
+
expect(icon.disabled).toBe(false);
|
|
53
|
+
expect(icon.getAttribute("aria-label")).toBe("Reconnect to app");
|
|
54
|
+
fireEvent.click(icon);
|
|
55
|
+
expect(onReconnect).toHaveBeenCalledTimes(1);
|
|
56
|
+
});
|
|
57
|
+
|
|
58
|
+
it("renders a disabled button when no onReconnect is provided", () => {
|
|
59
|
+
const { getByTestId } = renderOverlay({
|
|
60
|
+
state: WebSocketState.CLOSED,
|
|
61
|
+
code: WebSocketClosedReason.KERNEL_DISCONNECTED,
|
|
62
|
+
reason: "kernel not found",
|
|
63
|
+
});
|
|
64
|
+
|
|
65
|
+
const button = getByTestId("disconnected-indicator");
|
|
66
|
+
expect((button as HTMLButtonElement).disabled).toBe(true);
|
|
67
|
+
});
|
|
68
|
+
|
|
69
|
+
it.each([
|
|
70
|
+
[
|
|
71
|
+
WebSocketClosedReason.MALFORMED_QUERY,
|
|
72
|
+
"the kernel did not recognize a request; please file a bug with marimo",
|
|
73
|
+
],
|
|
74
|
+
[
|
|
75
|
+
WebSocketClosedReason.KERNEL_STARTUP_ERROR,
|
|
76
|
+
"Failed to start kernel sandbox",
|
|
77
|
+
],
|
|
78
|
+
])(
|
|
79
|
+
"renders a disabled button for non-recoverable close reason %s",
|
|
80
|
+
(code, reason) => {
|
|
81
|
+
const onReconnect = vi.fn();
|
|
82
|
+
const { getByTestId } = renderOverlay(
|
|
83
|
+
{ state: WebSocketState.CLOSED, code, reason },
|
|
84
|
+
onReconnect,
|
|
85
|
+
);
|
|
86
|
+
|
|
87
|
+
const button = getByTestId("disconnected-indicator") as HTMLButtonElement;
|
|
88
|
+
expect(button.disabled).toBe(true);
|
|
89
|
+
fireEvent.click(button);
|
|
90
|
+
expect(onReconnect).not.toHaveBeenCalled();
|
|
91
|
+
},
|
|
92
|
+
);
|
|
93
|
+
|
|
94
|
+
it("does not render the disconnect icon when another tab has taken over", () => {
|
|
95
|
+
const onReconnect = vi.fn();
|
|
96
|
+
const { queryByTestId } = renderOverlay(
|
|
97
|
+
{
|
|
98
|
+
state: WebSocketState.CLOSED,
|
|
99
|
+
code: WebSocketClosedReason.ALREADY_RUNNING,
|
|
100
|
+
reason: "another browser tab is already connected to the kernel",
|
|
101
|
+
canTakeover: true,
|
|
102
|
+
},
|
|
103
|
+
onReconnect,
|
|
104
|
+
);
|
|
105
|
+
|
|
106
|
+
expect(queryByTestId("disconnected-indicator")).toBeNull();
|
|
107
|
+
});
|
|
108
|
+
});
|
|
@@ -7,16 +7,26 @@ import { Tooltip } from "@/components/ui/tooltip";
|
|
|
7
7
|
import { notebookScrollToRunning } from "@/core/cells/actions";
|
|
8
8
|
import { onlyScratchpadIsRunningAtom } from "@/core/cells/cells";
|
|
9
9
|
import { viewStateAtom } from "@/core/mode";
|
|
10
|
-
import {
|
|
10
|
+
import {
|
|
11
|
+
type ConnectionStatus,
|
|
12
|
+
WebSocketClosedReason,
|
|
13
|
+
WebSocketState,
|
|
14
|
+
} from "@/core/websocket/types";
|
|
11
15
|
import { cn } from "@/utils/cn";
|
|
12
16
|
|
|
13
17
|
export const StatusOverlay: React.FC<{
|
|
14
18
|
connection: ConnectionStatus;
|
|
15
19
|
isRunning: boolean;
|
|
16
|
-
|
|
20
|
+
onReconnect?: () => void;
|
|
21
|
+
}> = ({ connection, isRunning, onReconnect }) => {
|
|
17
22
|
const { mode } = useAtomValue(viewStateAtom);
|
|
18
23
|
const isClosed = connection.state === WebSocketState.CLOSED;
|
|
19
24
|
const isOpen = connection.state === WebSocketState.OPEN;
|
|
25
|
+
// Only KERNEL_DISCONNECTED is recoverable by a retry. Other terminal
|
|
26
|
+
// reasons (MALFORMED_QUERY, KERNEL_STARTUP_ERROR) would deterministically
|
|
27
|
+
// fail the same way; ALREADY_RUNNING is handled by `LockedIcon` below.
|
|
28
|
+
const canReconnect =
|
|
29
|
+
isClosed && connection.code === WebSocketClosedReason.KERNEL_DISCONNECTED;
|
|
20
30
|
|
|
21
31
|
return (
|
|
22
32
|
<>
|
|
@@ -28,7 +38,11 @@ export const StatusOverlay: React.FC<{
|
|
|
28
38
|
)}
|
|
29
39
|
>
|
|
30
40
|
{isOpen && isRunning && <RunningIcon />}
|
|
31
|
-
{isClosed && !connection.canTakeover &&
|
|
41
|
+
{isClosed && !connection.canTakeover && (
|
|
42
|
+
<DisconnectedIcon
|
|
43
|
+
onReconnect={canReconnect ? onReconnect : undefined}
|
|
44
|
+
/>
|
|
45
|
+
)}
|
|
32
46
|
{isClosed && connection.canTakeover && <LockedIcon />}
|
|
33
47
|
</div>
|
|
34
48
|
</>
|
|
@@ -37,13 +51,33 @@ export const StatusOverlay: React.FC<{
|
|
|
37
51
|
|
|
38
52
|
const topLeftStatus = "print:hidden pointer-events-auto hover:cursor-pointer";
|
|
39
53
|
|
|
40
|
-
const DisconnectedIcon
|
|
41
|
-
|
|
42
|
-
|
|
43
|
-
|
|
44
|
-
|
|
45
|
-
|
|
46
|
-
|
|
54
|
+
const DisconnectedIcon: React.FC<{ onReconnect?: () => void }> = ({
|
|
55
|
+
onReconnect,
|
|
56
|
+
}) => {
|
|
57
|
+
const disabled = !onReconnect;
|
|
58
|
+
return (
|
|
59
|
+
<Tooltip
|
|
60
|
+
content={
|
|
61
|
+
disabled ? "App disconnected" : "App disconnected — click to reconnect"
|
|
62
|
+
}
|
|
63
|
+
>
|
|
64
|
+
{/* Wrapper span keeps the tooltip reachable when the button is
|
|
65
|
+
disabled — a disabled <button> swallows pointer events. */}
|
|
66
|
+
<span tabIndex={disabled ? 0 : -1}>
|
|
67
|
+
<button
|
|
68
|
+
type="button"
|
|
69
|
+
className={cn(topLeftStatus, "bg-transparent border-0 p-0")}
|
|
70
|
+
aria-label={disabled ? "App disconnected" : "Reconnect to app"}
|
|
71
|
+
data-testid="disconnected-indicator"
|
|
72
|
+
onClick={onReconnect}
|
|
73
|
+
disabled={disabled}
|
|
74
|
+
>
|
|
75
|
+
<UnlinkIcon className="w-[25px] h-[25px] text-(--red-11)" />
|
|
76
|
+
</button>
|
|
77
|
+
</span>
|
|
78
|
+
</Tooltip>
|
|
79
|
+
);
|
|
80
|
+
};
|
|
47
81
|
|
|
48
82
|
const LockedIcon = () => (
|
|
49
83
|
<Tooltip content="Notebook locked">
|
package/src/core/edit-app.tsx
CHANGED
|
@@ -79,7 +79,7 @@ export const EditApp: React.FC<AppProps> = ({
|
|
|
79
79
|
};
|
|
80
80
|
}, []);
|
|
81
81
|
|
|
82
|
-
const { connection } = useMarimoKernelConnection({
|
|
82
|
+
const { connection, reconnect } = useMarimoKernelConnection({
|
|
83
83
|
autoInstantiate: userConfig.runtime.auto_instantiate,
|
|
84
84
|
setCells: (cells, layout) => {
|
|
85
85
|
setCells(cells);
|
|
@@ -147,6 +147,7 @@ export const EditApp: React.FC<AppProps> = ({
|
|
|
147
147
|
connection={connection}
|
|
148
148
|
isRunning={isRunning}
|
|
149
149
|
width={appConfig.width}
|
|
150
|
+
onReconnect={reconnect}
|
|
150
151
|
>
|
|
151
152
|
<AppHeader
|
|
152
153
|
connection={connection}
|
|
@@ -9,6 +9,7 @@ import {
|
|
|
9
9
|
} from "rpc-anywhere";
|
|
10
10
|
import type { NotificationPayload } from "@/core/kernel/messages";
|
|
11
11
|
import type { ParentSchema } from "@/core/wasm/rpc";
|
|
12
|
+
import { shouldLoadDuckDBPackages } from "@/core/wasm/utils";
|
|
12
13
|
import { TRANSPORT_ID } from "@/core/wasm/worker/constants";
|
|
13
14
|
import { getPyodideVersion } from "@/core/wasm/worker/getPyodideVersion";
|
|
14
15
|
import { MessageBuffer } from "@/core/wasm/worker/message-buffer";
|
|
@@ -85,8 +86,8 @@ const requestHandler = createRPCRequestHandler({
|
|
|
85
86
|
loadPackages: async (code: string) => {
|
|
86
87
|
await pyodideReadyPromise; // Make sure loading is done
|
|
87
88
|
|
|
88
|
-
if (code
|
|
89
|
-
// Add pandas and duckdb to the code
|
|
89
|
+
if (shouldLoadDuckDBPackages(code)) {
|
|
90
|
+
// Add pandas and duckdb to the code for mo.sql and for remote duckdb sources
|
|
90
91
|
code = `import pandas\n${code}`;
|
|
91
92
|
code = `import duckdb\n${code}`;
|
|
92
93
|
code = `import sqlglot\n${code}`;
|
package/src/core/run-app.tsx
CHANGED
|
@@ -38,7 +38,7 @@ export const RunApp: React.FC<AppProps> = ({ appConfig }) => {
|
|
|
38
38
|
};
|
|
39
39
|
}, []);
|
|
40
40
|
|
|
41
|
-
const { connection } = useMarimoKernelConnection({
|
|
41
|
+
const { connection, reconnect } = useMarimoKernelConnection({
|
|
42
42
|
autoInstantiate: true,
|
|
43
43
|
setCells: setCells,
|
|
44
44
|
sessionId: getSessionId(),
|
|
@@ -84,6 +84,7 @@ export const RunApp: React.FC<AppProps> = ({ appConfig }) => {
|
|
|
84
84
|
connection={connection}
|
|
85
85
|
isRunning={isRunning}
|
|
86
86
|
width={appConfig.width}
|
|
87
|
+
onReconnect={reconnect}
|
|
87
88
|
>
|
|
88
89
|
<AppHeader connection={connection} className="sm:pt-8">
|
|
89
90
|
{galleryHref && (
|
|
@@ -0,0 +1,34 @@
|
|
|
1
|
+
/* Copyright 2026 Marimo. All rights reserved. */
|
|
2
|
+
|
|
3
|
+
import { describe, expect, it } from "vitest";
|
|
4
|
+
import { shouldLoadDuckDBPackages } from "../utils";
|
|
5
|
+
|
|
6
|
+
describe("shouldLoadDuckDBPackages", () => {
|
|
7
|
+
it("loads for mo.sql", () => {
|
|
8
|
+
expect(shouldLoadDuckDBPackages('df = mo.sql("SELECT 1")')).toBe(true);
|
|
9
|
+
});
|
|
10
|
+
|
|
11
|
+
it("loads for duckdb imports and usage", () => {
|
|
12
|
+
expect(shouldLoadDuckDBPackages("import duckdb")).toBe(true);
|
|
13
|
+
expect(shouldLoadDuckDBPackages("from duckdb import sql")).toBe(true);
|
|
14
|
+
expect(shouldLoadDuckDBPackages("import pandas, duckdb")).toBe(true);
|
|
15
|
+
expect(shouldLoadDuckDBPackages("rows = duckdb.sql('SELECT 1')")).toBe(
|
|
16
|
+
true,
|
|
17
|
+
);
|
|
18
|
+
});
|
|
19
|
+
|
|
20
|
+
it("loads when package discovery found duckdb", () => {
|
|
21
|
+
expect(
|
|
22
|
+
shouldLoadDuckDBPackages("print('hello')", new Set(["duckdb"])),
|
|
23
|
+
).toBe(true);
|
|
24
|
+
});
|
|
25
|
+
|
|
26
|
+
it("does not load for incidental duckdb text", () => {
|
|
27
|
+
expect(shouldLoadDuckDBPackages("name = 'duckdb'")).toBe(false);
|
|
28
|
+
expect(shouldLoadDuckDBPackages("# import duckdb")).toBe(false);
|
|
29
|
+
});
|
|
30
|
+
|
|
31
|
+
it("does not load without mo.sql, duckdb usage, or discovery", () => {
|
|
32
|
+
expect(shouldLoadDuckDBPackages("print('hello')")).toBe(false);
|
|
33
|
+
});
|
|
34
|
+
});
|
package/src/core/wasm/utils.ts
CHANGED
|
@@ -10,3 +10,17 @@ export function isWasm(): boolean {
|
|
|
10
10
|
document.querySelector("marimo-wasm") !== null
|
|
11
11
|
);
|
|
12
12
|
}
|
|
13
|
+
|
|
14
|
+
const DUCKDB_USAGE_PATTERN =
|
|
15
|
+
/(^|\n)\s*(?:import\s+[^\n#]*\bduckdb\b|from\s+duckdb\b|[^\n#]*\bduckdb\s*\.)/;
|
|
16
|
+
|
|
17
|
+
export function shouldLoadDuckDBPackages(
|
|
18
|
+
code: string,
|
|
19
|
+
foundPackages?: ReadonlySet<string>,
|
|
20
|
+
): boolean {
|
|
21
|
+
return (
|
|
22
|
+
code.includes("mo.sql") ||
|
|
23
|
+
DUCKDB_USAGE_PATTERN.test(code) ||
|
|
24
|
+
foundPackages?.has("duckdb") === true
|
|
25
|
+
);
|
|
26
|
+
}
|
|
@@ -9,6 +9,7 @@ import { WasmFileSystem } from "./fs";
|
|
|
9
9
|
import { getMarimoWheel } from "./getMarimoWheel";
|
|
10
10
|
import { t } from "./tracer";
|
|
11
11
|
import type { SerializedBridge, WasmController } from "./types";
|
|
12
|
+
import { shouldLoadDuckDBPackages } from "../utils";
|
|
12
13
|
|
|
13
14
|
const MAKE_SNAPSHOT = false;
|
|
14
15
|
|
|
@@ -163,8 +164,8 @@ export class DefaultWasmController implements WasmController {
|
|
|
163
164
|
private async loadNotebookDeps(code: string, foundPackages: Set<string>) {
|
|
164
165
|
const pyodide = this.requirePyodide;
|
|
165
166
|
|
|
166
|
-
if (code
|
|
167
|
-
// We need pandas and duckdb for mo.sql
|
|
167
|
+
if (shouldLoadDuckDBPackages(code, foundPackages)) {
|
|
168
|
+
// We need pandas and duckdb for mo.sql and for remote duckdb sources
|
|
168
169
|
code = `import pandas\n${code}`;
|
|
169
170
|
code = `import duckdb\n${code}`;
|
|
170
171
|
code = `import sqlglot\n${code}`;
|
|
@@ -34,6 +34,7 @@ import type {
|
|
|
34
34
|
SerializedBridge,
|
|
35
35
|
WasmController,
|
|
36
36
|
} from "./types";
|
|
37
|
+
import { shouldLoadDuckDBPackages } from "../utils";
|
|
37
38
|
|
|
38
39
|
/**
|
|
39
40
|
* Web worker responsible for running the notebook.
|
|
@@ -141,8 +142,8 @@ const requestHandler = createRPCRequestHandler({
|
|
|
141
142
|
const span = t.startSpan("loadPackages");
|
|
142
143
|
await pyodideReadyPromise; // Make sure loading is done
|
|
143
144
|
|
|
144
|
-
if (code
|
|
145
|
-
// Add pandas and duckdb to the code
|
|
145
|
+
if (shouldLoadDuckDBPackages(code)) {
|
|
146
|
+
// Add pandas and duckdb to the code for mo.sql and for remote duckdb sources
|
|
146
147
|
code = `import pandas\n${code}`;
|
|
147
148
|
code = `import duckdb\n${code}`;
|
|
148
149
|
code = `import sqlglot\n${code}`;
|
|
@@ -0,0 +1,155 @@
|
|
|
1
|
+
/* Copyright 2026 Marimo. All rights reserved. */
|
|
2
|
+
// @vitest-environment jsdom
|
|
3
|
+
|
|
4
|
+
import { act, renderHook } from "@testing-library/react";
|
|
5
|
+
import { createStore, Provider as JotaiProvider } from "jotai";
|
|
6
|
+
import type React from "react";
|
|
7
|
+
import { ErrorBoundary } from "react-error-boundary";
|
|
8
|
+
import { beforeEach, describe, expect, it, vi } from "vitest";
|
|
9
|
+
|
|
10
|
+
vi.mock("@/core/websocket/useWebSocket", async () => {
|
|
11
|
+
const actual =
|
|
12
|
+
await vi.importActual<typeof import("../useWebSocket")>("../useWebSocket");
|
|
13
|
+
return {
|
|
14
|
+
...actual,
|
|
15
|
+
useConnectionTransport: vi.fn(),
|
|
16
|
+
};
|
|
17
|
+
});
|
|
18
|
+
|
|
19
|
+
vi.mock("@/core/runtime/config", async () => {
|
|
20
|
+
const actual = await vi.importActual<typeof import("@/core/runtime/config")>(
|
|
21
|
+
"@/core/runtime/config",
|
|
22
|
+
);
|
|
23
|
+
return {
|
|
24
|
+
...actual,
|
|
25
|
+
useRuntimeManager: vi.fn(),
|
|
26
|
+
};
|
|
27
|
+
});
|
|
28
|
+
|
|
29
|
+
import { useRuntimeManager } from "@/core/runtime/config";
|
|
30
|
+
import { connectionAtom } from "../../network/connection";
|
|
31
|
+
import type { SessionId } from "../../kernel/session";
|
|
32
|
+
import { WebSocketClosedReason, WebSocketState } from "../types";
|
|
33
|
+
import { useMarimoKernelConnection } from "../useMarimoKernelConnection";
|
|
34
|
+
import { useConnectionTransport } from "../useWebSocket";
|
|
35
|
+
|
|
36
|
+
interface MockTransport {
|
|
37
|
+
readyState: 0 | 1 | 2 | 3;
|
|
38
|
+
retryCount: number;
|
|
39
|
+
reconnect: ReturnType<typeof vi.fn>;
|
|
40
|
+
close: ReturnType<typeof vi.fn>;
|
|
41
|
+
send: ReturnType<typeof vi.fn>;
|
|
42
|
+
addEventListener: ReturnType<typeof vi.fn>;
|
|
43
|
+
removeEventListener: ReturnType<typeof vi.fn>;
|
|
44
|
+
}
|
|
45
|
+
|
|
46
|
+
function makeTransport(
|
|
47
|
+
readyState: 0 | 1 | 2 | 3 = WebSocket.CLOSED,
|
|
48
|
+
): MockTransport {
|
|
49
|
+
return {
|
|
50
|
+
readyState,
|
|
51
|
+
retryCount: 0,
|
|
52
|
+
reconnect: vi.fn(),
|
|
53
|
+
close: vi.fn(),
|
|
54
|
+
send: vi.fn(),
|
|
55
|
+
addEventListener: vi.fn(),
|
|
56
|
+
removeEventListener: vi.fn(),
|
|
57
|
+
};
|
|
58
|
+
}
|
|
59
|
+
|
|
60
|
+
function makeRuntimeManager(isHealthy = vi.fn().mockResolvedValue(true)) {
|
|
61
|
+
return {
|
|
62
|
+
isHealthy,
|
|
63
|
+
getWsURL: () => new URL("ws://localhost/ws"),
|
|
64
|
+
waitForHealthy: vi.fn().mockResolvedValue(undefined),
|
|
65
|
+
isSameOrigin: true,
|
|
66
|
+
};
|
|
67
|
+
}
|
|
68
|
+
|
|
69
|
+
describe("useMarimoKernelConnection.reconnect()", () => {
|
|
70
|
+
let transport: MockTransport;
|
|
71
|
+
let isHealthy: ReturnType<typeof vi.fn>;
|
|
72
|
+
let store: ReturnType<typeof createStore>;
|
|
73
|
+
|
|
74
|
+
beforeEach(() => {
|
|
75
|
+
transport = makeTransport(WebSocket.CLOSED);
|
|
76
|
+
isHealthy = vi.fn().mockResolvedValue(true);
|
|
77
|
+
store = createStore();
|
|
78
|
+
store.set(connectionAtom, {
|
|
79
|
+
state: WebSocketState.CLOSED,
|
|
80
|
+
code: WebSocketClosedReason.KERNEL_DISCONNECTED,
|
|
81
|
+
reason: "kernel not found",
|
|
82
|
+
});
|
|
83
|
+
vi.mocked(useConnectionTransport).mockReturnValue(transport);
|
|
84
|
+
vi.mocked(useRuntimeManager).mockReturnValue(
|
|
85
|
+
makeRuntimeManager(isHealthy) as unknown as ReturnType<
|
|
86
|
+
typeof useRuntimeManager
|
|
87
|
+
>,
|
|
88
|
+
);
|
|
89
|
+
});
|
|
90
|
+
|
|
91
|
+
function renderUseHook() {
|
|
92
|
+
const wrapper: React.FC<React.PropsWithChildren> = ({ children }) => (
|
|
93
|
+
<JotaiProvider store={store}>
|
|
94
|
+
<ErrorBoundary fallback={null}>{children}</ErrorBoundary>
|
|
95
|
+
</JotaiProvider>
|
|
96
|
+
);
|
|
97
|
+
return renderHook(
|
|
98
|
+
() =>
|
|
99
|
+
useMarimoKernelConnection({
|
|
100
|
+
sessionId: "test-session" as SessionId,
|
|
101
|
+
autoInstantiate: false,
|
|
102
|
+
setCells: () => {},
|
|
103
|
+
}),
|
|
104
|
+
{ wrapper },
|
|
105
|
+
);
|
|
106
|
+
}
|
|
107
|
+
|
|
108
|
+
it("is a no-op when the transport is already OPEN", async () => {
|
|
109
|
+
transport.readyState = WebSocket.OPEN;
|
|
110
|
+
const { result } = renderUseHook();
|
|
111
|
+
await act(async () => {
|
|
112
|
+
await result.current.reconnect();
|
|
113
|
+
});
|
|
114
|
+
expect(isHealthy).not.toHaveBeenCalled();
|
|
115
|
+
expect(transport.reconnect).not.toHaveBeenCalled();
|
|
116
|
+
});
|
|
117
|
+
|
|
118
|
+
it("is a no-op when the transport is already CONNECTING", async () => {
|
|
119
|
+
transport.readyState = WebSocket.CONNECTING;
|
|
120
|
+
const { result } = renderUseHook();
|
|
121
|
+
await act(async () => {
|
|
122
|
+
await result.current.reconnect();
|
|
123
|
+
});
|
|
124
|
+
expect(isHealthy).not.toHaveBeenCalled();
|
|
125
|
+
expect(transport.reconnect).not.toHaveBeenCalled();
|
|
126
|
+
});
|
|
127
|
+
|
|
128
|
+
it("probes /health and reconnects when the runtime is healthy", async () => {
|
|
129
|
+
isHealthy.mockResolvedValue(true);
|
|
130
|
+
const { result } = renderUseHook();
|
|
131
|
+
await act(async () => {
|
|
132
|
+
await result.current.reconnect();
|
|
133
|
+
});
|
|
134
|
+
expect(isHealthy).toHaveBeenCalledOnce();
|
|
135
|
+
expect(transport.reconnect).toHaveBeenCalledOnce();
|
|
136
|
+
expect(store.get(connectionAtom)).toEqual({
|
|
137
|
+
state: WebSocketState.CONNECTING,
|
|
138
|
+
});
|
|
139
|
+
});
|
|
140
|
+
|
|
141
|
+
it("transitions to CLOSED and does not call ws.reconnect when the probe fails", async () => {
|
|
142
|
+
isHealthy.mockResolvedValue(false);
|
|
143
|
+
const { result } = renderUseHook();
|
|
144
|
+
await act(async () => {
|
|
145
|
+
await result.current.reconnect();
|
|
146
|
+
});
|
|
147
|
+
expect(isHealthy).toHaveBeenCalledOnce();
|
|
148
|
+
expect(transport.reconnect).not.toHaveBeenCalled();
|
|
149
|
+
expect(store.get(connectionAtom)).toEqual({
|
|
150
|
+
state: WebSocketState.CLOSED,
|
|
151
|
+
code: WebSocketClosedReason.KERNEL_DISCONNECTED,
|
|
152
|
+
reason: "kernel not found",
|
|
153
|
+
});
|
|
154
|
+
});
|
|
155
|
+
});
|
|
@@ -0,0 +1,137 @@
|
|
|
1
|
+
/* Copyright 2026 Marimo. All rights reserved. */
|
|
2
|
+
|
|
3
|
+
import { afterEach, describe, expect, it, vi } from "vitest";
|
|
4
|
+
import { Logger } from "@/utils/Logger";
|
|
5
|
+
import { WebSocketClosedReason, WebSocketState } from "../types";
|
|
6
|
+
import { classifyCloseEvent } from "../useMarimoKernelConnection";
|
|
7
|
+
import { MAX_RETRIES } from "../useWebSocket";
|
|
8
|
+
|
|
9
|
+
function classify(
|
|
10
|
+
reason: string | undefined,
|
|
11
|
+
retryCount = 0,
|
|
12
|
+
maxRetries = MAX_RETRIES,
|
|
13
|
+
) {
|
|
14
|
+
return classifyCloseEvent({ reason }, { retryCount, maxRetries });
|
|
15
|
+
}
|
|
16
|
+
|
|
17
|
+
describe("classifyCloseEvent", () => {
|
|
18
|
+
describe("transient closes (default branch)", () => {
|
|
19
|
+
it("retries when retryCount < maxRetries", () => {
|
|
20
|
+
const decision = classify(undefined, 0);
|
|
21
|
+
expect(decision.kind).toBe("retry");
|
|
22
|
+
expect(decision.status).toEqual({ state: WebSocketState.CONNECTING });
|
|
23
|
+
});
|
|
24
|
+
|
|
25
|
+
it("retries on each intermediate close event during a retry storm", () => {
|
|
26
|
+
for (let n = 0; n < MAX_RETRIES; n++) {
|
|
27
|
+
const decision = classify(undefined, n);
|
|
28
|
+
expect(decision.kind).toBe("retry");
|
|
29
|
+
expect(decision.status).toEqual({ state: WebSocketState.CONNECTING });
|
|
30
|
+
}
|
|
31
|
+
});
|
|
32
|
+
|
|
33
|
+
it("transitions to CLOSED when retryCount reaches maxRetries", () => {
|
|
34
|
+
const decision = classify(undefined, MAX_RETRIES);
|
|
35
|
+
expect(decision.kind).toBe("gave-up");
|
|
36
|
+
expect(decision.status).toEqual({
|
|
37
|
+
state: WebSocketState.CLOSED,
|
|
38
|
+
code: WebSocketClosedReason.KERNEL_DISCONNECTED,
|
|
39
|
+
reason: "kernel not found",
|
|
40
|
+
});
|
|
41
|
+
});
|
|
42
|
+
|
|
43
|
+
it("transitions to CLOSED when retryCount exceeds maxRetries", () => {
|
|
44
|
+
const decision = classify(undefined, MAX_RETRIES + 5);
|
|
45
|
+
expect(decision.kind).toBe("gave-up");
|
|
46
|
+
});
|
|
47
|
+
|
|
48
|
+
it("treats unknown reason strings as transient and logs a warning", () => {
|
|
49
|
+
const logger = vi.spyOn(Logger, "warn").mockImplementation(() => {});
|
|
50
|
+
const decision = classify("something-else", 3);
|
|
51
|
+
expect(decision.kind).toBe("retry");
|
|
52
|
+
expect(logger).toHaveBeenCalled();
|
|
53
|
+
logger.mockRestore();
|
|
54
|
+
});
|
|
55
|
+
});
|
|
56
|
+
|
|
57
|
+
afterEach(() => {
|
|
58
|
+
vi.restoreAllMocks();
|
|
59
|
+
});
|
|
60
|
+
|
|
61
|
+
describe("terminal closes (server-initiated)", () => {
|
|
62
|
+
it("MARIMO_ALREADY_CONNECTED → terminal + closeTransport, with takeover", () => {
|
|
63
|
+
const decision = classify("MARIMO_ALREADY_CONNECTED", 0);
|
|
64
|
+
expect(decision.kind).toBe("terminal");
|
|
65
|
+
expect(decision.status).toMatchObject({
|
|
66
|
+
state: WebSocketState.CLOSED,
|
|
67
|
+
code: WebSocketClosedReason.ALREADY_RUNNING,
|
|
68
|
+
canTakeover: true,
|
|
69
|
+
});
|
|
70
|
+
if (decision.kind === "terminal") {
|
|
71
|
+
expect(decision.closeTransport).toBe(true);
|
|
72
|
+
}
|
|
73
|
+
});
|
|
74
|
+
|
|
75
|
+
it.each([
|
|
76
|
+
"MARIMO_WRONG_KERNEL_ID",
|
|
77
|
+
"MARIMO_NO_FILE_KEY",
|
|
78
|
+
"MARIMO_NO_SESSION_ID",
|
|
79
|
+
"MARIMO_NO_SESSION",
|
|
80
|
+
"MARIMO_SHUTDOWN",
|
|
81
|
+
])("%s → terminal with KERNEL_DISCONNECTED, closes transport", (reason) => {
|
|
82
|
+
const decision = classify(reason, 0);
|
|
83
|
+
expect(decision.kind).toBe("terminal");
|
|
84
|
+
expect(decision.status).toMatchObject({
|
|
85
|
+
state: WebSocketState.CLOSED,
|
|
86
|
+
code: WebSocketClosedReason.KERNEL_DISCONNECTED,
|
|
87
|
+
});
|
|
88
|
+
if (decision.kind === "terminal") {
|
|
89
|
+
expect(decision.closeTransport).toBe(true);
|
|
90
|
+
}
|
|
91
|
+
});
|
|
92
|
+
|
|
93
|
+
it("MARIMO_MALFORMED_QUERY → terminal but does NOT close transport", () => {
|
|
94
|
+
const decision = classify("MARIMO_MALFORMED_QUERY", 0);
|
|
95
|
+
expect(decision.kind).toBe("terminal");
|
|
96
|
+
expect(decision.status).toMatchObject({
|
|
97
|
+
state: WebSocketState.CLOSED,
|
|
98
|
+
code: WebSocketClosedReason.MALFORMED_QUERY,
|
|
99
|
+
});
|
|
100
|
+
if (decision.kind === "terminal") {
|
|
101
|
+
expect(decision.closeTransport).toBe(false);
|
|
102
|
+
}
|
|
103
|
+
});
|
|
104
|
+
|
|
105
|
+
it("MARIMO_KERNEL_STARTUP_ERROR → terminal + closeTransport", () => {
|
|
106
|
+
const decision = classify("MARIMO_KERNEL_STARTUP_ERROR", 0);
|
|
107
|
+
expect(decision.kind).toBe("terminal");
|
|
108
|
+
expect(decision.status).toMatchObject({
|
|
109
|
+
state: WebSocketState.CLOSED,
|
|
110
|
+
code: WebSocketClosedReason.KERNEL_STARTUP_ERROR,
|
|
111
|
+
});
|
|
112
|
+
if (decision.kind === "terminal") {
|
|
113
|
+
expect(decision.closeTransport).toBe(true);
|
|
114
|
+
}
|
|
115
|
+
});
|
|
116
|
+
|
|
117
|
+
it("terminal closes ignore retryCount entirely", () => {
|
|
118
|
+
const decision = classify("MARIMO_SHUTDOWN", 99);
|
|
119
|
+
expect(decision.kind).toBe("terminal");
|
|
120
|
+
});
|
|
121
|
+
});
|
|
122
|
+
|
|
123
|
+
describe("retry budget exhaustion", () => {
|
|
124
|
+
it("yields retry on attempts 1..maxRetries-1 and gave-up on the final close", () => {
|
|
125
|
+
const states: string[] = [];
|
|
126
|
+
for (let attempt = 1; attempt <= MAX_RETRIES; attempt++) {
|
|
127
|
+
states.push(classify(undefined, attempt - 1).kind);
|
|
128
|
+
}
|
|
129
|
+
states.push(classify(undefined, MAX_RETRIES).kind);
|
|
130
|
+
|
|
131
|
+
expect(states).toEqual([
|
|
132
|
+
...Array.from({ length: MAX_RETRIES }, () => "retry"),
|
|
133
|
+
"gave-up",
|
|
134
|
+
]);
|
|
135
|
+
});
|
|
136
|
+
});
|
|
137
|
+
});
|