@marimo-team/islands 0.16.5 → 0.17.0

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

@@ -195,35 +195,36 @@ describe("createNotebookLens", () => {
 });
 
 describe("NotebookLanguageServerClient", () => {
-  let plugins: any[] = [];
   let mockClient: Mocked<ILanguageServerClient>;
   let notebookClient: NotebookLanguageServerClient;
 
   beforeEach(() => {
-    plugins = [];
     mockClient = {
-      plugins,
       ready: true,
       capabilities: {},
       initializePromise: Promise.resolve(),
+      clientCapabilities: {},
+      completionItemResolve: vi.fn(),
       initialize: vi.fn(),
       close: vi.fn(),
-      detachPlugin: vi.fn().mockImplementation((plugin) => {
-        plugins = plugins.filter((p) => p !== plugin);
-      }),
-      attachPlugin: vi.fn().mockImplementation((plugin) => {
-        plugins.push(plugin);
-      }),
+      onNotification: vi.fn(),
       textDocumentDidOpen: vi.fn(),
       textDocumentDidChange: vi.fn(),
       textDocumentHover: vi.fn(),
       textDocumentCompletion: vi.fn(),
+      textDocumentDefinition: vi.fn(),
+      textDocumentPrepareRename: vi.fn(),
+      textDocumentCodeAction: vi.fn(),
+      textDocumentSignatureHelp: vi.fn(),
       textDocumentRename: vi.fn(),
-      processNotification: vi.fn(),
-      notify: vi.fn(),
-      request: vi.fn(),
-    } as unknown as Mocked<ILanguageServerClient>;
-    notebookClient = new NotebookLanguageServerClient(mockClient, {});
+    };
+    (mockClient as any).processNotification = vi.fn();
+    (mockClient as any).notify = vi.fn();
+    notebookClient = new NotebookLanguageServerClient(mockClient, {}, () => ({
+      [Cells.cell1]: new EditorView({ doc: "# this is a comment" }),
+      [Cells.cell2]: new EditorView({ doc: "import math\nimport numpy" }),
+      [Cells.cell3]: new EditorView({ doc: "print(math.sqrt(4))" }),
+    }));
 
     // Mock the atom instead of the instance method
     vi.spyOn(store, "get").mockImplementation((atom) => {
@@ -404,6 +405,12 @@ describe("NotebookLanguageServerClient", () => {
       });
       expect(mockView3.state.doc.toString()).toBe("print(math.sqrt(4))");
 
+      (notebookClient as any).getNotebookEditors = () => ({
+        [Cells.cell1]: mockView1,
+        [Cells.cell2]: mockView2,
+        [Cells.cell3]: mockView3,
+      });
+
       // Setup rename params
       const renameParams: LSP.RenameParams = {
         textDocument: { uri: CellDocumentUri.of(Cells.cell2) },

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

@@ -1,6 +1,5 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 
-import type { LanguageServerPlugin } from "@marimo-team/codemirror-languageserver";
 import type * as LSP from "vscode-languageserver-protocol";
 import { Objects } from "@/utils/objects";
 import type { ILanguageServerClient } from "./types";
@@ -24,6 +23,26 @@ export class FederatedLanguageServerClient implements ILanguageServerClient {
     this.documentUri = getLSPDocument();
   }
 
+  onNotification(
+    listener: (n: {
+      jsonrpc: "2.0";
+      id?: null | undefined;
+      method: "textDocument/publishDiagnostics";
+      params: LSP.PublishDiagnosticsParams;
+    }) => void,
+  ): () => boolean {
+    const callbacks: Array<() => boolean> = [];
+    for (const client of this.clients) {
+      callbacks.push(client.onNotification(listener));
+    }
+    return () => {
+      for (const cb of callbacks) {
+        cb();
+      }
+      return true;
+    };
+  }
+
   get clientCapabilities(): LSP.ClientCapabilities | undefined {
     const capabilities = this.clients
       .map((client) => {
@@ -153,14 +172,6 @@ export class FederatedLanguageServerClient implements ILanguageServerClient {
     return null;
   }
 
-  attachPlugin(plugin: LanguageServerPlugin): void {
-    this.clients.forEach((client) => client.attachPlugin(plugin));
-  }
-
-  detachPlugin(plugin: LanguageServerPlugin): void {
-    this.clients.forEach((client) => client.detachPlugin(plugin));
-  }
-
   // Merge completions from all clients
   async textDocumentCompletion(
     params: LSP.CompletionParams,

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/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

@@ -183,6 +183,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 +210,7 @@ export const UserConfigSchema = z
     keymap: {},
     runtime: {},
     display: {},
+    diagnostics: {},
     experimental: {},
     server: {},
     ai: {},
@@ -307,6 +314,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/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"];
 

package/src/core/network/requests-network.ts

@@ -409,5 +409,19 @@ export function createNetworkRequests(): EditRequests & RunRequests {
         })
         .then(handleResponse);
     },
+    clearCache: async () => {
+      return getClient()
+        .POST("/api/cache/clear", {
+          body: {},
+        })
+        .then(handleResponseReturnNull);
+    },
+    getCacheInfo: async () => {
+      return getClient()
+        .POST("/api/cache/info", {
+          body: {},
+        })
+        .then(handleResponseReturnNull);
+    },
   };
 }

package/src/core/network/requests-static.ts

@@ -86,5 +86,7 @@ export function createStaticRequests(): EditRequests & RunRequests {
     listSecretKeys: throwNotInEditMode,
     writeSecret: throwNotInEditMode,
     invokeAiTool: throwNotInEditMode,
+    clearCache: throwNotInEditMode,
+    getCacheInfo: throwNotInEditMode,
   };
 }

package/src/core/network/requests-toasting.ts

@@ -64,6 +64,8 @@ export function createErrorToastingRequests(
     listSecretKeys: "Failed to fetch secrets",
     writeSecret: "Failed to write secret",
     invokeAiTool: "Failed to invoke AI tool",
+    clearCache: "Failed to clear cache",
+    getCacheInfo: "", // No toast
   };
 
   const handlers = {} as EditRequests & RunRequests;

package/src/core/network/types.ts

@@ -99,6 +99,8 @@ export type OpenTutorialRequest = schemas["OpenTutorialRequest"];
 export type TutorialId = OpenTutorialRequest["tutorialId"];
 export type InvokeAiToolRequest = schemas["InvokeAiToolRequest"];
 export type InvokeAiToolResponse = schemas["InvokeAiToolResponse"];
+export type ClearCacheRequest = schemas["ClearCacheRequest"];
+export type GetCacheInfoRequest = schemas["GetCacheInfoRequest"];
 
 /**
  * Requests sent to the BE during run/edit mode.
@@ -188,6 +190,9 @@ export interface EditRequests {
   writeSecret: (request: CreateSecretRequest) => Promise<null>;
   // AI Tool requests
   invokeAiTool: (request: InvokeAiToolRequest) => Promise<InvokeAiToolResponse>;
+  // Cache requests
+  clearCache: () => Promise<null>;
+  getCacheInfo: () => Promise<null>;
 }
 
 export type RequestKey = keyof (EditRequests & RunRequests);