chore(asyncio): migrate remaining background-task sites to asyncio_utils (#9596)

Author: mscolnick

marimo/_server/ai/mcp/client.py

@@ -20,6 +20,7 @@
     MCPTransportRegistry,
 )
 from marimo._server.ai.mcp.types import MCPToolArgs
+from marimo._utils.asyncio_utils import cancel_and_wait, supervised_task
 
 if TYPE_CHECKING:
     from anyio.streams.memory import (
@@ -262,8 +263,9 @@ async def _connection_lifecycle(self, server_name: str) -> None:
                 await self._discover_tools(connection)
 
                 if server_name not in self.health_check_tasks:
-                    self.health_check_tasks[server_name] = asyncio.create_task(
-                        self._monitor_server_health(server_name)
+                    self.health_check_tasks[server_name] = supervised_task(
+                        self._monitor_server_health(server_name),
+                        name=f"mcp.health.{server_name}",
                     )
 
                 # Signal that connection is established
@@ -331,9 +333,12 @@ async def connect_to_server(self, server_name: str) -> bool:
             self._update_server_status(server_name, MCPServerStatus.CONNECTING)
             self._remove_server_tools(server_name)
 
-            # Create task to run existing connection logic
+            # Create task to run existing connection logic. Not supervised:
+            # this task is awaited in disconnect_from_server(), so supervisor
+            # logging would duplicate the awaiter's error handling.
             connection_task = asyncio.create_task(
-                self._connection_lifecycle(server_name)
+                self._connection_lifecycle(server_name),
+                name=f"mcp.lifecycle.{server_name}",
             )
             connection.connection_task = connection_task
 
@@ -790,12 +795,7 @@ async def _cancel_health_monitoring(
         if server_name is not None:
             # Cancel single server monitoring
             if server_name in self.health_check_tasks:
-                task = self.health_check_tasks[server_name]
-                task.cancel()
-                try:
-                    await task
-                except asyncio.CancelledError:
-                    pass
+                await cancel_and_wait(self.health_check_tasks[server_name])
                 del self.health_check_tasks[server_name]
                 LOGGER.debug(f"Cancelled health monitoring for {server_name}")
         else:

marimo/_server/api/endpoints/terminal.py

@@ -19,6 +19,7 @@
 from marimo._server.codes import WebSocketCodes
 from marimo._server.router import APIRouter
 from marimo._session.model import SessionMode
+from marimo._utils.asyncio_utils import cancel_and_wait
 from marimo._utils.platform import is_pyodide, is_windows
 
 if TYPE_CHECKING:
@@ -331,12 +332,7 @@ async def _write_to_pty(
 
 async def _cancel_tasks(tasks: Iterable[asyncio.Task[Any]]) -> None:
     for task in tasks:
-        if not task.done():
-            task.cancel()
-            try:
-                await task
-            except asyncio.CancelledError:
-                pass
+        await cancel_and_wait(task)
 
 
 def supports_terminal() -> bool:

marimo/_server/api/middleware.py

@@ -41,6 +41,7 @@
 from marimo._server.uvicorn_utils import close_uvicorn
 from marimo._session.model import SessionMode
 from marimo._tracer import server_tracer
+from marimo._utils.asyncio_utils import supervised_task
 from marimo._utils.print import print_tabbed
 
 if TYPE_CHECKING:
@@ -700,8 +701,9 @@ def __init__(
         self.app_state.timeout_tracker = time.time()
         self.timeout_duration_minutes = timeout_duration_minutes
 
-        # Hold a strong reference so the monitor task isn't GC'd.
-        self._monitor_task = asyncio.create_task(self.monitor())
+        self._monitor_task = supervised_task(
+            self.monitor(), name="timeout.monitor"
+        )
 
     async def __call__(
         self, scope: Scope, receive: Receive, send: Send

marimo/_server/lsp.py

@@ -25,6 +25,7 @@
     LspServerStatus,
 )
 from marimo._tracer import server_tracer
+from marimo._utils.asyncio_utils import supervised_task
 from marimo._utils.net import find_free_port
 from marimo._utils.paths import marimo_package_path
 from marimo._utils.platform import is_windows
@@ -178,8 +179,9 @@ async def _start_internal(self) -> AlertNotification | None:
                 self._health_check_task is None
                 or self._health_check_task.done()
             ):
-                self._health_check_task = asyncio.create_task(
-                    self._monitor_process_health()
+                self._health_check_task = supervised_task(
+                    self._monitor_process_health(),
+                    name=f"lsp.health.{self.id}",
                 )
 
         except Exception as e:

marimo/_server/rtc/doc.py

@@ -7,6 +7,7 @@
 from marimo import _loggers
 from marimo._server.workspace import FileKey
 from marimo._types.ids import CellId_t
+from marimo._utils.asyncio_utils import supervised_task
 
 if TYPE_CHECKING:
     from loro import LoroDoc
@@ -143,8 +144,9 @@ async def remove_client(
 
         # Create the cleaner task outside the lock to avoid deadlocks
         if should_create_cleaner:
-            self.loro_docs_cleaners[file_key] = asyncio.create_task(
-                self._clean_loro_doc(file_key, 60.0)
+            self.loro_docs_cleaners[file_key] = supervised_task(
+                self._clean_loro_doc(file_key, 60.0),
+                name=f"rtc.cleaner.{file_key}",
             )
 
     async def _do_remove_doc(self, file_key: FileKey) -> None: