Add Debug Adapter Protocol (DAP) support to marimo LSP

Enables debugging marimo notebooks in VS Code by forwarding DAP messages
between the VS Code debugger and the LSP server. Just a stub now to
wire up everything.

Author: manzt

extension/package.json

@@ -20,6 +20,35 @@
       "title": "Marimo",
       "properties": {}
     },
+    "debuggers": [
+      {
+        "type": "marimo",
+        "label": "Marimo Debugger",
+        "languages": ["python"],
+        "configurationAttributes": {
+          "launch": {
+            "properties": {
+              "notebookUri": {
+                "type": "string",
+                "description": "URI of the notebook to debug"
+              }
+            }
+          }
+        },
+        "initialConfigurations": [
+          {
+            "type": "marimo",
+            "request": "launch",
+            "name": "Debug Marimo Notebook"
+          }
+        ]
+      }
+    ],
+    "breakpoints": [
+      {
+        "language": "python"
+      }
+    ],
     "notebooks": [
       {
         "type": "marimo-notebook",

extension/src/debugAdapter.ts

@@ -0,0 +1,105 @@
+import * as vscode from "vscode";
+import * as lsp from "vscode-languageclient";
+import type { BaseLanguageClient } from "vscode-languageclient";
+
+import { notebookType } from "./types.ts";
+import { Logger } from "./logging.ts";
+import { executeCommand } from "./commands.ts";
+
+export function debugAdapter(
+  client: BaseLanguageClient,
+  options: { signal: AbortSignal },
+) {
+  Logger.info("Debug.Init", "Registering debug adapter");
+
+  const disposeFactory = vscode.debug.registerDebugAdapterDescriptorFactory(
+    "marimo",
+    {
+      createDebugAdapterDescriptor: createDebugAdapterDescriptor.bind(
+        null,
+        client,
+      ),
+    },
+  );
+
+  const disposeProvider = vscode.debug.registerDebugConfigurationProvider(
+    "marimo",
+    {
+      resolveDebugConfiguration(_folder, config) {
+        Logger.info("Debug.Config", "Resolving debug configuration", {
+          config,
+        });
+
+        const notebook = vscode.window.activeNotebookEditor?.notebook;
+        if (!notebook || notebook.notebookType !== notebookType) {
+          Logger.warn("Debug.Config", "No active marimo notebook found");
+          return undefined;
+        }
+        config.type = "marimo";
+        config.name = config.name ?? "Debug Marimo";
+        config.request = config.request ?? "launch";
+        config.notebookUri = notebook.uri.toString();
+
+        Logger.info("Debug.Config", "Configuration resolved", {
+          notebookUri: config.notebookUri,
+          type: config.type,
+          request: config.request,
+        });
+        return config;
+      },
+    },
+  );
+
+  options.signal.addEventListener("abort", () => {
+    Logger.info("Debug.Cleanup", "Disposing debug adapter");
+    disposeFactory.dispose();
+    disposeProvider.dispose();
+  });
+}
+
+function createDebugAdapterDescriptor(
+  client: lsp.BaseLanguageClient,
+  session: vscode.DebugSession,
+): vscode.DebugAdapterDescriptor {
+  Logger.info("Debug.Factory", "Creating debug adapter", {
+    sessionId: session.id,
+    name: session.name,
+    type: session.type,
+    configuration: session.configuration,
+  });
+
+  const sendMessage = new vscode.EventEmitter<vscode.DebugProtocolMessage>();
+  const disposer = client.onNotification(
+    "marimo/dap",
+    ({ sessionId, message }) => {
+      Logger.debug("Debug.Receive", "Received DAP response from LSP", {
+        sessionId,
+        message,
+      });
+      if (sessionId === session.id) {
+        sendMessage.fire(message);
+      }
+    },
+  );
+
+  return new vscode.DebugAdapterInlineImplementation({
+    onDidSendMessage: sendMessage.event,
+    handleMessage(message) {
+      Logger.debug("Debug.Send", "Sending DAP message to LSP", {
+        sessionId: session.id,
+        message,
+      });
+      executeCommand(client, {
+        command: "marimo.dap",
+        params: {
+          sessionId: session.id,
+          notebookUri: session.configuration.notebookUri,
+          message,
+        },
+      });
+    },
+    dispose() {
+      disposer.dispose();
+    },
+  });
+}

extension/src/extension.ts

@@ -1,6 +1,7 @@
 import * as vscode from "vscode";
 import * as lsp from "vscode-languageclient/node";
 
+import { debugAdapter } from "./debugAdapter.ts";
 import { kernelManager } from "./kernelManager.ts";
 import { languageClient } from "./languageClient.ts";
 import { notebookSerializer } from "./notebookSerializer.ts";
@@ -17,6 +18,7 @@ export async function activate(context: vscode.ExtensionContext) {
 
   const controller = new AbortController();
   const client = languageClient({ signal: controller.signal });
+  debugAdapter(client, { signal: controller.signal });
   kernelManager(client, { signal: controller.signal });
   notebookSerializer(client, { signal: controller.signal });
 

extension/src/types.ts

@@ -1,3 +1,4 @@
+import * as vscode from "vscode";
 import { components as Api } from "@marimo-team/openapi/src/api";
 import type { NotebookSerialization } from "./schemas.ts";
 
@@ -19,6 +20,11 @@ export type RequestMap = {
   >;
   "marimo.serialize": { notebook: NotebookSerialization };
   "marimo.deserialize": { source: string };
+  "marimo.dap": {
+    sessionId: string;
+    notebookUri: string;
+    message: vscode.DebugProtocolMessage;
+  };
 };
 
 export const notebookType = "marimo-notebook";

src/marimo_lsp/debug_adapter.py

@@ -0,0 +1,71 @@
+"""Handler for DAP messages."""
+
+from __future__ import annotations
+
+import typing
+
+import attrs
+import cattrs
+
+from marimo_lsp.loggers import get_logger
+
+if typing.TYPE_CHECKING:
+    from pygls.lsp.server import LanguageServer
+
+    from marimo_lsp.session_manager import LspSessionManager
+
+logger = get_logger()
+converter = cattrs.Converter()
+
+
+@attrs.define
+class DapRequestMessage:
+    """
+    A generic DAP (Debug Adapter Protocol) request message.
+
+    DAP requests follow a standard structure where the command field
+    determines the action, and arguments contain command-specific parameters
+    that require further parsing based on the command type.
+    """
+
+    seq: int
+    """Sequence number of the message."""
+
+    type: typing.Literal["request"]
+    """Message type - always 'request' for DAP requests."""
+
+    command: str
+    """The command to execute (e.g., 'initialize', 'launch', 'setBreakpoints')."""
+
+    arguments: dict | None
+    """Command-specific arguments. Should be parsed further in ./debug_adapter.py"""
+
+
+def handle_debug_adapter_request(
+    ls: LanguageServer,
+    manager: LspSessionManager,
+    *,
+    notebook_uri: str,
+    session_id: str,
+    message: dict,
+) -> None:
+    """Handle DAP requests."""
+    request = converter.structure(message, DapRequestMessage)
+    logger.debug(f"Debug.Send {session_id=}, {request=}")
+
+    session = manager.get_session(notebook_uri)
+    assert session, f"No session in workspace for {notebook_uri}"
+
+    ls.protocol.notify(
+        "marimo/dap",
+        {
+            "sessionId": session_id,
+            "message": {
+                "type": "response",
+                "request_seq": request.seq,
+                "success": True,
+                "command": request.command,
+                "request": {},
+            },
+        },
+    )

src/marimo_lsp/models.py

@@ -103,6 +103,20 @@ class ConvertRequest(BaseRequest):
     """The identifier for the text document to convert"""
 
 
+@attrs.define
+class DebugAdapterRequest(BaseRequest):
+    """A forwarded DAP request."""
+
+    session_id: str
+    """A UUID for the debug session."""
+
+    notebook_uri: str
+    """The URI of the notebook."""
+
+    message: dict
+    """They DAP message."""
+
+
 def _camel_to_snake(name: str) -> str:
     s1 = re.sub(r"(.)([A-Z][a-z]+)", r"\1_\2", name)
     return re.sub(r"([A-Z]+)([A-Z][a-z]*)", r"\1_\2", s1).lower()

src/marimo_lsp/server.py

@@ -22,6 +22,7 @@
 from marimo_lsp.loggers import get_logger
 from marimo_lsp.models import (
     ConvertRequest,
+    DebugAdapterRequest,
     DeserializeRequest,
     RunRequest,
     SerializeRequest,
@@ -145,6 +146,21 @@ async def deserialize(args: DeserializeRequest):
         converter = MarimoConvert.from_py(args.source)
         return dataclasses.asdict(converter.to_ir())
 
+    @server.command("marimo.dap")
+    async def dap(ls: LanguageServer, params: DebugAdapterRequest):
+        """Handle DAP messages forwarded from VS Code extension."""
+        from marimo_lsp.debug_adapter import (  # noqa: PLC0415
+            handle_debug_adapter_request,
+        )
+
+        return handle_debug_adapter_request(
+            ls=ls,
+            manager=manager,
+            session_id=params.session_id,
+            notebook_uri=params.notebook_uri,
+            message=params.message,
+        )
+
     @server.feature(
         lsp.TEXT_DOCUMENT_CODE_ACTION,
         lsp.CodeActionOptions(