fix: skew potection token in run-mode when auth is enabled (#5168)

Bug fix where we were using the random skew potection token instead of
one based on the code, in run mode.

Author: mscolnick

marimo/_server/sessions.py

@@ -782,26 +782,34 @@ def __init__(
         # since this will contain config-level overrides
         self._config_manager = config_manager
 
+        def _get_code() -> str:
+            app = file_router.get_single_app_file_manager(
+                default_width=self._config_manager.default_width,
+                default_sql_output=self._config_manager.default_sql_output,
+            ).app
+            return "".join(code for code in app.cell_manager.codes())
+
         # Auth token and Skew-protection token
-        if auth_token is not None:
-            self.auth_token = auth_token
-            self.skew_protection_token = SkewProtectionToken.random()
-        elif mode == SessionMode.EDIT:
+        if mode == SessionMode.EDIT:
             # In edit mode, if no auth token is provided,
             # generate a random token
-            self.auth_token = AuthToken.random()
+            self.auth_token = (
+                AuthToken.random() if auth_token is None else auth_token
+            )
             self.skew_protection_token = SkewProtectionToken.random()
         else:
-            app = file_router.get_single_app_file_manager(
-                default_width=self._config_manager.default_width,
-                default_sql_output=self._config_manager.default_sql_output,
-            ).app
-            codes = "".join(code for code in app.cell_manager.codes())
+            source_code = _get_code()
             # Because run-mode is read-only and we could have multiple
             # servers for the same app (going to sleep or autoscaling),
             # we default to a token based on the app's code
-            self.auth_token = AuthToken.from_code(codes)
-            self.skew_protection_token = SkewProtectionToken.from_code(codes)
+            self.auth_token = (
+                AuthToken.from_code(source_code)
+                if auth_token is None
+                else auth_token
+            )
+            self.skew_protection_token = SkewProtectionToken.from_code(
+                source_code
+            )
 
     def app_manager(self, key: MarimoFileKey) -> AppFileManager:
         """

marimo/_server/tokens.py

@@ -30,7 +30,7 @@ def random() -> AuthToken:
 
     @staticmethod
     def from_code(code: str) -> AuthToken:
-        return AuthToken(str(hash(code)))
+        return AuthToken(str(hash("auth:" + code)))
 
     @staticmethod
     def is_empty(token: AuthToken) -> bool:
@@ -51,7 +51,7 @@ def __init__(self, token: str) -> None:
 
     @staticmethod
     def from_code(code: str) -> SkewProtectionToken:
-        return SkewProtectionToken(str(hash(code)))
+        return SkewProtectionToken(str(hash("skew:" + code)))
 
     @staticmethod
     def random() -> SkewProtectionToken:

pyproject.toml

@@ -280,7 +280,7 @@ extra-dependencies = [
     "ipython~=8.12.3",
     # testing gen ai
     "openai>=1.55.3",
-    "anthropic==0.34.1",
+    "anthropic==0.52.2",
     "google-generativeai==0.8.5",
     "boto3>=1.38.25",
     "litellm>=1.70.0",

tests/_server/test_session_manager.py

@@ -16,6 +16,7 @@
     Session,
     SessionManager,
 )
+from marimo._server.tokens import AuthToken, SkewProtectionToken
 from marimo._types.ids import SessionId
 
 if TYPE_CHECKING:
@@ -254,3 +255,140 @@ async def test_create_session_with_script_config_overrides(
     )
 
     session.close()
+
+
+def test_session_manager_auth_token_edit_mode_with_provided_token():
+    """Test that provided auth token is used in EDIT mode"""
+    provided_token = AuthToken("custom-edit-token")
+    session_manager = SessionManager(
+        file_router=AppFileRouter.new_file(),
+        mode=SessionMode.EDIT,
+        development_mode=False,
+        quiet=False,
+        include_code=True,
+        lsp_server=MagicMock(spec=LspServer),
+        config_manager=get_default_config_manager(current_path=None),
+        cli_args={},
+        argv=None,
+        auth_token=provided_token,
+        redirect_console_to_browser=False,
+        ttl_seconds=None,
+    )
+
+    assert session_manager.auth_token is provided_token
+    assert str(session_manager.auth_token) == "custom-edit-token"
+    assert session_manager.skew_protection_token is not None
+
+
+def test_session_manager_auth_token_edit_mode_without_provided_token():
+    """Test that random auth token is generated in EDIT mode when none provided"""
+    session_manager = SessionManager(
+        file_router=AppFileRouter.new_file(),
+        mode=SessionMode.EDIT,
+        development_mode=False,
+        quiet=False,
+        include_code=True,
+        lsp_server=MagicMock(spec=LspServer),
+        config_manager=get_default_config_manager(current_path=None),
+        cli_args={},
+        argv=None,
+        auth_token=None,
+        redirect_console_to_browser=False,
+        ttl_seconds=None,
+    )
+
+    # Should generate a random token (we can't predict the value, but it should exist)
+    assert session_manager.auth_token is not None
+    assert str(session_manager.auth_token) != ""
+    # Verify it's a random token by checking length (AuthToken.random() uses token_urlsafe(16))
+    assert len(str(session_manager.auth_token)) > 10
+    assert session_manager.skew_protection_token is not None
+
+
+def test_session_manager_auth_token_run_mode_with_provided_token():
+    """Test that provided auth token is used in RUN mode"""
+    provided_token = AuthToken("custom-run-token")
+    session_manager = SessionManager(
+        file_router=AppFileRouter.new_file(),
+        mode=SessionMode.RUN,
+        development_mode=False,
+        quiet=False,
+        include_code=True,
+        lsp_server=MagicMock(spec=LspServer),
+        config_manager=get_default_config_manager(current_path=None),
+        cli_args={},
+        argv=None,
+        auth_token=provided_token,
+        redirect_console_to_browser=False,
+        ttl_seconds=None,
+    )
+
+    assert session_manager.auth_token is provided_token
+    assert str(session_manager.auth_token) == "custom-run-token"
+    assert str(session_manager.skew_protection_token) == str(
+        SkewProtectionToken.from_code("")
+    )
+
+
+def test_session_manager_auth_token_run_mode_without_provided_token(
+    tmp_path: Path,
+):
+    """Test that code-based auth token is generated in RUN mode when none provided"""
+    # Create a simple marimo file
+    notebook_content = dedent(
+        """\
+        import marimo
+
+        app = marimo.App()
+
+        @app.cell
+        def test_cell():
+            "hello"
+            return
+        """
+    )
+
+    file_path = tmp_path / "test_notebook.py"
+    file_path.write_text(notebook_content)
+
+    session_manager = SessionManager(
+        file_router=AppFileRouter.infer(str(file_path)),
+        mode=SessionMode.RUN,
+        development_mode=False,
+        quiet=False,
+        include_code=True,
+        lsp_server=MagicMock(spec=LspServer),
+        config_manager=get_default_config_manager(current_path=None),
+        cli_args={},
+        argv=None,
+        auth_token=None,
+        redirect_console_to_browser=False,
+        ttl_seconds=None,
+    )
+
+    # Should generate a deterministic token based on code
+    assert session_manager.auth_token is not None
+    assert str(session_manager.auth_token) != ""
+    assert str(session_manager.skew_protection_token) != ""
+
+    # Create another session manager with the same code - should have same token
+    session_manager2 = SessionManager(
+        file_router=AppFileRouter.infer(str(file_path)),
+        mode=SessionMode.RUN,
+        development_mode=False,
+        quiet=False,
+        include_code=True,
+        lsp_server=MagicMock(spec=LspServer),
+        config_manager=get_default_config_manager(current_path=None),
+        cli_args={},
+        argv=None,
+        auth_token=None,
+        redirect_console_to_browser=False,
+        ttl_seconds=None,
+    )
+
+    # Should have the same deterministic token
+    assert str(session_manager.auth_token) == str(session_manager2.auth_token)
+    assert str(session_manager.skew_protection_token) == str(
+        session_manager2.skew_protection_token
+    )