Add marimo.restartKernel command

Adds the ability to restart an active marimo kernel from both the UI and
quick action.

Author: manzt

extension/package.json

@@ -156,6 +156,12 @@
         "shortTitle": "On cell change: Lazy",
         "category": "marimo",
         "icon": "$(symbol-interface)"
+      },
+      {
+        "command": "marimo.restartKernel",
+        "title": "Restart Notebook Kernel",
+        "shortTitle": "Restart Kernel",
+        "category": "marimo"
       }
     ],
     "viewsContainers": {
@@ -216,15 +222,20 @@
           "when": "notebookType == 'marimo-notebook' && marimo.notebook.hasStaleCells",
           "group": "navigation@1"
         },
+        {
+          "command": "marimo.restartKernel",
+          "when": "notebookType == 'marimo-notebook'",
+          "group": "navigation@2"
+        },
         {
           "command": "marimo.toggleOnCellChangeAutoRun",
           "when": "notebookType == 'marimo-notebook' && marimo.config.runtime.on_cell_change != 'lazy'",
-          "group": "navigation@2"
+          "group": "navigation@3"
         },
         {
           "command": "marimo.toggleOnCellChangeLazy",
           "when": "notebookType == 'marimo-notebook' && marimo.config.runtime.on_cell_change == 'lazy'",
-          "group": "navigation@2"
+          "group": "navigation@3"
         }
       ]
     }

extension/src/constants.ts

@@ -7,6 +7,7 @@ export type MarimoCommand =
   | "marimo.publishMarimoNotebook"
   | "marimo.publishMarimoNotebookGist"
   | "marimo.refreshPackages"
+  | "marimo.restartKernel"
   | "marimo.runStale"
   | "marimo.showMarimoMenu"
   | "marimo.toggleOnCellChangeAutoRun"
@@ -21,5 +22,5 @@ export type MarimoView =
 export const NOTEBOOK_TYPE = "marimo-notebook";
 
 export type MarimoContextKey =
-  | "marimo.notebook.hasStaleCells"
-  | "marimo.config.runtime.on_cell_change";
+  | "marimo.config.runtime.on_cell_change"
+  | "marimo.notebook.hasStaleCells";

extension/src/layers/RegisterCommands.ts

@@ -3,7 +3,9 @@ import { Cause, Chunk, Effect, Either, Layer, Option } from "effect";
 import { decodeCellMetadata, isStaleCellMetadata } from "../schemas.ts";
 import { ConfigContextManager } from "../services/config/ConfigContextManager.ts";
 import { MarimoConfigurationService } from "../services/config/MarimoConfigurationService.ts";
+import { ExecutionRegistry } from "../services/ExecutionRegistry.ts";
 import { GitHubClient } from "../services/GitHubClient.ts";
+import { LanguageClient } from "../services/LanguageClient.ts";
 import { NotebookSerializer } from "../services/NotebookSerializer.ts";
 import { OutputChannel } from "../services/OutputChannel.ts";
 import { VsCode } from "../services/VsCode.ts";
@@ -17,10 +19,12 @@ export const RegisterCommandsLive = Layer.scopedDiscard(
   Effect.gen(function* () {
     const gh = yield* GitHubClient;
     const code = yield* VsCode;
+    const client = yield* LanguageClient;
     const channel = yield* OutputChannel;
     const serializer = yield* NotebookSerializer;
     const configService = yield* MarimoConfigurationService;
     const configContextManager = yield* ConfigContextManager;
+    const executions = yield* ExecutionRegistry;
 
     yield* code.commands.registerCommand(
       "marimo.newMarimoNotebook",
@@ -78,6 +82,61 @@ export const RegisterCommandsLive = Layer.scopedDiscard(
         configContextManager,
       }),
     );
+
+    yield* code.commands.registerCommand(
+      "marimo.restartKernel",
+      Effect.gen(function* () {
+        const editor = yield* code.window.getActiveNotebookEditor();
+        if (Option.isNone(editor)) {
+          yield* code.window.showInformationMessage(
+            "No marimo notebook is currently open",
+          );
+          return;
+        }
+
+        yield* code.window.withProgress(
+          {
+            location: code.ProgressLocation.Window,
+            title: "Restarting kernel",
+            cancellable: true,
+          },
+          Effect.fnUntraced(function* (progress) {
+            progress.report({ message: "Closing session..." });
+
+            const result = yield* client
+              .executeCommand({
+                command: "marimo.api",
+                params: {
+                  method: "close_session",
+                  params: {
+                    notebookUri: getNotebookUri(editor.value.notebook),
+                    inner: {},
+                  },
+                },
+              })
+              .pipe(Effect.either);
+
+            if (Either.isLeft(result)) {
+              yield* Effect.logFatal("Failed to restart kernel", result.left);
+              yield* showErrorAndPromptLogs("Failed to restart kernel.", {
+                channel,
+                code,
+              });
+              return;
+            }
+
+            yield* executions.handleInterrupted(editor.value);
+
+            progress.report({ message: "Kernel restarted." });
+            yield* Effect.sleep("500 millis");
+          }),
+        );
+
+        yield* code.window.showInformationMessage(
+          "Kernel restarted successfully",
+        );
+      }),
+    );
   }),
 );
 

extension/src/types.ts

@@ -64,6 +64,9 @@ type ListPackagesRequest = {};
 type DependencyTreeRequest = {};
 // biome-ignore lint/complexity/noBannedTypes: We need this for over the wire
 type GetConfigurationRequest = {};
+// biome-ignore lint/complexity/noBannedTypes: We need this for over the wire
+type CloseSessionRequest = {};
+
 interface UpdateConfigurationRequest {
   config: Record<string, unknown>;
 }
@@ -81,6 +84,7 @@ type MarimoApiMethodMap = {
   get_dependency_tree: SessionScoped<DependencyTreeRequest>;
   get_configuration: NotebookScoped<GetConfigurationRequest>;
   update_configuration: NotebookScoped<UpdateConfigurationRequest>;
+  close_session: NotebookScoped<CloseSessionRequest>;
 };
 
 type ApiRequest<K extends keyof MarimoApiMethodMap> = {

src/marimo_lsp/api.py

@@ -16,6 +16,7 @@
 from marimo_lsp.debug_adapter import handle_debug_adapter_request
 from marimo_lsp.loggers import get_logger
 from marimo_lsp.models import (
+    CloseSessionRequest,
     DebugAdapterRequest,
     DependencyTreeRequest,
     DeserializeRequest,
@@ -111,6 +112,14 @@ async def interrupt(
         logger.warning(f"No session found for {args.notebook_uri}")
 
 
+async def close_session(
+    manager: LspSessionManager,
+    args: NotebookCommand[CloseSessionRequest],
+):
+    logger.info(f"close_session for {args.notebook_uri}")
+    manager.close_session(args.notebook_uri)
+
+
 async def get_package_list(
     manager: LspSessionManager,
     args: SessionCommand[ListPackagesRequest],
@@ -258,6 +267,12 @@ async def handle_api_command(  # noqa: C901, PLR0911
             msgspec.convert(params, type=NotebookCommand[DebugAdapterRequest]),
         )
 
+    if method == "close_session":
+        return await close_session(
+            manager,
+            msgspec.convert(params, type=NotebookCommand[CloseSessionRequest]),
+        )
+
     if method == "serialize":
         return await serialize(msgspec.convert(params, type=SerializeRequest))
 

src/marimo_lsp/models.py

@@ -88,6 +88,10 @@ class GetConfigurationRequest(msgspec.Struct, rename="camel"):
     """A request to get the current configuration."""
 
 
+class CloseSessionRequest(msgspec.Struct, rename="camel"):
+    """A request to close the current session."""
+
+
 class UpdateConfigurationRequest(msgspec.Struct, rename="camel"):
     """A request to update the user configuration."""