refactor(runtime): split kernel command dispatch into router + callback bundles by mscolnick · Pull Request #9577 · marimo-team/marimo

mscolnick · 2026-05-18T13:18:21Z

Extract the in-line request_handler cached_property and the six *Callbacks
classes out of runtime.py (~1270 LOC removed) into a small RequestRouter,
a KernelRequestHandlers bundle for handlers that need kernel-side context
wrappers, and per-domain callback modules under _runtime/callbacks/ that
satisfy a KernelCallback protocol and self-register their command bindings.

…ck bundles Extract the in-line request_handler cached_property and the six *Callbacks classes out of runtime.py (~1270 LOC removed) into a small RequestRouter, a KernelRequestHandlers bundle for handlers that need kernel-side context wrappers, and per-domain callback modules under _runtime/callbacks/ that satisfy a KernelCallback protocol and self-register their command bindings.

vercel · 2026-05-18T13:18:25Z

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

Project	Deployment	Actions	Updated (UTC)
marimo-docs	Ready	Preview, Comment	May 18, 2026 7:09pm

Copilot

Pull request overview

Pure refactor of marimo/_runtime/runtime.py: the inline request_handler cached_property and the six per-domain *Callbacks classes are extracted out of the ~3000-LOC kernel module into a small RequestRouter, a KernelRequestHandlers bundle for handlers that need kernel-side context wrappers, and one callback module per domain under marimo/_runtime/callbacks/. Each callback bundle implements a runtime-checkable KernelCallback protocol and self-registers its command bindings on the router. Behaviour and command bindings are preserved; tests are reorganized to match the new layout and a new test_request_router.py covers the router and registration contracts.

Changes:

Add RequestRouter (request_router.py) and KernelRequestHandlers (kernel_request_handlers.py); Kernel.__init__ now constructs the router and lets each callback register(router).
Move DatasetCallbacks, ExternalStorageCallbacks, SqlCallbacks, SecretsCallbacks, PackagesCallbacks, CacheCallbacks into marimo/_runtime/callbacks/* behind a shared KernelCallback protocol; trim imports in runtime.py accordingly.
Update tests: replace the old request_handler-caching test with router dispatch-table checks, add tests/_runtime/test_request_router.py, and re-target broadcast_notification patch paths to marimo._runtime.callbacks.packages.

Reviewed changes

Copilot reviewed 15 out of 15 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
marimo/_runtime/runtime.py	Removes inline callback classes and `request_handler`; wires `RequestRouter` + `KernelRequestHandlers` + callback bundles in `__init__`; `handle_message` now calls `router.dispatch`.
marimo/_runtime/request_router.py	New small dispatcher mapping command type → async handler.
marimo/_runtime/kernel_request_handlers.py	New bundle of kernel-owned handlers (instantiate/run/sync/scratchpad/UI/model/etc.) that wrap kernel methods with request-context and completion notifications.
marimo/_runtime/callbacks/init.py	Public re-exports for all callback classes and `KernelCallback`.
marimo/_runtime/callbacks/protocol.py	New runtime-checkable `KernelCallback` protocol with a `register(router)` method.
marimo/_runtime/callbacks/datasets.py	Moved `DatasetCallbacks`; self-registers dataset/SQL preview commands.
marimo/_runtime/callbacks/external_storage.py	Moved `ExternalStorageCallbacks`; self-registers storage list/download.
marimo/_runtime/callbacks/sql.py	Moved `SqlCallbacks`; self-registers `ValidateSQLCommand`.
marimo/_runtime/callbacks/secrets.py	Moved `SecretsCallbacks`; self-registers secret list/refresh.
marimo/_runtime/callbacks/packages.py	Moved `PackagesCallbacks`; self-registers install command, wraps with `CompletedRunNotification`.
marimo/_runtime/callbacks/cache.py	Moved `CacheCallbacks`; self-registers clear/info cache commands.
tests/_runtime/test_request_router.py	New tests for router dispatch and per-callback registration contracts.
tests/_runtime/test_runtime.py	Replaces `request_handler` caching test with router dispatch-table size check.
tests/_runtime/test_runtime_secrets.py	Updates `SecretsCallbacks` import to new module path.
tests/_runtime/test_manage_script_metadata.py	Re-targets `broadcast_notification` patch paths to `marimo._runtime.callbacks.packages`.

cubic-dev-ai

3 issues found across 15 files

Prompt for AI agents (unresolved issues)


Check if these issues are valid — if so, understand the root cause of each and fix them. If appropriate, use sub-agents to investigate and fix each issue separately.


<file name="marimo/_runtime/kernel_request_handlers.py">

<violation number="1" location="marimo/_runtime/kernel_request_handlers.py:150">
P1: This truthiness check drops valid model updates when `state` is empty (`{}`). Check for `None` instead of truthiness so falsy-but-valid states are processed.</violation>
</file>

<file name="marimo/_runtime/callbacks/datasets.py">

<violation number="1" location="marimo/_runtime/callbacks/datasets.py:170">
P2: `preview_datasource_connection` incorrectly requires catalog support, so query-only SQL engines are rejected and no connection notification is sent.</violation>
</file>

<file name="marimo/_runtime/callbacks/cache.py">

<violation number="1" location="marimo/_runtime/callbacks/cache.py:36">
P2: `clear_cache` only clears `BasePersistenceLoader` instances, so in-memory cache contexts are not purged.</violation>
</file>

Architecture diagram

sequenceDiagram
    participant Client as Frontend Client
    participant API as API Handler
    participant Router as RequestRouter
    participant Krh as KernelRequestHandlers
    participant CaCb as CacheCallbacks
    participant DsCb as DatasetCallbacks
    participant EsCb as ExternalStorageCallbacks
    participant PkCb as PackagesCallbacks
    participant ScCb as SecretsCallbacks
    participant SqCb as SqlCallbacks
    participant Kernel as Kernel

    Note over Client,Kernel: Command dispatch flow (all 27 commands)

    Client->>API: Send command message
    API->>Router: dispatch(command)
    
    alt Kernel-owned commands (e.g., ExecuteCellsCommand, UpdateUIElementCommand, etc.)
        Router->>Krh: registered handler
        Note over Krh: Wraps with http_request_context<br/>and broadcast_notification
        Krh->>Kernel: delegate to kernel method
        Kernel-->>Krh: result
        Krh-->>API: CompletedRunNotification
    else Cache commands (ClearCacheCommand, GetCacheInfoCommand)
        Router->>CaCb: registered handler
        CaCb->>Kernel: access kernel.globals
        CaCb->>CaCb: clear or gather cache info
        CaCb-->>API: CacheClearedNotification / CacheInfoNotification
    else Dataset commands (PreviewDatasetColumnCommand, etc.)
        Router->>DsCb: registered handler
        DsCb->>Kernel: get_sql_connection()
        DsCb->>DsCb: get_column_preview / engine catalog ops
        alt Success
            DsCb-->>API: DataColumnPreviewNotification / SQLTablePreviewNotification
        else Error
            DsCb-->>API: Error notification
        end
    else External storage commands (StorageListEntriesCommand, StorageDownloadCommand)
        Router->>EsCb: registered handler
        EsCb->>Kernel: access kernel.globals
        EsCb->>EsCb: _get_storage_backend()
        EsCb->>EsCb: list_entries / download / sign URL
        EsCb-->>API: StorageEntriesNotification / StorageDownloadReadyNotification
    else Package commands (InstallPackagesCommand)
        Router->>PkCb: registered handler
        PkCb->>Kernel: module_registry, enqueue_control_request
        PkCb->>PkCb: install_missing_packages()
        PkCb-->>API: InstallingPackageAlertNotification / CompletedRunNotification
    else Secrets commands (ListSecretKeysCommand, RefreshSecretsCommand)
        Router->>ScCb: registered handler
        ScCb->>Kernel: user_config, _original_environ
        ScCb->>ScCb: get_secret_keys / load_dotenv
        ScCb-->>API: SecretKeysResultNotification
    else SQL commands (ValidateSQLCommand)
        Router->>SqCb: registered handler
        SqCb->>Kernel: get_sql_connection()
        SqCb->>SqCb: parse_sql / engine.execute_in_explain_mode
        alt Success
            SqCb-->>API: ValidateSQLResultNotification
        else Error
            SqCb-->>API: Error notification
        end
    else Unknown command
        Router--xAPI: Raises ValueError
    end

    Note over Router: All callbacks self-register via<br/>KernelCallback protocol in Kernel.__init__

_{Reply with feedback, questions, or to request a fix.

Re-trigger cubic}

cubic-dev-ai · 2026-05-18T13:23:19Z

+        # caused Model+UI interleaving that the batch merger
+        # couldn't collapse (different types), leading to every
+        # drag tick getting its own full cell re-execution.
+        if ui_element_id and state:


P1: This truthiness check drops valid model updates when state is empty ({}). Check for None instead of truthiness so falsy-but-valid states are processed.

Prompt for AI agents

Check if this issue is valid — if so, understand the root cause and fix it. At marimo/_runtime/kernel_request_handlers.py, line 150: <comment>This truthiness check drops valid model updates when `state` is empty (`{}`). Check for `None` instead of truthiness so falsy-but-valid states are processed.</comment> <file context> @@ -0,0 +1,185 @@ + # caused Model+UI interleaving that the batch merger + # couldn't collapse (different types), leading to every + # drag tick getting its own full cell re-execution. + if ui_element_id and state: + await self._kernel.set_ui_element_value( + UpdateUIElementCommand.from_ids_and_values( </file context>

Suggested change

if ui_element_id and state:

if ui_element_id is not None and state is not None:

this is what the previous code was, so will handle in a follow-up to keep this to just a refactor

Understood—keeping this PR refactor-only and deferring behavior changes makes sense. Thanks for the context.

cubic-dev-ai · 2026-05-18T13:23:19Z

+            schema=schema_name,
+        )
+
+        engine, error = self.get_engine_catalog(variable_name)


P2: preview_datasource_connection incorrectly requires catalog support, so query-only SQL engines are rejected and no connection notification is sent.

Prompt for AI agents

Check if this issue is valid — if so, understand the root cause and fix it. At marimo/_runtime/callbacks/datasets.py, line 170: <comment>`preview_datasource_connection` incorrectly requires catalog support, so query-only SQL engines are rejected and no connection notification is sent.</comment> <file context> @@ -0,0 +1,348 @@ + schema=schema_name, + ) + + engine, error = self.get_engine_catalog(variable_name) + if error is not None or engine is None: + broadcast_notification( </file context>

this is what the previous code was, so will handle in a follow-up to keep this to just a refactor

Got it—thanks for clarifying scope; out of scope for this refactor and fine to handle in a follow-up.

cubic-dev-ai · 2026-05-18T13:23:19Z

+        saved = 0
+        for obj in ctx.globals.values():
+            if isinstance(obj, CacheContext):
+                if isinstance(obj.loader, BasePersistenceLoader):


P2: clear_cache only clears BasePersistenceLoader instances, so in-memory cache contexts are not purged.

Prompt for AI agents

Check if this issue is valid — if so, understand the root cause and fix it. At marimo/_runtime/callbacks/cache.py, line 36: <comment>`clear_cache` only clears `BasePersistenceLoader` instances, so in-memory cache contexts are not purged.</comment> <file context> @@ -0,0 +1,67 @@ + saved = 0 + for obj in ctx.globals.values(): + if isinstance(obj, CacheContext): + if isinstance(obj.loader, BasePersistenceLoader): + obj.loader.clear() + </file context>

this is what the previous code was, so will handle in a follow-up to keep this to just a refactor

Got it—keeping this PR scoped to a refactor and handling this in a follow-up.

The explicit-set version already lives in test_request_router.py against the same kernel fixture; the count-based duplicate added nothing.

Replace the three hand-maintained tables (per-callback bindings, kernel-owned set, and the 27-command full dispatch table) with checks against `typing.get_args(CommandMessage)` and a single `NOT_ROUTED` allowlist. Tests now self-adjust when commands are added, removed, or moved between callbacks, and report missing/extra/conflicting handlers by name instead of by count.

github-actions · 2026-05-18T19:35:18Z

🚀 Development release published. You may be able to view the changes at https://marimo.app?v=0.23.7-dev44

Resolves conflict with PR marimo-team#9577 (refactor: split kernel command dispatch into router + callback bundles). The request_handler cached_property was extracted from runtime.py into kernel_request_handlers.py — wire the UpdateQueryParamsCommand through the new KernelRequestHandlers.register path instead of the old in-line handler.register block.

Copilot AI review requested due to automatic review settings May 18, 2026 13:18

Copilot started reviewing on behalf of mscolnick May 18, 2026 13:19 View session

vercel Bot deployed to Preview May 18, 2026 13:19 View deployment

Copilot AI reviewed May 18, 2026

View reviewed changes

Comment thread tests/_runtime/test_runtime.py Outdated

cubic-dev-ai Bot reviewed May 18, 2026

View reviewed changes

test: drop redundant router dispatch-table test

379ac74

The explicit-set version already lives in test_request_router.py against the same kernel fixture; the count-based duplicate added nothing.

mscolnick requested a review from manzt May 18, 2026 13:39

mscolnick added the enhancement New feature or request label May 18, 2026

vercel Bot deployed to Preview May 18, 2026 13:41 View deployment

vercel Bot deployed to Preview May 18, 2026 19:09 View deployment

manzt approved these changes May 18, 2026

View reviewed changes

mscolnick merged commit c09cb3d into main May 18, 2026
44 checks passed

mscolnick deleted the ms/feature/kernel-callbacks-router branch May 18, 2026 19:29

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

refactor(runtime): split kernel command dispatch into router + callback bundles#9577

refactor(runtime): split kernel command dispatch into router + callback bundles#9577
mscolnick merged 3 commits into
mainfrom
ms/feature/kernel-callbacks-router

mscolnick commented May 18, 2026

vercel Bot commented May 18, 2026 •

edited

Loading

Copilot AI left a comment

Uh oh!

cubic-dev-ai Bot left a comment

cubic-dev-ai Bot May 18, 2026

mscolnick May 18, 2026

cubic-dev-ai Bot May 18, 2026

cubic-dev-ai Bot May 18, 2026

mscolnick May 18, 2026

cubic-dev-ai Bot May 18, 2026

cubic-dev-ai Bot May 18, 2026

mscolnick May 18, 2026

cubic-dev-ai Bot May 18, 2026

Uh oh!

github-actions Bot commented May 18, 2026

Labels

3 participants

	if ui_element_id and state:
	if ui_element_id is not None and state is not None:

Uh oh!

Conversation

mscolnick commented May 18, 2026

vercel Bot commented May 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

cubic-dev-ai Bot left a comment

Choose a reason for hiding this comment

cubic-dev-ai Bot May 18, 2026

Choose a reason for hiding this comment

mscolnick May 18, 2026

Choose a reason for hiding this comment

cubic-dev-ai Bot May 18, 2026

Choose a reason for hiding this comment

cubic-dev-ai Bot May 18, 2026

Choose a reason for hiding this comment

mscolnick May 18, 2026

Choose a reason for hiding this comment

cubic-dev-ai Bot May 18, 2026

Choose a reason for hiding this comment

cubic-dev-ai Bot May 18, 2026

Choose a reason for hiding this comment

mscolnick May 18, 2026

Choose a reason for hiding this comment

cubic-dev-ai Bot May 18, 2026

Choose a reason for hiding this comment

Uh oh!

github-actions Bot commented May 18, 2026

Labels

3 participants

vercel Bot commented May 18, 2026 •

edited

Loading