refactor: SessionConsumer as a SessionExtension #7570

mscolnick · 2025-12-19T21:31:25Z

This is further refactoring to make SessionConsumer a type of extension (plus some).

vercel · 2025-12-19T21:31:28Z

The latest updates on your projects. Learn more about Vercel for GitHub.

Project	Deployment	Review	Updated (UTC)
marimo-docs	Ready	Preview, Comment	Dec 22, 2025 4:19pm

Copilot

Pull request overview

This PR refactors SessionConsumer to become a type of SessionExtension, enabling a more composable architecture where consumers can be treated as extensions to the session lifecycle. Key changes include renaming lifecycle methods (on_start → on_attach, on_stop → on_detach, write_operation → notify) and introducing a new NotificationListenerExtension to handle kernel message distribution.

Moved SessionConsumer from marimo._server.model to marimo._server.consumer as an abstract base class extending SessionExtension
Introduced NotificationListenerExtension to handle kernel notification distribution, extracting this logic from the session core
Updated all consumers and tests to use the new lifecycle method names and notify signature

Reviewed changes

Copilot reviewed 23 out of 23 changed files in this pull request and generated 6 comments.

Show a summary per file

File	Description
marimo/_server/consumer.py	New file defining SessionConsumer as an ABC extending SessionExtension with updated method signatures
marimo/_server/sessions/extensions/types.py	New file defining the SessionExtension protocol with on_attach/on_detach lifecycle hooks
marimo/_server/sessions/extensions/extensions.py	Moved SessionExtension definition to types.py and added NotificationListenerExtension to handle kernel message distribution
marimo/_server/sessions/session.py	Updated session to treat consumers as extensions, removed message_distributor logic (moved to extension), and renamed write_operation to notify
marimo/_server/sessions/room.py	Simplified room.close() and updated broadcast to work with KernelMessage instead of MessageOperation
marimo/_server/api/endpoints/ws.py	Updated WebSocketHandler to implement new SessionConsumer interface with on_attach/on_detach and notify methods
marimo/_server/export/init.py	Updated RunUntilCompletionSessionConsumer to implement new SessionConsumer interface
marimo/_server/sessions/utils.py	Updated to serialize operations before notifying consumers
marimo/_server/sessions/types.py	Updated Session protocol to use notify instead of write_operation, moved imports to TYPE_CHECKING
tests/_server/test_sessions.py	Updated all test assertions to use new method names (on_attach, on_detach, notify)
tests/_server/test_session_manager.py	Updated imports to use SessionConsumer from new location
tests/_server/sessions/test_file_change_handler.py	Updated mock session to use notify instead of write_operation
marimo/_server/start.py	Updated to use session.notify instead of session.write_operation
marimo/_server/sessions/file_change_handler.py	Updated to use session.notify and moved Session import to TYPE_CHECKING
marimo/_server/sessions/session_manager.py	Updated to use session.notify and moved SessionConsumer import
marimo/_server/sessions/events.py	Added on_notification_sent event and emit_notification_sent method
marimo/_server/model.py	Removed SessionConsumer protocol (moved to consumer.py)
marimo/_server/api/endpoints/execution.py	Updated to use session.notify
marimo/_server/api/endpoints/editing.py	Updated to use session.notify
marimo/_server/api/deps.py	Moved Session import to TYPE_CHECKING
marimo/_ai/_tools/tools/datasource.py	Moved Session import to TYPE_CHECKING
marimo/_ai/_tools/base.py	Moved Session and SessionManager imports to TYPE_CHECKING
marimo/_server/sessions/extensions/init.py	New file with copyright header

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot · 2025-12-19T21:50:40Z

marimo/_server/api/endpoints/ws.py



-class WebsocketHandler(SessionConsumer):
+class WebsocketHandler(SessionConsumer, SessionExtension):


WebsocketHandler inherits from both SessionConsumer and SessionExtension, but SessionConsumer already inherits from SessionExtension (as defined in marimo/_server/consumer.py). The explicit SessionExtension inheritance is redundant and can be removed.

Copilot · 2025-12-19T21:50:40Z

marimo/_server/consumer.py

+    """
+    Consumer for a session. This extends the SessionExtension interface.
+
+    This allows use to communicate with a session via different


Typo in docstring: "use to" should be "us to"

Suggested change

This allows use to communicate with a session via different

This allows us to communicate with a session via different

Copilot · 2025-12-19T21:50:41Z

marimo/_server/sessions/room.py

-                consumer.write_operation(operation)
+                consumer.notify(notification)

    def close(self) -> None:


The close method no longer calls on_detach for consumers, making it inconsistent with remove_consumer which does call on_detach. While this works correctly in practice because session.close() calls _detach_extensions() first, the room.close() method is now asymmetric with room.add_consumer/remove_consumer, which could be confusing for maintainability.

Suggested change

def close(self) -> None:

def close(self) -> None:

# Detach all consumers before clearing the room to keep behavior

# consistent with remove_consumer.

for consumer in list(self.consumers.keys()):

consumer.on_detach()

Copilot · 2025-12-19T21:50:41Z

marimo/_server/sessions/session.py

+        # Consumers are also extensions, so we want to attach them to the session
+        self.extensions.append(session_consumer)
+        session_consumer.on_attach(self, self._event_bus)


The on_attach method is called twice for the initial session consumer during initialization. The flow is: init calls connect_consumer (line 169) which adds the consumer to extensions and explicitly calls on_attach (lines 306-307), then _attach_extensions is called (line 173) which calls on_attach again for all extensions including the newly added consumer. This double attachment could cause issues such as duplicate event subscriptions or resource initialization.

Suggested change

# Consumers are also extensions, so we want to attach them to the session

self.extensions.append(session_consumer)

session_consumer.on_attach(self, self._event_bus)

# Consumers are also extensions; register them so the extension

# attachment mechanism can attach them appropriately.

self.extensions.append(session_consumer)

Copilot · 2025-12-19T21:50:41Z

marimo/_server/sessions/session.py

+        # HACK: Ideally we don't need to reach into this extension directly
+        for extension in self.extensions:
+            if isinstance(extension, NotificationListenerExtension):
+                extension.distributor.flush()


The flush_messages method accesses extension.distributor without checking if the attribute exists. If on_attach fails or hasn't been called yet, this will raise an AttributeError. Consider adding a check or using getattr with a default value.

Suggested change

extension.distributor.flush()

distributor = getattr(extension, "distributor", None)

if distributor is not None:

distributor.flush()

marimo/_server/export/__init__.py

Light2Dark · 2025-12-22T14:42:16Z

marimo/_server/api/endpoints/ws.py

+    def on_attach(self, session: Session, event_bus: SessionEventBus) -> None:
+        del session
+        del event_bus
+        return None


Suggested change

return None

probably something that should be caught by linters instead of humans

https://docs.astral.sh/ruff/rules/useless-return/, can add this, but need to clean up some old code

Light2Dark · 2025-12-22T14:44:55Z

marimo/_server/export/__init__.py

+        ) -> None:
+            del session
+            del event_bus
+            return None


Suggested change

return None

marimo/_server/sessions/session_manager.py

Light2Dark · 2025-12-22T15:07:59Z

marimo/_server/consumer.py

+    def notify(self, notification: KernelMessage) -> None: ...
+
+    @abstractmethod
+    def connection_state(self) -> ConnectionState: ...


nit

Suggested change

def connection_state(self) -> ConnectionState: ...

def get_connection_state(self) -> ConnectionState: ...

i can do this in a followup for renaming

manzt · 2025-12-22T16:33:40Z

marimo/_server/api/endpoints/files.py

-    session = app_state.require_current_session()
-    prev_path = session.app_file_manager.path
-
-    session.app_file_manager.rename(body.filename)
-    new_path = session.app_file_manager.path
-
-    if prev_path and new_path:
-        app_state.session_manager.recents.rename(prev_path, new_path)
-    elif new_path:
-        app_state.session_manager.recents.touch(new_path)


nice to see these conditionals moved

manzt · 2025-12-22T16:35:12Z

marimo/_server/sessions/session.py

        ttl_seconds: Optional[int],
+        extensions: list[SessionExtension] | None = None,


tangent: are we fine to use | syntax now for Optional? Would be nice to be consistent with that across the codebase.

ideally we can move it all to | at some point. (easy by flipping on with a ruff rule). some type inference stuff like our custom parse_raw may require Optional in the dataclasses

refactor: SessionConsumer as a SessionExtension

affd620

vercel bot deployed to Preview December 19, 2025 21:32 View deployment

mscolnick requested review from Copilot, dmadisetti and manzt December 19, 2025 21:38

unsub

de110e5

vercel bot deployed to Preview December 19, 2025 21:40 View deployment

Copilot started reviewing on behalf of mscolnick December 19, 2025 21:43 View session

start

d1c528b

vercel bot deployed to Preview December 19, 2025 21:45 View deployment

Copilot AI reviewed Dec 19, 2025

View reviewed changes

vercel bot deployed to Preview December 19, 2025 22:04 View deployment

mscolnick added 2 commits December 19, 2025 23:03

file watcher

3397859

fixes

146b3c4

mscolnick force-pushed the ms/session-consumer-as-extension branch from 7b1f952 to 146b3c4 Compare December 20, 2025 04:29

vercel bot deployed to Preview December 20, 2025 04:29 View deployment

fixes

055b49f

vercel bot deployed to Preview December 20, 2025 15:50 View deployment

fix

c24f7ed

vercel bot deployed to Preview December 21, 2025 20:05 View deployment

delete unused

f97f333

vercel bot deployed to Preview December 22, 2025 14:51 View deployment

Light2Dark reviewed Dec 22, 2025

View reviewed changes

Light2Dark previously approved these changes Dec 22, 2025

View reviewed changes

fixes

ab1245c

mscolnick dismissed Light2Dark’s stale review via ab1245c December 22, 2025 16:18

vercel bot deployed to Preview December 22, 2025 16:19 View deployment

Light2Dark approved these changes Dec 22, 2025

View reviewed changes

mscolnick merged commit 5b14c18 into main Dec 22, 2025
38 of 41 checks passed

mscolnick deleted the ms/session-consumer-as-extension branch December 22, 2025 16:38

manzt reviewed Dec 22, 2025

View reviewed changes

dmadisetti mentioned this pull request Jan 7, 2026

fix: relative path fixes for home mode #7722

Merged

dmadisetti added the internal A refactor or improvement that is not user facing label Jan 8, 2026



		class WebsocketHandler(SessionConsumer):
		class WebsocketHandler(SessionConsumer, SessionExtension):

	This allows use to communicate with a session via different
	This allows us to communicate with a session via different

-    def close(self) -> None:
+    def close(self) -> None:
+        # Detach all consumers before clearing the room to keep behavior
+        # consistent with remove_consumer.
+        for consumer in list(self.consumers.keys()):
+            consumer.on_detach()

	def connection_state(self) -> ConnectionState: ...
	def get_connection_state(self) -> ConnectionState: ...

		ttl_seconds: Optional[int],
		extensions: list[SessionExtension] \| None = None,

refactor: SessionConsumer as a SessionExtension #7570

refactor: SessionConsumer as a SessionExtension #7570

Uh oh!

Conversation

mscolnick commented Dec 19, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

vercel bot commented Dec 19, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Copilot AI Dec 19, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Dec 19, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Dec 19, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Dec 19, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Dec 19, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mscolnick Dec 22, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

mscolnick commented Dec 19, 2025 •

edited

Loading

vercel bot commented Dec 19, 2025 •

edited

Loading

mscolnick Dec 22, 2025 •

edited

Loading