Refactor session extensions to use EventAwareExtension base class (#8678)

Author: mscolnick

marimo/_internal/session/extensions.py

@@ -10,10 +10,16 @@
     ReplayExtension,
     SessionViewExtension,
 )
-from marimo._session.extensions.types import SessionExtension
+from marimo._session.extensions.types import (
+    EventAwareExtension,
+    ExtensionRegistry,
+    SessionExtension,
+)
 
 __all__ = [
     "CachingExtension",
+    "EventAwareExtension",
+    "ExtensionRegistry",
     "HeartbeatExtension",
     "LoggingExtension",
     "NotificationListenerExtension",

marimo/_server/scratchpad.py

@@ -18,7 +18,7 @@
 from marimo._messaging.notification import CellNotification
 from marimo._messaging.serde import deserialize_kernel_message
 from marimo._runtime.scratch import SCRATCH_CELL_ID
-from marimo._session.events import SessionEventBus, SessionEventListener
+from marimo._session.extensions.types import EventAwareExtension
 
 if TYPE_CHECKING:
     from collections.abc import AsyncGenerator
@@ -73,7 +73,7 @@ def _format_sse(event: str, data: Any) -> str:
 # -- Listeners ----------------------------------------------------------------
 
 
-class ScratchCellListener(SessionEventListener):
+class ScratchCellListener(EventAwareExtension):
     """Listens for scratch cell notifications via an asyncio.Queue.
 
     Supports both SSE streaming (via ``stream()``) and simple blocking
@@ -82,19 +82,10 @@ class ScratchCellListener(SessionEventListener):
     """
 
     def __init__(self) -> None:
+        super().__init__()
         self._queue: asyncio.Queue[CellNotification | None] = asyncio.Queue()
         self.timed_out = False
 
-    def on_attach(self, session: Session, event_bus: SessionEventBus) -> None:
-        del session
-        self._event_bus = event_bus
-        event_bus.subscribe(self)
-
-    def on_detach(self) -> None:
-        if hasattr(self, "_event_bus"):
-            self._event_bus.unsubscribe(self)
-            del self._event_bus
-
     def on_notification_sent(
         self, session: Session, notification: KernelMessage
     ) -> None:

marimo/_server/session/listeners.py

@@ -8,16 +8,12 @@
 from __future__ import annotations
 
 from marimo._server.recents import RecentFilesManager
-from marimo._session.events import (
-    SessionEventBus,
-    SessionEventListener,
-)
-from marimo._session.extensions.types import SessionExtension
+from marimo._session.extensions.types import EventAwareExtension
 from marimo._session.session import Session
 from marimo._types.ids import SessionId
 
 
-class RecentsTrackerListener(SessionExtension, SessionEventListener):
+class RecentsTrackerListener(EventAwareExtension):
     """Event listener that tracks recently accessed files."""
 
     def __init__(self, recents_manager: RecentFilesManager) -> None:
@@ -26,20 +22,8 @@ def __init__(self, recents_manager: RecentFilesManager) -> None:
         Args:
             recents_manager: Manager for recent files
         """
+        super().__init__()
         self._recents = recents_manager
-        self._event_bus: SessionEventBus | None = None
-
-    def on_attach(self, session: Session, event_bus: SessionEventBus) -> None:
-        """Attach the recents tracker listener to a session."""
-        del session
-        self._event_bus = event_bus
-        self._event_bus.subscribe(self)
-
-    def on_detach(self) -> None:
-        """Detach the recents tracker listener from a session."""
-        if self._event_bus:
-            self._event_bus.unsubscribe(self)
-            self._event_bus = None
 
     async def on_session_created(self, session: Session) -> None:
         """Update recent files when a session is created."""

marimo/_session/events.py

@@ -13,6 +13,8 @@
 from marimo._types.ids import ConsumerId, SessionId
 
 if TYPE_CHECKING:
+    from collections.abc import Awaitable, Callable
+
     from marimo._runtime import commands
     from marimo._session.session import Session
 
@@ -96,76 +98,84 @@ def unsubscribe(self, listener: SessionEventListener) -> None:
         if listener in self._listeners:
             self._listeners.remove(listener)
 
-    async def emit_session_created(self, session: Session) -> None:
-        """Emit a session created event."""
-        for listener in self._listeners:
+    def _emit(
+        self,
+        event_name: str,
+        call: Callable[[SessionEventListener], None],
+    ) -> None:
+        """Dispatch a synchronous event to all listeners."""
+        for listener in list(self._listeners):
             try:
-                await listener.on_session_created(session)
+                call(listener)
             except Exception as e:
                 LOGGER.error(
-                    "Error handling session created event for listener %s: %s",
+                    "Error handling %s for listener %s: %s",
+                    event_name,
                     listener,
                     e,
                 )
-                continue
 
-    async def emit_session_closed(self, session: Session) -> None:
-        """Emit a session closed event."""
-        for listener in self._listeners:
+    async def _emit_async(
+        self,
+        event_name: str,
+        call: Callable[[SessionEventListener], Awaitable[None]],
+    ) -> None:
+        """Dispatch an async event to all listeners."""
+        for listener in list(self._listeners):
             try:
-                await listener.on_session_closed(session)
+                await call(listener)
             except Exception as e:
                 LOGGER.error(
-                    "Error handling session closed event for listener %s: %s",
+                    "Error handling %s for listener %s: %s",
+                    event_name,
                     listener,
                     e,
                 )
-                continue
+
+    async def emit_session_created(self, session: Session) -> None:
+        """Emit a session created event."""
+        await self._emit_async(
+            "session_created",
+            lambda listener: listener.on_session_created(session),
+        )
+
+    async def emit_session_closed(self, session: Session) -> None:
+        """Emit a session closed event."""
+        await self._emit_async(
+            "session_closed",
+            lambda listener: listener.on_session_closed(session),
+        )
 
     async def emit_session_resumed(
         self, session: Session, old_id: SessionId
     ) -> None:
         """Emit a session resumed event."""
-        for listener in self._listeners:
-            try:
-                await listener.on_session_resumed(session, old_id)
-            except Exception as e:
-                LOGGER.error(
-                    "Error handling session resumed event for listener %s: %s",
-                    listener,
-                    e,
-                )
-                continue
+        await self._emit_async(
+            "session_resumed",
+            lambda listener: listener.on_session_resumed(session, old_id),
+        )
 
     async def emit_session_notebook_renamed(
         self, session: Session, old_path: str | None
     ) -> None:
         """Emit a session renamed event."""
-        for listener in self._listeners:
-            try:
-                await listener.on_session_notebook_renamed(session, old_path)
-            except Exception as e:
-                LOGGER.error(
-                    "Error handling session notebook renamed event for listener %s: %s",
-                    listener,
-                    e,
-                )
-                continue
+        await self._emit_async(
+            "session_notebook_renamed",
+            lambda listener: listener.on_session_notebook_renamed(
+                session, old_path
+            ),
+        )
 
     def emit_notification_sent(
         self, session: Session, notification: KernelMessage
     ) -> None:
         """Emit a notification sent event."""
-        for listener in self._listeners:
-            try:
-                listener.on_notification_sent(session, notification)
-            except Exception as e:
-                LOGGER.error(
-                    "Error handling notification sent event for listener %s: %s",
-                    listener,
-                    e,
-                )
-                continue
+        self._emit(
+            "notification_sent",
+            lambda listener: listener.on_notification_sent(
+                session, notification
+            ),
+        )
 
     def emit_received_command(
         self,
@@ -174,28 +184,16 @@ def emit_received_command(
         from_consumer_id: Optional[ConsumerId],
     ) -> None:
         """Emit a received command event."""
-        for listener in self._listeners:
-            try:
-                listener.on_received_command(
-                    session, request, from_consumer_id
-                )
-            except Exception as e:
-                LOGGER.error(
-                    "Error handling received command event for listener %s: %s",
-                    listener,
-                    e,
-                )
-                continue
+        self._emit(
+            "received_command",
+            lambda listener: listener.on_received_command(
+                session, request, from_consumer_id
+            ),
+        )
 
     def emit_received_stdin(self, session: Session, stdin: str) -> None:
         """Emit a received stdin event."""
-        for listener in self._listeners:
-            try:
-                listener.on_received_stdin(session, stdin)
-            except Exception as e:
-                LOGGER.error(
-                    "Error handling received stdin event for listener %s: %s",
-                    listener,
-                    e,
-                )
-                continue
+        self._emit(
+            "received_stdin",
+            lambda listener: listener.on_received_stdin(session, stdin),
+        )