@marimo-team/frontend 0.16.1 → 0.16.3

package/src/core/datasets/request-registry.ts

@@ -1,10 +1,16 @@
 /* Copyright 2024 Marimo. All rights reserved. */
-import type { SQLTableListPreview, SQLTablePreview } from "../kernel/messages";
+import type {
+  SQLTableListPreview,
+  SQLTablePreview,
+  ValidateSQLResult,
+} from "../kernel/messages";
+import { CachingRequestRegistry } from "../network/CachingRequestRegistry";
 import { DeferredRequestRegistry } from "../network/DeferredRequestRegistry";
 import { getRequestClient } from "../network/requests";
 import type {
   PreviewSQLTableListRequest,
   PreviewSQLTableRequest,
+  ValidateSQLRequest,
 } from "../network/types";
 
 // We make a request to the backend to preview the table, passing in Engine, DB, Schema, and Table
@@ -32,3 +38,20 @@ export const PreviewSQLTableList = new DeferredRequestRegistry<
     ...req,
   });
 });
+
+export const ValidateSQL = new CachingRequestRegistry(
+  new DeferredRequestRegistry<
+    Omit<ValidateSQLRequest, "requestId">,
+    ValidateSQLResult
+  >("validate-sql", async (requestId, req) => {
+    const client = getRequestClient();
+    await client.validateSQL({
+      requestId: requestId,
+      ...req,
+    });
+  }),
+  {
+    // Only keep the last 3 validation results
+    maxSize: 3,
+  },
+);

package/src/core/dom/events.ts

@@ -43,7 +43,7 @@ export const MarimoIncomingMessageEvent = defineCustomEvent(
 )<{
   objectId: UIElementId;
   message: unknown;
-  buffers: DataView[] | undefined;
+  buffers: readonly DataView[];
 }>();
 export type MarimoIncomingMessageEventType = ReturnType<
   typeof MarimoIncomingMessageEvent.create

package/src/core/dom/outline.ts

@@ -57,7 +57,7 @@ export function headingToIdentifier(heading: Element): OutlineItem["by"] {
   return { path: `//${heading.tagName}[contains(., "${name}")]` };
 }
 
-export function mergeOutlines(outlines: Array<Outline | null>): Outline {
+export function mergeOutlines(outlines: (Outline | null)[]): Outline {
   return {
     items: outlines.filter(Boolean).flatMap((outline) => outline.items),
   };
@@ -104,7 +104,7 @@ export function canCollapseOutline(outline: Outline | null): boolean {
  */
 export function findCollapseRange(
   startIndex: number,
-  outlines: Array<Outline | null>,
+  outlines: (Outline | null)[],
 ): [number, number] | null {
   // Higher header is the lowest value
   const getHighestHeader = (outline: Outline) => {

package/src/core/dom/uiregistry.ts

@@ -1,8 +1,5 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 
-import { byteStringToDataView } from "@/utils/data-views";
-import type { Base64String } from "@/utils/json/base64";
-import { typedAtob } from "@/utils/json/base64";
 import { Logger } from "@/utils/Logger";
 import type { CellId, UIElementId } from "../cells/ids";
 import {
@@ -136,15 +133,12 @@ export class UIElementRegistry {
   broadcastMessage(
     objectId: UIElementId,
     message: unknown,
-    buffers: Base64String[] | undefined | null,
+    buffers: readonly DataView[],
   ): void {
     const entry = this.entries.get(objectId);
     if (entry === undefined) {
       Logger.warn("UIElementRegistry missing entry", objectId);
     } else {
-      const toDataView = (base64: Base64String) => {
-        return byteStringToDataView(typedAtob(base64));
-      };
       entry.elements.forEach((element) => {
         element.dispatchEvent(
           MarimoIncomingMessageEvent.create({
@@ -153,7 +147,7 @@ export class UIElementRegistry {
             detail: {
               objectId: objectId,
               message: message,
-              buffers: buffers ? buffers.map(toDataView) : undefined,
+              buffers: buffers,
             },
           }),
         );

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

@@ -16,6 +16,10 @@ describe("getImportCode", () => {
   });
 });
 
+const opts = {
+  aiEnabled: true,
+};
+
 describe("getAutoFixes", () => {
   it("returns wrap in function fix for multiple-defs error", () => {
     const error: MarimoError = {
@@ -24,7 +28,7 @@ describe("getAutoFixes", () => {
       cells: ["foo"],
     };
 
-    const fixes = getAutoFixes(error);
+    const fixes = getAutoFixes(error, opts);
     expect(fixes).toHaveLength(1);
     expect(fixes[0].title).toBe("Fix: Wrap in a function");
   });
@@ -37,11 +41,25 @@ describe("getAutoFixes", () => {
       raising_cell: null,
     };
 
-    const fixes = getAutoFixes(error);
+    const fixes = getAutoFixes(error, opts);
     expect(fixes).toHaveLength(1);
     expect(fixes[0].title).toBe("Fix: Add 'import numpy as np'");
   });
 
+  it("returns sql fix for sql-error error", () => {
+    const error: MarimoError = {
+      type: "sql-error",
+      msg: "syntax error",
+      sql_statement: "SELECT * FROM table",
+    };
+
+    const fixes = getAutoFixes(error, opts);
+    expect(fixes).toHaveLength(1);
+    expect(fixes[0].title).toBe("Fix with AI");
+
+    expect(getAutoFixes(error, { aiEnabled: false })).toHaveLength(0);
+  });
+
   it("returns no fixes for NameError with unknown import", () => {
     const error: MarimoError = {
       type: "exception",
@@ -50,7 +68,7 @@ describe("getAutoFixes", () => {
       raising_cell: null,
     };
 
-    expect(getAutoFixes(error)).toHaveLength(0);
+    expect(getAutoFixes(error, opts)).toHaveLength(0);
   });
 
   it("returns no fixes for other error types", () => {
@@ -59,6 +77,6 @@ describe("getAutoFixes", () => {
       msg: "invalid syntax",
     };
 
-    expect(getAutoFixes(error)).toHaveLength(0);
+    expect(getAutoFixes(error, opts)).toHaveLength(0);
   });
 });

package/src/core/errors/errors.ts

@@ -12,10 +12,19 @@ export interface AutoFix {
     addCodeBelow: (code: string) => void;
     editor: EditorView | undefined;
     cellId: CellId;
+    setAiCompletionCell?: (cell: {
+      cellId: CellId;
+      initialPrompt?: string;
+    }) => void;
   }) => Promise<void>;
 }
 
-export function getAutoFixes(error: MarimoError): AutoFix[] {
+export function getAutoFixes(
+  error: MarimoError,
+  opts: {
+    aiEnabled: boolean;
+  },
+): AutoFix[] {
   if (error.type === "multiple-defs") {
     return [
       {
@@ -57,6 +66,25 @@ export function getAutoFixes(error: MarimoError): AutoFix[] {
     ];
   }
 
+  if (error.type === "sql-error") {
+    // Only show AI fix if AI is enabled
+    if (!opts.aiEnabled) {
+      return [];
+    }
+    return [
+      {
+        title: "Fix with AI",
+        description: "Fix the SQL statement",
+        onFix: async (ctx) => {
+          ctx.setAiCompletionCell?.({
+            cellId: ctx.cellId,
+            initialPrompt: `Fix the SQL statement: ${error.msg}`,
+          });
+        },
+      },
+    ];
+  }
+
   return [];
 }
 

package/src/core/errors/state.ts

@@ -7,7 +7,7 @@ import { generateUUID } from "@/utils/uuid";
 import type { Banner } from "../kernel/messages";
 
 interface BannerState {
-  banners: Array<Identified<Banner>>;
+  banners: Identified<Banner>[];
 }
 
 const { valueAtom: bannersAtom, useActions } = createReducerAndAtoms(

package/src/core/islands/bridge.ts

@@ -153,6 +153,7 @@ export class IslandsPyodideBridge implements RunRequests, EditRequests {
   previewSQLTable = throwNotImplemented;
   previewSQLTableList = throwNotImplemented;
   previewDataSourceConnection = throwNotImplemented;
+  validateSQL = throwNotImplemented;
   openFile = throwNotImplemented;
   sendListFiles = throwNotImplemented;
   sendSearchFiles = throwNotImplemented;

package/src/core/islands/main.ts

@@ -16,7 +16,7 @@ import { renderHTML } from "@/plugins/core/RenderHTML";
 import { initializePlugins } from "@/plugins/plugins";
 import { logNever } from "@/utils/assertNever";
 import { Functions } from "@/utils/functions";
-import type { Base64String } from "@/utils/json/base64";
+import { safeExtractSetUIElementMessageBuffers } from "@/utils/json/base64";
 import { jsonParseWithSpecialChar } from "@/utils/json/json-parser";
 import { Logger } from "@/utils/Logger";
 import {
@@ -119,6 +119,7 @@ export async function initialize() {
       case "sql-table-list-preview":
       case "datasets":
       case "data-source-connections":
+      case "validate-sql-result":
       case "secret-keys-result":
       case "startup-logs":
         // Unsupported
@@ -144,7 +145,7 @@ export async function initialize() {
         UI_ELEMENT_REGISTRY.broadcastMessage(
           msg.data.ui_element as UIElementId,
           msg.data.message,
-          msg.data.buffers as Base64String[],
+          safeExtractSetUIElementMessageBuffers(msg.data),
         );
         return;
 

package/src/core/islands/parse.ts

@@ -88,9 +88,7 @@ export function parseMarimoIslandApps(): MarimoIslandApp[] {
   return [...apps.values()];
 }
 
-export function createMarimoFile(app: {
-  cells: Array<{ code: string }>;
-}): string {
+export function createMarimoFile(app: { cells: { code: string }[] }): string {
   const lines = [
     "import marimo",
     "app = marimo.App()",

package/src/core/kernel/messages.ts

@@ -24,7 +24,7 @@ export type OutputChannel = schemas["CellChannel"];
 export type CellOutput = schemas["CellOutput"];
 export type MarimoError = Extract<
   CellOutput["data"],
-  Array<{ type: string }>
+  { type: string }[]
 >[number];
 export type OutputMessage = schemas["CellOutput"];
 export type CompletionOption = schemas["CompletionResult"]["options"][0];
@@ -38,6 +38,7 @@ export type DataColumnPreview = OperationMessageData<"data-column-preview">;
 export type SQLTablePreview = OperationMessageData<"sql-table-preview">;
 export type SQLTableListPreview =
   OperationMessageData<"sql-table-list-preview">;
+export type ValidateSQLResult = OperationMessageData<"validate-sql-result">;
 export type SecretKeysResult = OperationMessageData<"secret-keys-result">;
 export type StartupLogs = OperationMessageData<"startup-logs">;
 export type CellMessage = OperationMessageData<"cell-op">;

package/src/core/network/CachingRequestRegistry.ts

@@ -0,0 +1,74 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import { LRUCache } from "@/utils/lru";
+import type {
+  DeferredRequestRegistry,
+  RequestId,
+} from "./DeferredRequestRegistry";
+
+type ToKey<REQ> = (request: REQ) => string;
+
+interface CachingOptions<REQ> {
+  toKey?: ToKey<REQ>;
+  maxSize?: number;
+}
+
+/**
+ * Light wrapper adding memoization and in-flight de-duplication on top of
+ * DeferredRequestRegistry, keyed by a string representation of the request.
+ */
+export class CachingRequestRegistry<REQ, RES> {
+  private delegate: DeferredRequestRegistry<REQ, RES>;
+  private toKey: ToKey<REQ>;
+  private cache: LRUCache<string, Promise<RES>>;
+
+  static jsonStringifySortKeys<T>(): ToKey<T> {
+    return (o: T) => {
+      if (typeof o !== "object" || o === null) {
+        return String(o);
+      }
+      return JSON.stringify(o, Object.keys(o).sort(), 2);
+    };
+  }
+
+  constructor(
+    delegate: DeferredRequestRegistry<REQ, RES>,
+    options: CachingOptions<REQ> = {},
+  ) {
+    this.delegate = delegate;
+    this.toKey =
+      options.toKey ?? CachingRequestRegistry.jsonStringifySortKeys();
+    const maxSize = options.maxSize ?? 128;
+    this.cache = new LRUCache<string, Promise<RES>>(maxSize);
+  }
+
+  /**
+   * Resolve via cache if present, else delegate; de-duplicates concurrent
+   * requests with the same key and stores successful results in the cache.
+   */
+  public request(req: REQ): Promise<RES> {
+    const key = this.toKey(req);
+
+    const cached = this.cache.get(key);
+    if (cached !== undefined) {
+      return cached;
+    }
+
+    const promise = this.delegate.request(req);
+    this.cache.set(key, promise);
+    return promise.catch((err) => {
+      this.cache.delete(key);
+      throw err;
+    });
+  }
+
+  // Path through to the delegate
+  public resolve(requestId: RequestId, response: RES) {
+    this.delegate.resolve(requestId, response);
+  }
+
+  // Path through to the delegate
+  public reject(requestId: RequestId, error: Error) {
+    this.delegate.reject(requestId, error);
+  }
+}

package/src/core/network/DeferredRequestRegistry.ts

@@ -45,7 +45,9 @@ export class DeferredRequestRegistry<REQ, RES> {
   async request(opts: REQ): Promise<RES> {
     if (this.opts.resolveExistingRequests) {
       const result = this.opts.resolveExistingRequests();
-      this.requests.forEach((deferred) => deferred.resolve(result));
+      for (const deferred of this.requests.values()) {
+        deferred.resolve(result);
+      }
       this.requests.clear();
     }
 

package/src/core/network/__tests__/CachingRequestRegistry.test.ts

@@ -0,0 +1,73 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+import { beforeEach, describe, expect, it, vi } from "vitest";
+import { CachingRequestRegistry } from "../CachingRequestRegistry";
+import {
+  DeferredRequestRegistry,
+  type RequestId,
+} from "../DeferredRequestRegistry";
+
+vi.mock("@/utils/uuid", () => ({
+  generateUUID: vi.fn().mockReturnValue("uuid"),
+}));
+
+describe("CachingRequestRegistry", () => {
+  const REQUEST_ID = "uuid" as RequestId;
+  let makeRequestMock = vi.fn();
+  let delegate: DeferredRequestRegistry<unknown, unknown>;
+  let caching: CachingRequestRegistry<unknown, unknown>;
+
+  beforeEach(() => {
+    makeRequestMock = vi.fn().mockResolvedValue(undefined);
+    delegate = new DeferredRequestRegistry("operation", makeRequestMock);
+    caching = new CachingRequestRegistry(delegate);
+  });
+
+  it("should cache successful responses for identical requests", async () => {
+    const req = { a: 1 };
+
+    const p1 = caching.request(req);
+    expect(makeRequestMock).toHaveBeenCalledTimes(1);
+    expect(makeRequestMock).toHaveBeenCalledWith(REQUEST_ID, req);
+
+    // Resolve first request
+    delegate.resolve(REQUEST_ID, "response");
+    await expect(p1).resolves.toBe("response");
+
+    // Second call with equivalent request gets served from cache
+    const p2 = caching.request({ a: 1 });
+    expect(makeRequestMock).toHaveBeenCalledTimes(1);
+    await expect(p2).resolves.toBe("response");
+  });
+
+  it("should de-duplicate in-flight requests with the same key", async () => {
+    const req = { q: "select *" };
+
+    const p1 = caching.request(req);
+    const p2 = caching.request({ q: "select *" });
+
+    // Only one network invocation while in-flight
+    expect(makeRequestMock).toHaveBeenCalledTimes(1);
+    expect(p1).toStrictEqual(p2);
+
+    // Resolve and ensure both resolve to same result
+    delegate.resolve(REQUEST_ID, "ok");
+    await expect(p1).resolves.toBe("ok");
+    await expect(p2).resolves.toBe("ok");
+  });
+
+  it("should not cache errors", async () => {
+    // First call rejects
+    makeRequestMock.mockRejectedValueOnce(new Error("boom"));
+
+    await expect(caching.request({ x: 1 })).rejects.toThrow("boom");
+    expect(makeRequestMock).toHaveBeenCalledTimes(1);
+
+    // Next call should attempt again (not cached)
+    const p2 = caching.request({ x: 1 });
+    expect(makeRequestMock).toHaveBeenCalledTimes(2);
+
+    // Resolve the second request
+    delegate.resolve(REQUEST_ID, "ok");
+    await expect(p2).resolves.toBe("ok");
+  });
+});

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

@@ -207,6 +207,13 @@ export function createNetworkRequests(): EditRequests & RunRequests {
         })
         .then(handleResponseReturnNull);
     },
+    validateSQL: (request) => {
+      return getClient()
+        .POST("/api/sql/validate", {
+          body: request,
+        })
+        .then(handleResponseReturnNull);
+    },
     openFile: async (request) => {
       await waitForConnectionOpen();
       await getClient()

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

@@ -57,6 +57,7 @@ export function createStaticRequests(): EditRequests & RunRequests {
     previewSQLTable: throwNotInEditMode,
     previewSQLTableList: throwNotInEditMode,
     previewDataSourceConnection: throwNotInEditMode,
+    validateSQL: throwNotInEditMode,
     openFile: throwNotInEditMode,
     getUsageStats: throwNotInEditMode,
     sendListFiles: throwNotInEditMode,

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

@@ -35,6 +35,7 @@ export function createErrorToastingRequests(
     previewSQLTable: "Failed to fetch SQL table",
     previewSQLTableList: "Failed to fetch SQL table list",
     previewDataSourceConnection: "Failed to preview data source connection",
+    validateSQL: "Failed to validate SQL",
     openFile: "Failed to open file",
     getUsageStats: "", // No toast
     sendListFiles: "Failed to list files",

package/src/core/network/types.ts

@@ -62,6 +62,7 @@ export type PreviewSQLTableRequest = schemas["PreviewSQLTableRequest"];
 export type PreviewSQLTableListRequest = schemas["PreviewSQLTableListRequest"];
 export type PreviewDataSourceConnectionRequest =
   schemas["PreviewDataSourceConnectionRequest"];
+export type ValidateSQLRequest = schemas["ValidateSQLRequest"];
 export type PdbRequest = schemas["PdbRequest"];
 export type ReadCodeResponse = schemas["ReadCodeResponse"];
 export type RecentFilesResponse = schemas["RecentFilesResponse"];
@@ -140,7 +141,8 @@ export interface EditRequests {
   previewDataSourceConnection: (
     request: PreviewDataSourceConnectionRequest,
   ) => Promise<null>;
-  openFile: (request: { path: string }) => Promise<null>;
+  validateSQL: (request: ValidateSQLRequest) => Promise<null>;
+  openFile: (request: { path: string; lineNumber?: number }) => Promise<null>;
   getUsageStats: () => Promise<UsageResponse>;
   // Debugger
   sendPdb: (request: PdbRequest) => Promise<null>;

package/src/core/variables/state.ts

@@ -38,11 +38,11 @@ const {
   },
   setMetadata: (
     state,
-    metadata: Array<{
+    metadata: {
       name: VariableName;
       value?: string | null;
       dataType?: string | null;
-    }>,
+    }[],
   ) => {
     const newVariables = { ...state };
     for (const { name, value, dataType } of metadata) {

package/src/core/wasm/__tests__/state.test.ts

@@ -13,7 +13,7 @@ import { hasAnyOutputAtom } from "../state";
 
 describe("hasAnyOutputAtom", () => {
   const createNotebookState = (
-    outputs: Array<OutputMessage | null>,
+    outputs: (OutputMessage | null)[],
   ): NotebookState => ({
     ...initialNotebookState(),
     cellIds: new MultiColumn([

package/src/core/wasm/bridge.ts

@@ -496,6 +496,11 @@ export class PyodideBridge implements RunRequests, EditRequests {
       return null;
     };
 
+  validateSQL: EditRequests["validateSQL"] = async (request) => {
+    await this.putControlRequest(request);
+    return null;
+  };
+
   sendModelValue: RunRequests["sendModelValue"] = async (request) => {
     await this.putControlRequest(request);
     return null;

package/src/core/websocket/useMarimoWebSocket.tsx

@@ -16,7 +16,10 @@ import {
 } from "@/plugins/impl/anywidget/model";
 import { logNever } from "@/utils/assertNever";
 import { prettyError } from "@/utils/errors";
-import type { Base64String, JsonString } from "@/utils/json/base64";
+import {
+  type JsonString,
+  safeExtractSetUIElementMessageBuffers,
+} from "@/utils/json/base64";
 import { jsonParseWithSpecialChar } from "@/utils/json/json-parser";
 import { Logger } from "@/utils/Logger";
 import { reloadSafe } from "@/utils/reload-safe";
@@ -32,6 +35,7 @@ import type { ConnectionName } from "../datasets/engines";
 import {
   PreviewSQLTable,
   PreviewSQLTableList,
+  ValidateSQL,
 } from "../datasets/request-registry";
 import { useDatasetsActions } from "../datasets/state";
 import { UI_ELEMENT_REGISTRY } from "../dom/uiregistry";
@@ -111,7 +115,7 @@ export function useMarimoWebSocket(opts: {
         const modelId = msg.data.model_id;
         const uiElement = msg.data.ui_element;
         const message = msg.data.message;
-        const buffers = (msg.data.buffers ?? []) as Base64String[];
+        const buffers = safeExtractSetUIElementMessageBuffers(msg.data);
 
         if (modelId && isMessageWidgetState(message)) {
           handleWidgetMessage({
@@ -238,6 +242,9 @@ export function useMarimoWebSocket(opts: {
       case "sql-table-list-preview":
         PreviewSQLTableList.resolve(msg.data.request_id as RequestId, msg.data);
         return;
+      case "validate-sql-result":
+        ValidateSQL.resolve(msg.data.request_id as RequestId, msg.data);
+        return;
       case "secret-keys-result":
         SECRETS_REGISTRY.resolve(msg.data.request_id as RequestId, msg.data);
         return;

package/src/custom.d.ts

@@ -20,5 +20,5 @@ interface JSON {
 
 // Improve type inference for Array.filter with BooleanConstructor
 interface Array<T> {
-  filter(predicate: BooleanConstructor): Array<NonNullable<T>>;
+  filter(predicate: BooleanConstructor): NonNullable<T>[];
 }

package/src/hooks/useCellRenderCount.ts

@@ -8,6 +8,7 @@ export function useCellRenderCount() {
 
       const currentCount = Number.parseInt(
         document.body.dataset.cellRenderCount || "0",
+        10,
       );
       document.body.dataset.cellRenderCount = (currentCount + 1).toString();
     },

package/src/hooks/useResizeHandle.ts

@@ -29,7 +29,10 @@ export const useResizeHandle = ({
       return;
     }
 
-    let width = Number.parseInt(window.getComputedStyle(resizableDiv).width);
+    let width = Number.parseInt(
+      window.getComputedStyle(resizableDiv).width,
+      10,
+    );
     let lastX = 0;
     let isResizing = false;
     let activeDirection: "left" | "right" | null = null;

package/src/plugins/core/RenderHTML.tsx

@@ -89,7 +89,7 @@ const replaceSrcScripts = (domNode: DOMNode): JSX.Element | undefined => {
       script.src = src;
       document.head.append(script);
     }
-    // eslint-disable-next-line react/jsx-no-useless-fragment
+    // biome-ignore lint/complexity/noUselessFragments: this is intentional
     return <></>;
   }
 };
@@ -163,7 +163,6 @@ export const renderHTML = ({ html, additionalReplacements = [] }: Options) => {
           return transformed;
         }
       }
-      // eslint-disable-next-line react/jsx-no-useless-fragment
       return reactNode as JSX.Element;
     },
   });