From db08cbed4b6f656ebf39adecde97fade315300bb Mon Sep 17 00:00:00 2001 From: Kiran Gadhave Date: Wed, 24 Jun 2026 13:06:12 -0700 Subject: [PATCH] refactor(server): centralize WebSocket close reasons in an enum Replace the MARIMO_* close-frame literals duplicated across the ws endpoints with a WebSocketCloseReason enum in codes.py, so a typo fails fast instead of silently hitting the frontend's retry path. Also drop two dead frontend reasons (MARIMO_WRONG_KERNEL_ID, MARIMO_MALFORMED_QUERY) that no backend emitter ever sent. --- .../editor/header/__tests__/status.test.tsx | 4 --- .../src/components/editor/header/status.tsx | 5 ++-- .../useMarimoKernelConnection.test.ts | 13 ---------- frontend/src/core/websocket/types.ts | 1 - .../websocket/useMarimoKernelConnection.tsx | 14 ---------- marimo/_server/api/endpoints/terminal.py | 4 +-- .../endpoints/ws/ws_connection_validator.py | 10 +++---- .../api/endpoints/ws/ws_session_connector.py | 7 ++--- marimo/_server/api/endpoints/ws_endpoint.py | 13 ++++++---- marimo/_server/codes.py | 21 ++++++++++++++- tests/_server/test_codes.py | 26 +++++++++++++++++++ 11 files changed, 67 insertions(+), 51 deletions(-) create mode 100644 tests/_server/test_codes.py diff --git a/frontend/src/components/editor/header/__tests__/status.test.tsx b/frontend/src/components/editor/header/__tests__/status.test.tsx index 579eb6234b8..84d90c9709c 100644 --- a/frontend/src/components/editor/header/__tests__/status.test.tsx +++ b/frontend/src/components/editor/header/__tests__/status.test.tsx @@ -67,10 +67,6 @@ describe("StatusOverlay disconnect indicator", () => { }); it.each([ - [ - WebSocketClosedReason.MALFORMED_QUERY, - "the kernel did not recognize a request; please file a bug with marimo", - ], [ WebSocketClosedReason.KERNEL_STARTUP_ERROR, "Failed to start kernel sandbox", diff --git a/frontend/src/components/editor/header/status.tsx b/frontend/src/components/editor/header/status.tsx index 3a5558cfc3a..56e1cbe62ed 100644 --- a/frontend/src/components/editor/header/status.tsx +++ b/frontend/src/components/editor/header/status.tsx @@ -22,9 +22,8 @@ export const StatusOverlay: React.FC<{ const { mode } = useAtomValue(viewStateAtom); const isClosed = connection.state === WebSocketState.CLOSED; const isOpen = connection.state === WebSocketState.OPEN; - // Only KERNEL_DISCONNECTED is recoverable by a retry. Other terminal - // reasons (MALFORMED_QUERY, KERNEL_STARTUP_ERROR) would deterministically - // fail the same way. + // Only KERNEL_DISCONNECTED is recoverable by a retry. KERNEL_STARTUP_ERROR + // would deterministically fail the same way. const canReconnect = isClosed && connection.code === WebSocketClosedReason.KERNEL_DISCONNECTED; diff --git a/frontend/src/core/websocket/__tests__/useMarimoKernelConnection.test.ts b/frontend/src/core/websocket/__tests__/useMarimoKernelConnection.test.ts index 4bb0a999f93..af9cc99d61f 100644 --- a/frontend/src/core/websocket/__tests__/useMarimoKernelConnection.test.ts +++ b/frontend/src/core/websocket/__tests__/useMarimoKernelConnection.test.ts @@ -32,7 +32,6 @@ describe("classifyCloseEvent", () => { describe("terminal closes (server-initiated)", () => { it.each([ - "MARIMO_WRONG_KERNEL_ID", "MARIMO_NO_FILE_KEY", "MARIMO_NO_SESSION_ID", "MARIMO_NO_SESSION", @@ -49,18 +48,6 @@ describe("classifyCloseEvent", () => { } }); - it("MARIMO_MALFORMED_QUERY → terminal but does NOT close transport", () => { - const decision = classify("MARIMO_MALFORMED_QUERY"); - expect(decision.kind).toBe("terminal"); - expect(decision.status).toMatchObject({ - state: WebSocketState.CLOSED, - code: WebSocketClosedReason.MALFORMED_QUERY, - }); - if (decision.kind === "terminal") { - expect(decision.closeTransport).toBe(false); - } - }); - it("MARIMO_KERNEL_STARTUP_ERROR → terminal + closeTransport", () => { const decision = classify("MARIMO_KERNEL_STARTUP_ERROR"); expect(decision.kind).toBe("terminal"); diff --git a/frontend/src/core/websocket/types.ts b/frontend/src/core/websocket/types.ts index 709ff0f9c76..1bb4af05303 100644 --- a/frontend/src/core/websocket/types.ts +++ b/frontend/src/core/websocket/types.ts @@ -14,7 +14,6 @@ export type WebSocketState = export const WebSocketClosedReason = { KERNEL_DISCONNECTED: "KERNEL_DISCONNECTED", - MALFORMED_QUERY: "MALFORMED_QUERY", KERNEL_STARTUP_ERROR: "KERNEL_STARTUP_ERROR", } as const; diff --git a/frontend/src/core/websocket/useMarimoKernelConnection.tsx b/frontend/src/core/websocket/useMarimoKernelConnection.tsx index 3547ea95548..dcb13f93203 100644 --- a/frontend/src/core/websocket/useMarimoKernelConnection.tsx +++ b/frontend/src/core/websocket/useMarimoKernelConnection.tsx @@ -82,12 +82,10 @@ const SUPPORTS_LAZY_KERNELS = true; // (marimo/_server/api/endpoints/ws_endpoint.py and ws/*.py). Keep in sync with // the backend literals. export type CloseReason = - | "MARIMO_WRONG_KERNEL_ID" | "MARIMO_NO_FILE_KEY" | "MARIMO_NO_SESSION_ID" | "MARIMO_NO_SESSION" | "MARIMO_SHUTDOWN" - | "MARIMO_MALFORMED_QUERY" | "MARIMO_KERNEL_STARTUP_ERROR" | typeof TRANSPORT_EXHAUSTED_REASON; @@ -107,7 +105,6 @@ export function classifyCloseEvent(event: { reason?: string }): CloseDecision { reason: "kernel not found", }, }; - case "MARIMO_WRONG_KERNEL_ID": case "MARIMO_NO_FILE_KEY": case "MARIMO_NO_SESSION_ID": case "MARIMO_NO_SESSION": @@ -121,17 +118,6 @@ export function classifyCloseEvent(event: { reason?: string }): CloseDecision { }, closeTransport: true, }; - case "MARIMO_MALFORMED_QUERY": - return { - kind: "terminal", - status: { - state: WebSocketState.CLOSED, - code: WebSocketClosedReason.MALFORMED_QUERY, - reason: - "the kernel did not recognize a request; please file a bug with marimo", - }, - closeTransport: false, - }; case "MARIMO_KERNEL_STARTUP_ERROR": return { kind: "terminal", diff --git a/marimo/_server/api/endpoints/terminal.py b/marimo/_server/api/endpoints/terminal.py index 91a3722dc3f..377b39e807d 100644 --- a/marimo/_server/api/endpoints/terminal.py +++ b/marimo/_server/api/endpoints/terminal.py @@ -16,7 +16,7 @@ from marimo import _loggers from marimo._server.api.auth import validate_auth from marimo._server.api.deps import AppState -from marimo._server.codes import WebSocketCodes +from marimo._server.codes import WebSocketCloseReason, WebSocketCodes from marimo._server.router import APIRouter from marimo._session.model import SessionMode from marimo._utils.asyncio_utils import cancel_and_wait @@ -354,7 +354,7 @@ async def websocket_endpoint(websocket: WebSocket) -> None: if app_state.enable_auth and not validate_auth(websocket): await websocket.close( - WebSocketCodes.UNAUTHORIZED, "MARIMO_UNAUTHORIZED" + WebSocketCodes.UNAUTHORIZED, WebSocketCloseReason.UNAUTHORIZED ) return diff --git a/marimo/_server/api/endpoints/ws/ws_connection_validator.py b/marimo/_server/api/endpoints/ws/ws_connection_validator.py index 21f65a3d038..83f3e014789 100644 --- a/marimo/_server/api/endpoints/ws/ws_connection_validator.py +++ b/marimo/_server/api/endpoints/ws/ws_connection_validator.py @@ -10,7 +10,7 @@ from marimo import _loggers from marimo._server.api.auth import validate_auth from marimo._server.api.deps import AppState -from marimo._server.codes import WebSocketCodes +from marimo._server.codes import WebSocketCloseReason, WebSocketCodes from marimo._server.workspace import MarimoFileKey from marimo._types.ids import SessionId @@ -47,7 +47,7 @@ async def validate_auth(self) -> bool: """ if self.app_state.enable_auth and not validate_auth(self.websocket): await self.websocket.close( - WebSocketCodes.UNAUTHORIZED, "MARIMO_UNAUTHORIZED" + WebSocketCodes.UNAUTHORIZED, WebSocketCloseReason.UNAUTHORIZED ) return False return True @@ -64,7 +64,7 @@ async def extract_connection_params( raw_session_id = self.app_state.query_params(SESSION_QUERY_PARAM_KEY) if raw_session_id is None: await self.websocket.close( - WebSocketCodes.NORMAL_CLOSE, "MARIMO_NO_SESSION_ID" + WebSocketCodes.NORMAL_CLOSE, WebSocketCloseReason.NO_SESSION_ID ) return None @@ -78,7 +78,7 @@ async def extract_connection_params( if file_key is None: await self.websocket.close( - WebSocketCodes.NORMAL_CLOSE, "MARIMO_NO_FILE_KEY" + WebSocketCodes.NORMAL_CLOSE, WebSocketCloseReason.NO_FILE_KEY ) return None @@ -112,7 +112,7 @@ async def extract_file_key_only(self) -> MarimoFileKey | None: if file_key is None: LOGGER.warning("RTC: Closing websocket - no file key") await self.websocket.close( - WebSocketCodes.NORMAL_CLOSE, "MARIMO_NO_FILE_KEY" + WebSocketCodes.NORMAL_CLOSE, WebSocketCloseReason.NO_FILE_KEY ) return None diff --git a/marimo/_server/api/endpoints/ws/ws_session_connector.py b/marimo/_server/api/endpoints/ws/ws_session_connector.py index 304ea7aaf1e..7de0ddd45eb 100644 --- a/marimo/_server/api/endpoints/ws/ws_session_connector.py +++ b/marimo/_server/api/endpoints/ws/ws_session_connector.py @@ -21,7 +21,7 @@ from marimo import _loggers from marimo._messaging.types import NoopStream from marimo._runtime.params import QueryParams -from marimo._server.codes import WebSocketCodes +from marimo._server.codes import WebSocketCloseReason, WebSocketCodes from marimo._server.models.models import InstantiateNotebookRequest from marimo._session.model import ConnectionState, SessionMode @@ -130,7 +130,8 @@ def _connect_kiosk(self) -> tuple[Session, ConnectionType]: if self.manager.mode is not SessionMode.EDIT: LOGGER.debug("Kiosk mode is only supported in edit mode") raise WebSocketDisconnect( - WebSocketCodes.FORBIDDEN, "MARIMO_KIOSK_NOT_ALLOWED" + WebSocketCodes.FORBIDDEN, + WebSocketCloseReason.KIOSK_NOT_ALLOWED, ) # Try to find session by ID first @@ -151,7 +152,7 @@ def _connect_kiosk(self) -> tuple[Session, ConnectionType]: self.params.file_key, ) raise WebSocketDisconnect( - WebSocketCodes.NORMAL_CLOSE, "MARIMO_NO_SESSION" + WebSocketCodes.NORMAL_CLOSE, WebSocketCloseReason.NO_SESSION ) LOGGER.debug("Connecting to kiosk session") diff --git a/marimo/_server/api/endpoints/ws_endpoint.py b/marimo/_server/api/endpoints/ws_endpoint.py index b92a9e167d9..ed544e856dd 100644 --- a/marimo/_server/api/endpoints/ws_endpoint.py +++ b/marimo/_server/api/endpoints/ws_endpoint.py @@ -42,7 +42,7 @@ SessionConnector, is_viewer_connection, ) -from marimo._server.codes import WebSocketCodes +from marimo._server.codes import WebSocketCloseReason, WebSocketCodes from marimo._server.router import APIRouter from marimo._server.rtc.doc import LoroDocManager from marimo._server.session_manager import SessionManager @@ -124,7 +124,8 @@ async def ws_sync( else: LOGGER.warning("RTC: Loro is not installed, closing websocket") await websocket.close( - WebSocketCodes.NORMAL_CLOSE, "MARIMO_LORO_NOT_INSTALLED" + WebSocketCodes.NORMAL_CLOSE, + WebSocketCloseReason.LORO_NOT_INSTALLED, ) return @@ -140,7 +141,9 @@ async def ws_sync( LOGGER.warning( f"RTC: Closing websocket - no session found for file key {file_key}" ) - await websocket.close(WebSocketCodes.FORBIDDEN, "MARIMO_NOT_ALLOWED") + await websocket.close( + WebSocketCodes.FORBIDDEN, WebSocketCloseReason.NOT_ALLOWED + ) return # Handle RTC connection @@ -461,7 +464,7 @@ async def _close_kernel_startup_error(self, error_message: str) -> None: # Then close with simple reason await self._safe_close( WebSocketCodes.UNEXPECTED_ERROR, - "MARIMO_KERNEL_STARTUP_ERROR", + WebSocketCloseReason.KERNEL_STARTUP_ERROR, ) def on_attach(self, session: Session, event_bus: SessionEventBus) -> None: @@ -478,7 +481,7 @@ def on_detach(self) -> None: if is_connected: task = asyncio.create_task( self._safe_close( - WebSocketCodes.NORMAL_CLOSE, "MARIMO_SHUTDOWN" + WebSocketCodes.NORMAL_CLOSE, WebSocketCloseReason.SHUTDOWN ) ) _background_tasks.add(task) diff --git a/marimo/_server/codes.py b/marimo/_server/codes.py index 9e4dce2674b..2e75f0a674f 100644 --- a/marimo/_server/codes.py +++ b/marimo/_server/codes.py @@ -1,7 +1,7 @@ # Copyright 2026 Marimo. All rights reserved. from __future__ import annotations -from enum import IntEnum +from enum import Enum, IntEnum class WebSocketCodes(IntEnum): @@ -10,3 +10,22 @@ class WebSocketCodes(IntEnum): FORBIDDEN = 1008 UNAUTHORIZED = 3000 UNEXPECTED_ERROR = 1011 + + +class WebSocketCloseReason(str, Enum): + """Reasons the backend sends in a WebSocket close frame. + + The value is the raw `CloseEvent.reason` string sent over the wire. The + frontend hand-mirrors the subset it handles in + `frontend/src/core/websocket/useMarimoKernelConnection.tsx`. + """ + + LORO_NOT_INSTALLED = "MARIMO_LORO_NOT_INSTALLED" + NOT_ALLOWED = "MARIMO_NOT_ALLOWED" + UNAUTHORIZED = "MARIMO_UNAUTHORIZED" + NO_SESSION_ID = "MARIMO_NO_SESSION_ID" + NO_FILE_KEY = "MARIMO_NO_FILE_KEY" + NO_SESSION = "MARIMO_NO_SESSION" + KIOSK_NOT_ALLOWED = "MARIMO_KIOSK_NOT_ALLOWED" + KERNEL_STARTUP_ERROR = "MARIMO_KERNEL_STARTUP_ERROR" + SHUTDOWN = "MARIMO_SHUTDOWN" diff --git a/tests/_server/test_codes.py b/tests/_server/test_codes.py new file mode 100644 index 00000000000..0c31887be03 --- /dev/null +++ b/tests/_server/test_codes.py @@ -0,0 +1,26 @@ +# Copyright 2026 Marimo. All rights reserved. +from __future__ import annotations + +from marimo._server.codes import WebSocketCloseReason + + +def test_websocket_close_reason_wire_values() -> None: + """Pin the raw close-frame strings sent over the wire. + + Renaming a member must not silently change the value the frontend + matches against in `useMarimoKernelConnection.tsx`. + """ + assert ( + WebSocketCloseReason.LORO_NOT_INSTALLED == "MARIMO_LORO_NOT_INSTALLED" + ) + assert WebSocketCloseReason.NOT_ALLOWED == "MARIMO_NOT_ALLOWED" + assert WebSocketCloseReason.UNAUTHORIZED == "MARIMO_UNAUTHORIZED" + assert WebSocketCloseReason.NO_SESSION_ID == "MARIMO_NO_SESSION_ID" + assert WebSocketCloseReason.NO_FILE_KEY == "MARIMO_NO_FILE_KEY" + assert WebSocketCloseReason.NO_SESSION == "MARIMO_NO_SESSION" + assert WebSocketCloseReason.KIOSK_NOT_ALLOWED == "MARIMO_KIOSK_NOT_ALLOWED" + assert ( + WebSocketCloseReason.KERNEL_STARTUP_ERROR + == "MARIMO_KERNEL_STARTUP_ERROR" + ) + assert WebSocketCloseReason.SHUTDOWN == "MARIMO_SHUTDOWN"