@marimo-team/islands 0.16.5 → 0.17.1

package/src/core/codemirror/lsp/notebook-lsp.ts

@@ -1,18 +1,20 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 
+import type { EditorView } from "@codemirror/view";
 import type * as LSP from "vscode-languageserver-protocol";
+import { getNotebook } from "@/core/cells/cells";
 import type { CellId } from "@/core/cells/ids";
 import { store } from "@/core/state/jotai";
 import { invariant } from "@/utils/invariant";
 import { Logger } from "@/utils/Logger";
 import { LRUCache } from "@/utils/lru";
+import { Objects } from "@/utils/objects";
 import { topologicalCodesAtom } from "../copilot/getCodes";
 import { createNotebookLens, type NotebookLens } from "./lens";
 import {
   CellDocumentUri,
   type ILanguageServerClient,
   isClientWithNotify,
-  isClientWithPlugins,
 } from "./types";
 import { getLSPDocument } from "./utils";
 
@@ -84,18 +86,32 @@ class Snapshotter {
   }
 }
 
+const defaultGetNotebookEditors = () => {
+  const evs = getNotebook().cellHandles;
+  return Objects.mapValues(evs, (r) => r.current?.editorViewOrNull);
+};
+
 export class NotebookLanguageServerClient implements ILanguageServerClient {
   public readonly documentUri: LSP.DocumentUri;
   private readonly client: ILanguageServerClient;
   private readonly snapshotter: Snapshotter;
+  private readonly getNotebookEditors: () => Record<
+    CellId,
+    EditorView | null | undefined
+  >;
 
   private static readonly SEEN_CELL_DOCUMENT_URIS = new Set<CellDocumentUri>();
 
   constructor(
     client: ILanguageServerClient,
     initialSettings: Record<string, unknown>,
+    getNotebookEditors: () => Record<
+      CellId,
+      EditorView | null | undefined
+    > = defaultGetNotebookEditors,
   ) {
     this.documentUri = getLSPDocument();
+    this.getNotebookEditors = getNotebookEditors;
 
     this.client = client;
     this.patchProcessNotification();
@@ -114,6 +130,17 @@ export class NotebookLanguageServerClient implements ILanguageServerClient {
     this.snapshotter = new Snapshotter(this.getNotebookCode.bind(this));
   }
 
+  onNotification(
+    listener: (n: {
+      jsonrpc: "2.0";
+      id?: null | undefined;
+      method: "textDocument/publishDiagnostics";
+      params: LSP.PublishDiagnosticsParams;
+    }) => void,
+  ): () => boolean {
+    return this.client.onNotification(listener);
+  }
+
   get ready(): boolean {
     return this.client.ready;
   }
@@ -150,16 +177,6 @@ export class NotebookLanguageServerClient implements ILanguageServerClient {
     this.client.close();
   }
 
-  // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  detachPlugin(plugin: any): void {
-    this.client.detachPlugin(plugin);
-  }
-
-  // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  attachPlugin(plugin: any): void {
-    this.client.attachPlugin(plugin);
-  }
-
   private getNotebookCode() {
     return store.get(topologicalCodesAtom);
   }
@@ -371,37 +388,26 @@ export class NotebookLanguageServerClient implements ILanguageServerClient {
     const newEdits = lens.getEditsForNewText(edit.newText);
     const editsToNewCode = new Map(newEdits.map((e) => [e.cellId, e.text]));
 
-    invariant(
-      isClientWithPlugins(this.client),
-      "Expected client with plugins.",
-    );
-
     // Update the code in the plugins manually
-    for (const plugin of this.client.plugins) {
-      const documentUri: string = plugin.documentUri;
-      if (!CellDocumentUri.is(documentUri)) {
-        Logger.warn("Invalid cell document URI", documentUri);
-        continue;
-      }
-
-      const cellId = CellDocumentUri.parse(documentUri);
+    const editors = this.getNotebookEditors();
+    for (const [cellId, ev] of Objects.entries(editors)) {
       const newCode = editsToNewCode.get(cellId);
       if (newCode == null) {
         Logger.warn("No new code for cell", cellId);
         continue;
       }
 
-      if (!plugin.view) {
-        Logger.warn("No view for plugin", plugin);
+      if (!ev) {
+        Logger.warn("No view for plugin", cellId);
         continue;
       }
 
       // Only update if it has changed
-      if (plugin.view.state.doc.toString() !== newCode) {
-        plugin.view.dispatch({
+      if (ev.state.doc.toString() !== newCode) {
+        ev.dispatch({
           changes: {
             from: 0,
-            to: plugin.view.state.doc.length,
+            to: ev.state.doc.length,
             insert: newCode,
           },
         });

package/src/core/codemirror/lsp/transports.ts

@@ -1,5 +1,5 @@
 /* Copyright 2024 Marimo. All rights reserved. */
-import { WebSocketTransport } from "@open-rpc/client-js";
+import { ReconnectingWebSocketTransport } from "@/core/lsp/transport";
 import { waitForConnectionOpen } from "../../network/connection";
 import { getRuntimeManager } from "../../runtime/config";
 
@@ -16,16 +16,10 @@ export function createTransport(
   serverName: "pylsp" | "basedpyright" | "copilot" | "ty",
 ) {
   const runtimeManager = getRuntimeManager();
-  const transport = new WebSocketTransport(
-    runtimeManager.getLSPURL(serverName).toString(),
-  );
-
-  // Override connect to ensure runtime is healthy
-  const originalConnect = transport.connect.bind(transport);
-  transport.connect = async () => {
-    await waitForConnectionOpen();
-    return originalConnect();
-  };
-
-  return transport;
+  return new ReconnectingWebSocketTransport({
+    getWsUrl: () => runtimeManager.getLSPURL(serverName).toString(),
+    waitForConnection: async () => {
+      await waitForConnectionOpen();
+    },
+  });
 }

package/src/core/codemirror/lsp/types.ts

@@ -1,6 +1,5 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 
-import type { EditorView } from "@codemirror/view";
 import type { LanguageServerClient } from "@marimo-team/codemirror-languageserver";
 import type { DocumentUri } from "vscode-languageserver-protocol";
 import type { CellId } from "@/core/cells/ids";
@@ -41,15 +40,3 @@ export function isClientWithNotify(
 } {
   return "notify" in client;
 }
-
-/**
- * Plugins are a @private on `LanguageServerClient`,
- * hiding public use with TypeScript.
- */
-export function isClientWithPlugins(
-  client: ILanguageServerClient,
-): client is ILanguageServerClient & {
-  plugins: { documentUri: string; view?: EditorView }[];
-} {
-  return "plugins" in client;
-}

package/src/core/codemirror/replace-editor-content.ts

@@ -0,0 +1,87 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import type { TransactionSpec } from "@codemirror/state";
+import type { EditorView } from "@codemirror/view";
+
+/**
+ * Replace the entire content of the editor with new content.
+ *
+ * When the editor has focus, this function attempts to preserve the cursor position
+ * by keeping it on the same line and column. If the line no longer exists or shrinks,
+ * the cursor is adjusted accordingly. This prevents the cursor from jumping during
+ * external updates (e.g., auto-formatting or external edits).
+ *
+ * When the editor doesn't have focus, the content is replaced without any special handling.
+ *
+ * @param editor - The CodeMirror editor view
+ * @param newContent - The new content to replace the editor with
+ * @param options - Optional configuration
+ * @param options.preserveCursor - Whether to preserve cursor position when focused (default: true)
+ * @param options.effects - Additional effects to apply to the transaction
+ * @param options.userEvent - User event string for the transaction
+ */
+export function replaceEditorContent(
+  editor: EditorView,
+  newContent: string,
+  options: {
+    preserveCursor?: boolean;
+    effects?: TransactionSpec["effects"];
+    userEvent?: string;
+  } = {},
+): void {
+  const { preserveCursor = true, effects, userEvent } = options;
+  const doc = editor.state.doc;
+  const currentContent = doc.toString();
+
+  // Noop if content is the same
+  if (currentContent === newContent) {
+    return;
+  }
+
+  // If editor has focus and we want to preserve cursor, try to maintain position
+  if (preserveCursor && editor.hasFocus) {
+    const cursorPos = editor.state.selection.main.head;
+
+    // Get the current line and column
+    const currentLine = doc.lineAt(cursorPos);
+    const lineNumber = currentLine.number; // 1-based line number
+    const columnInLine = cursorPos - currentLine.from; // 0-based column in line
+
+    // Calculate new cursor position based on line-aware logic
+    const newDoc = editor.state.update({
+      changes: { from: 0, to: doc.length, insert: newContent },
+    }).state.doc;
+
+    let newCursorPos: number;
+
+    if (lineNumber <= newDoc.lines) {
+      // Line still exists in new document
+      const newLine = newDoc.line(lineNumber);
+      const newLineLength = newLine.length;
+
+      // Keep same column, but clamp to line length if line shrank
+      newCursorPos = newLine.from + Math.min(columnInLine, newLineLength);
+    } else {
+      // Line no longer exists, move to the end of the last line
+      const lastLine = newDoc.line(newDoc.lines);
+      newCursorPos = lastLine.to;
+    }
+
+    // Apply changes with preserved cursor position
+    editor.dispatch({
+      changes: { from: 0, to: doc.length, insert: newContent },
+      selection: {
+        anchor: newCursorPos,
+      },
+      effects,
+      userEvent,
+    });
+  } else {
+    // Editor doesn't have focus or we don't want to preserve cursor
+    editor.dispatch({
+      changes: { from: 0, to: doc.length, insert: newContent },
+      effects,
+      userEvent,
+    });
+  }
+}

package/src/core/config/__tests__/config-schema.test.ts

@@ -57,6 +57,7 @@ test("default UserConfig - empty", () => {
         "activate_on_typing": true,
         "copilot": false,
       },
+      "diagnostics": {},
       "display": {
         "cell_output": "above",
         "code_editor_font_size": 14,
@@ -123,6 +124,7 @@ test("default UserConfig - one level", () => {
         "activate_on_typing": true,
         "copilot": false,
       },
+      "diagnostics": {},
       "display": {
         "cell_output": "above",
         "code_editor_font_size": 14,
@@ -172,6 +174,7 @@ test("default UserConfig - one level", () => {
       runtime: {},
       display: {},
       experimental: {},
+      diagnostics: {},
     }),
   ).toEqual(UserConfigSchema.parse({}));
 });

package/src/core/config/config-schema.ts

@@ -44,6 +44,9 @@ export const DEFAULT_AI_MODEL = "openai/gpt-4o";
  */
 const AUTO_DOWNLOAD_FORMATS = ["html", "markdown", "ipynb"] as const;
 
+const COPILOT_MODES = ["manual", "ask", "agent"] as const;
+export type CopilotMode = (typeof COPILOT_MODES)[number];
+
 const AiConfigSchema = z
   .object({
     api_key: z.string().optional(),
@@ -151,7 +154,7 @@ export const UserConfigSchema = z
     ai: z
       .looseObject({
         rules: z.string().prefault(""),
-        mode: z.enum(["manual", "ask"]).prefault("manual"),
+        mode: z.enum(COPILOT_MODES).prefault("manual"),
         inline_tooltip: z.boolean().prefault(false),
         open_ai: AiConfigSchema.optional(),
         anthropic: AiConfigSchema.optional(),
@@ -183,6 +186,12 @@ export const UserConfigSchema = z
       // Pass through so that we don't remove any extra keys that the user has added.
       .prefault(() => ({})),
     server: z.looseObject({}).prefault(() => ({})),
+    diagnostics: z
+      .looseObject({
+        enabled: z.boolean().optional(),
+        sql_linter: z.boolean().optional(),
+      })
+      .prefault(() => ({})),
     sharing: z
       .looseObject({
         html: z.boolean().optional(),
@@ -204,6 +213,7 @@ export const UserConfigSchema = z
     keymap: {},
     runtime: {},
     display: {},
+    diagnostics: {},
     experimental: {},
     server: {},
     ai: {},
@@ -307,6 +317,7 @@ export function defaultUserConfig(): UserConfig {
     keymap: {},
     runtime: {},
     display: {},
+    diagnostics: {},
     experimental: {},
     server: {},
     ai: {},

package/src/core/config/feature-flag.tsx

@@ -11,9 +11,8 @@ export interface ExperimentalFeatures {
   rtc_v2: boolean;
   performant_table_charts: boolean;
   chat_modes: boolean;
-  sql_linter: boolean;
+  cache_panel: boolean;
   external_agents: boolean;
-  sql_mode: boolean;
   // Add new feature flags here
 }
 
@@ -23,9 +22,8 @@ const defaultValues: ExperimentalFeatures = {
   rtc_v2: false,
   performant_table_charts: false,
   chat_modes: false,
-  sql_linter: true,
+  cache_panel: false,
   external_agents: import.meta.env.DEV,
-  sql_mode: false,
 };
 
 export function getFeatureFlag<T extends keyof ExperimentalFeatures>(

package/src/core/errors/__tests__/errors.test.ts

@@ -57,6 +57,7 @@ describe("getAutoFixes", () => {
     expect(fixes).toHaveLength(1);
     expect(fixes[0].title).toBe("Fix with AI");
 
+    // No fixes without AI
     expect(getAutoFixes(error, { aiEnabled: false })).toHaveLength(0);
   });
 

package/src/core/errors/errors.ts

@@ -1,21 +1,26 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 import type { EditorView } from "@codemirror/view";
 import { invariant } from "@/utils/invariant";
+import { getDatasourceContext } from "../ai/context/providers/datasource";
+import type { AiCompletionCell } from "../ai/state";
 import type { CellId } from "../cells/ids";
 import type { MarimoError } from "../kernel/messages";
 import { wrapInFunction } from "./utils";
 
+interface AIFix {
+  setAiCompletionCell: (opts: AiCompletionCell) => void;
+  triggerFix: boolean;
+}
+
 export interface AutoFix {
   title: string;
   description: string;
+  fixType: "manual" | "ai";
   onFix: (ctx: {
     addCodeBelow: (code: string) => void;
     editor: EditorView | undefined;
     cellId: CellId;
-    setAiCompletionCell?: (cell: {
-      cellId: CellId;
-      initialPrompt?: string;
-    }) => void;
+    aiFix?: AIFix;
   }) => Promise<void>;
 }
 
@@ -31,6 +36,7 @@ export function getAutoFixes(
         title: "Fix: Wrap in a function",
         description:
           "Make this cell's variables local by wrapping the cell in a function.",
+        fixType: "manual",
         onFix: async (ctx) => {
           invariant(ctx.editor, "Editor is null");
           const code = wrapInFunction(ctx.editor.state.doc.toString());
@@ -59,6 +65,7 @@ export function getAutoFixes(
       {
         title: `Fix: Add '${cellCode}'`,
         description: "Add a new cell for the missing import",
+        fixType: "manual",
         onFix: async (ctx) => {
           ctx.addCodeBelow(cellCode);
         },
@@ -75,10 +82,17 @@ export function getAutoFixes(
       {
         title: "Fix with AI",
         description: "Fix the SQL statement",
+        fixType: "ai",
         onFix: async (ctx) => {
-          ctx.setAiCompletionCell?.({
+          const datasourceContext = getDatasourceContext(ctx.cellId);
+          let initialPrompt = `Fix the SQL statement: ${error.msg}.`;
+          if (datasourceContext) {
+            initialPrompt += `\nDatabase schema: ${datasourceContext}`;
+          }
+          ctx.aiFix?.setAiCompletionCell({
             cellId: ctx.cellId,
-            initialPrompt: `Fix the SQL statement: ${error.msg}`,
+            initialPrompt: initialPrompt,
+            triggerImmediately: ctx.aiFix.triggerFix,
           });
         },
       },

package/src/core/hotkeys/hotkeys.ts

@@ -239,6 +239,18 @@ const DEFAULT_HOT_KEY = {
     group: "Editing",
     key: "F2",
   },
+  "cell.copyLineUp": {
+    name: "Copy line(s) up",
+    group: "Editing",
+    key: "Alt-Shift-ArrowUp",
+    editable: false,
+  },
+  "cell.copyLineDown": {
+    name: "Copy line(s) down",
+    group: "Editing",
+    key: "Alt-Shift-ArrowDown",
+    editable: false,
+  },
 
   // Markdown
   "markdown.bold": {

package/src/core/islands/bridge.ts

@@ -180,6 +180,8 @@ export class IslandsPyodideBridge implements RunRequests, EditRequests {
   listSecretKeys = throwNotImplemented;
   writeSecret = throwNotImplemented;
   invokeAiTool = throwNotImplemented;
+  clearCache = throwNotImplemented;
+  getCacheInfo = throwNotImplemented;
 
   private async putControlRequest(operation: object) {
     await this.rpc.proxy.request.bridge({

package/src/core/islands/main.ts

@@ -185,6 +185,10 @@ export async function initialize() {
         return;
       case "reconnected":
         return;
+      case "cache-cleared":
+        return;
+      case "cache-info-fetched":
+        return;
       default:
         logNever(msg.data);
     }

package/src/core/kernel/messages.ts

@@ -43,6 +43,7 @@ export type SecretKeysResult = OperationMessageData<"secret-keys-result">;
 export type StartupLogs = OperationMessageData<"startup-logs">;
 export type CellMessage = OperationMessageData<"cell-op">;
 export type Capabilities = OperationMessageData<"kernel-ready">["capabilities"];
+export type CacheInfoFetched = OperationMessageData<"cache-info-fetched">;
 
 export type MessageOperationUnion = schemas["KnownUnions"]["operation"];