chore(asyncio): hardening pass — utilities, deprecated API cleanup, lifespan fix (#9552)

Author: mscolnick

marimo/_server/ai/tools/tool_manager.py

@@ -306,11 +306,7 @@ async def _call_handler(
 
         if inspect.iscoroutinefunction(handler):
             return await handler(arguments)
-        else:
-            # Run sync function in thread pool to avoid blocking
-            return await asyncio.get_event_loop().run_in_executor(
-                None, handler, arguments
-            )
+        return await asyncio.to_thread(handler, arguments)
 
     async def _invoke_mcp_tool(
         self, tool_name: str, arguments: FunctionArgs

marimo/_server/api/endpoints/execution.py

@@ -2,7 +2,6 @@
 from __future__ import annotations
 
 import asyncio
-import contextlib
 from typing import TYPE_CHECKING
 from uuid import uuid4
 
@@ -33,6 +32,7 @@
 from marimo._server.uvicorn_utils import close_uvicorn
 from marimo._server.workspace import MarimoFileKey
 from marimo._types.ids import ConsumerId
+from marimo._utils.asyncio_utils import cancel_and_wait
 
 if TYPE_CHECKING:
     from collections.abc import AsyncGenerator
@@ -335,9 +335,7 @@ async def sse_generator() -> AsyncGenerator[str, None]:
 
                 yield build_done_event(session, listener)
         finally:
-            disconnect_task.cancel()
-            with contextlib.suppress(asyncio.CancelledError):
-                await disconnect_task
+            await cancel_and_wait(disconnect_task)
 
     return StreamingResponse(sse_generator(), media_type="text/event-stream")
 

marimo/_server/api/endpoints/ws_endpoint.py

@@ -369,7 +369,7 @@ def _close() -> None:
                     self.manager.close_session(self.params.session_id)
 
             if session is not None:
-                cancellation_handle = asyncio.get_event_loop().call_later(
+                cancellation_handle = asyncio.get_running_loop().call_later(
                     session.ttl_seconds, _close
                 )
                 self.cancel_close_handle = cancellation_handle

marimo/_server/api/interrupt.py

@@ -21,7 +21,7 @@ class InterruptHandler:
     def __init__(self, quiet: bool, shutdown: Callable[[], None]) -> None:
         self.quiet = quiet
         self.shutdown = shutdown
-        self.loop = asyncio.get_event_loop()
+        self.loop = asyncio.get_running_loop()
         self.original_handler = signal.getsignal(signal.SIGINT)
         self._time_of_last_confirmation: float | None = None
 

marimo/_server/api/lifespans.py

@@ -27,6 +27,7 @@
 from marimo._server.uvicorn_utils import close_uvicorn
 from marimo._server.workspace import NEW_FILE
 from marimo._session.model import SessionMode
+from marimo._utils.asyncio_utils import cancel_and_wait, supervised_task
 from marimo._utils.subprocess import cancel_pending_reaps
 
 if TYPE_CHECKING:
@@ -57,18 +58,15 @@ async def lsp(app: Starlette) -> AsyncIterator[None]:
 
     LOGGER.debug("Language Servers are enabled")
     # Start LSP server in background to avoid blocking server startup
-    task = asyncio.create_task(session_mgr.start_lsp_server())
-    background_tasks.add(task)  # Keep a reference to prevent GC
-    task.add_done_callback(background_tasks.discard)  # Clean up when done
+    task = supervised_task(
+        session_mgr.start_lsp_server(),
+        name="lsp.start",
+        registry=background_tasks,
+    )
 
     yield
 
-    # Shutdown
-    task.cancel()
-    try:
-        await task
-    except asyncio.CancelledError:
-        pass
+    await cancel_and_wait(task)
 
 
 @contextlib.asynccontextmanager
@@ -119,24 +117,30 @@ async def background_connect_mcp_servers() -> MCPClient | None:
             LOGGER.warning(f"Failed to connect MCP servers: {e}")
             return None
 
-    task = asyncio.create_task(background_connect_mcp_servers())
-    background_tasks.add(task)  # Keep a reference to prevent GC
-    task.add_done_callback(background_tasks.discard)  # Clean up when done
+    # Awaited below — opt out of supervisor logging to avoid duplicate logs.
+    task = supervised_task(
+        background_connect_mcp_servers(),
+        name="mcp.connect",
+        registry=background_tasks,
+        on_exception=lambda _exc: None,
+    )
 
     yield
 
-    # Shutdown
-    task.cancel()
+    await cancel_and_wait(task)
+    if task.cancelled():
+        return
+
+    mcp_client = task.result()
+    if not mcp_client:
+        return
+
     try:
-        mcp_client = await task
-        if mcp_client:
-            LOGGER.info("Disconnecting from all MCP servers")
-            await mcp_client.disconnect_from_all_servers()
-            LOGGER.info("Successfully disconnected from all MCP servers")
-    except asyncio.CancelledError:
-        pass
+        LOGGER.info("Disconnecting from all MCP servers")
+        await mcp_client.disconnect_from_all_servers()
+        LOGGER.info("Successfully disconnected from all MCP servers")
     except Exception as e:
-        LOGGER.error(f"Error during MCP cleanup: {e}")
+        LOGGER.error(f"Error during MCP disconnect: {e}")
 
 
 @contextlib.asynccontextmanager

marimo/_server/api/middleware.py

@@ -343,7 +343,7 @@ async def send(
         self, request: _URLRequest, stream: bool = False, max_retries: int = 2
     ) -> _AsyncHTTPResponse:
         del stream
-        loop = asyncio.get_event_loop()
+        loop = asyncio.get_running_loop()
 
         body = await self._collect_body(request)
 

marimo/_server/export/exporter.py

@@ -746,7 +746,7 @@ async def _save_file(
         filepath = export_dir / download_name
 
         # Run blocking file I/O in thread pool
-        loop = asyncio.get_event_loop()
+        loop = asyncio.get_running_loop()
         await loop.run_in_executor(
             self._executor, self._write_file_sync, filepath, content
         )

marimo/_server/session_manager.py

@@ -8,9 +8,8 @@
 
 from __future__ import annotations
 
-import asyncio
 from pathlib import Path
-from typing import TYPE_CHECKING, TypeVar
+from typing import TYPE_CHECKING
 
 from marimo import _loggers
 from marimo._cli.sandbox import SandboxMode
@@ -48,10 +47,11 @@
 from marimo._session.session_repository import SessionRepository
 from marimo._session.types import KernelState
 from marimo._types.ids import ConsumerId, SessionId
+from marimo._utils.asyncio_utils import fire_and_forget
 from marimo._utils.file_watcher import FileWatcherManager
 
 if TYPE_CHECKING:
-    from collections.abc import Awaitable, Coroutine, Mapping
+    from collections.abc import Mapping
 
     from marimo._session.notebook import AppFileManager
 
@@ -241,7 +241,10 @@ def create_session(
         self._repository.add_sync(session_id, session)
 
         # Emit session created event (triggers file watcher attachment, recents, etc.)
-        run_async(self._event_bus.emit_session_created(session))
+        fire_and_forget(
+            self._event_bus.emit_session_created(session),
+            name="session.created",
+        )
 
         return session
 
@@ -333,10 +336,11 @@ def maybe_resume_session(
         if resumed_session:
             # Emit resume event (use new_session_id as both old and new since
             # the strategy already updated it)
-            run_async(
+            fire_and_forget(
                 self._event_bus.emit_session_resumed(
                     resumed_session, new_session_id
-                )
+                ),
+                name="session.resumed",
             )
 
         return resumed_session
@@ -389,7 +393,10 @@ def close_session(self, session_id: SessionId) -> bool:
         if session is None:
             return False
 
-        run_async(self._event_bus.emit_session_closed(session))
+        fire_and_forget(
+            self._event_bus.emit_session_closed(session),
+            name="session.closed",
+        )
 
         session.close()
         return True
@@ -418,28 +425,3 @@ def should_send_code_to_frontend(self) -> bool:
     def get_active_connection_count(self) -> int:
         """Get the number of sessions with active connections."""
         return len(self._repository.get_active_sessions())
-
-
-T = TypeVar("T")
-
-
-def run_async(coro: Coroutine[None, None, T] | Awaitable[T]) -> T:
-    """Run an async coroutine, handling various event loop states.
-
-    1. Event loop is running: create a task
-    2. Event loop exists but not running: run_until_complete
-    3. No event loop: create one with asyncio.run
-    """
-    try:
-        loop = asyncio.get_event_loop()
-        if loop.is_running():
-            # Create a task and return it (fire and forget)
-            # Note: This doesn't wait for completion
-            task = asyncio.create_task(coro)  # type: ignore
-            return task  # type: ignore
-        else:
-            # Run to completion
-            return loop.run_until_complete(coro)  # type: ignore
-    except RuntimeError:
-        # No event loop exists, create one
-        return asyncio.run(coro)  # type: ignore

marimo/_server/utils.py

@@ -2,12 +2,20 @@
 from __future__ import annotations
 
 import asyncio
-import sys
 from typing import TYPE_CHECKING, Any, TypeVar
 
+from marimo._utils.asyncio_utils import initialize_asyncio
+
 if TYPE_CHECKING:
     from collections.abc import Coroutine
 
+__all__ = [
+    "asyncio_run",
+    "initialize_asyncio",
+    "initialize_fd_limit",
+    "initialize_mimetypes",
+]
+
 
 def initialize_mimetypes() -> None:
     import mimetypes
@@ -34,16 +42,6 @@ def initialize_mimetypes() -> None:
     )
 
 
-def initialize_asyncio() -> None:
-    """Platform-specific initialization of asyncio.
-
-    Sessions use the `add_reader()` API, which is only available in the
-    SelectorEventLoop policy; Windows uses the Proactor by default.
-    """
-    if sys.platform == "win32":
-        asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
-
-
 def initialize_fd_limit(limit: int) -> None:
     """Raise the limit on open file descriptors.
 

marimo/_session/extensions/extensions.py

@@ -144,11 +144,12 @@ async def _heartbeat() -> None:
                 await _check_alive()
 
         try:
-            loop = asyncio.get_event_loop()
-            self.heartbeat_task = loop.create_task(_heartbeat())
+            loop = asyncio.get_running_loop()
         except RuntimeError:
-            # This can happen if there is no event loop running
+            # No loop (tests, scripts) — nothing to schedule against.
             self.heartbeat_task = None
+            return
+        self.heartbeat_task = loop.create_task(_heartbeat())
 
     def _stop(self) -> None:
         """Stop the heartbeat monitoring."""