@marimo-team/frontend 0.15.1-dev33 → 0.15.1-dev35

package/src/core/ai/context/__tests__/utils.test.ts

@@ -0,0 +1,174 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+
+import { describe, expect, it } from "vitest";
+import { contextToXml } from "../utils";
+
+describe("contextToXml", () => {
+  it("should convert basic context to XML", () => {
+    const context = {
+      type: "data",
+      data: {
+        name: "dataset1",
+        source: "memory",
+      },
+    };
+
+    const result = contextToXml(context);
+    expect(result).toBe('<data name="dataset1" source="memory"></data>');
+  });
+
+  it("should handle context with details", () => {
+    const context = {
+      type: "variable",
+      data: {
+        name: "my_var",
+        dataType: "string",
+      },
+      details: "This is a string variable",
+    };
+
+    const result = contextToXml(context);
+    expect(result).toBe(
+      '<variable name="my_var" dataType="string">This is a string variable</variable>',
+    );
+  });
+
+  it("should escape XML characters in attributes", () => {
+    const context = {
+      type: "data",
+      data: {
+        name: "dataset<>&\"'",
+        description: "Contains special chars",
+      },
+    };
+
+    const result = contextToXml(context);
+    expect(result).toMatchInlineSnapshot(
+      `"<data name="dataset&lt;&gt;&"'" description="Contains special chars"></data>"`,
+    );
+  });
+
+  it("should escape XML characters in details", () => {
+    const context = {
+      type: "error",
+      data: {
+        name: "error1",
+      },
+      details: "Error message with <tags> & \"quotes\" and 'apostrophes'",
+    };
+
+    const result = contextToXml(context);
+    expect(result).toMatchInlineSnapshot(
+      `"<error name="error1">Error message with &lt;tags&gt; & "quotes" and 'apostrophes'</error>"`,
+    );
+  });
+
+  it("should handle undefined values in data", () => {
+    const context = {
+      type: "variable",
+      data: {
+        name: "my_var",
+        dataType: undefined,
+        value: "test",
+      },
+    };
+
+    const result = contextToXml(context);
+    expect(result).toBe('<variable name="my_var" value="test"></variable>');
+  });
+
+  it("should handle empty data object", () => {
+    const context = {
+      type: "empty",
+      data: {},
+    };
+
+    const result = contextToXml(context);
+    expect(result).toBe("<empty></empty>");
+  });
+
+  it("should handle numeric values", () => {
+    const context = {
+      type: "metric",
+      data: {
+        count: 42,
+        percentage: 85.5,
+        isActive: true,
+      },
+    };
+
+    const result = contextToXml(context);
+    expect(result).toBe(
+      '<metric count="42" percentage="85.5" isActive="true"></metric>',
+    );
+  });
+
+  it("should handle complex nested data", () => {
+    const context = {
+      type: "complex",
+      data: {
+        name: "test",
+        config: JSON.stringify({ key: "value", nested: { prop: "test" } }),
+      },
+      details: "Complex configuration data",
+    };
+
+    const result = contextToXml(context);
+    expect(result).toMatchInlineSnapshot(
+      `"<complex name="test" config="{"key":"value","nested":{"prop":"test"}}">Complex configuration data</complex>"`,
+    );
+  });
+
+  it("should handle boolean values", () => {
+    const context = {
+      type: "flags",
+      data: {
+        enabled: true,
+        visible: false,
+      },
+    };
+
+    const result = contextToXml(context);
+    expect(result).toBe('<flags enabled="true" visible="false"></flags>');
+  });
+
+  it("should handle null values", () => {
+    const context = {
+      type: "nullable",
+      data: {
+        name: "test",
+        value: null,
+      },
+    };
+
+    const result = contextToXml(context);
+    expect(result).toBe('<nullable name="test" value="null"></nullable>');
+  });
+
+  it("should handle multiline details", () => {
+    const context = {
+      type: "multiline",
+      data: {
+        name: "test",
+      },
+      details: "Line 1\nLine 2\nLine 3",
+    };
+
+    const result = contextToXml(context);
+    expect(result).toBe(
+      '<multiline name="test">Line 1\nLine 2\nLine 3</multiline>',
+    );
+  });
+
+  it("should handle special characters in type name", () => {
+    const context = {
+      type: "data-source",
+      data: {
+        name: "test",
+      },
+    };
+
+    const result = contextToXml(context);
+    expect(result).toBe('<data-source name="test"></data-source>');
+  });
+});

package/src/core/ai/context/context.ts

@@ -3,6 +3,7 @@
 import { allTablesAtom } from "@/core/datasets/data-source-connections";
 import type { JotaiStore } from "@/core/state/jotai";
 import { variablesAtom } from "@/core/variables/state";
+import { ErrorContextProvider } from "./providers/error";
 import { TableContextProvider } from "./providers/tables";
 import { VariableContextProvider } from "./providers/variable";
 import { AIContextRegistry } from "./registry";
@@ -12,5 +13,6 @@ export function getAIContextRegistry(store: JotaiStore) {
   const variables = store.get(variablesAtom);
   return new AIContextRegistry()
     .register(new TableContextProvider(tablesMap))
-    .register(new VariableContextProvider(variables, tablesMap));
+    .register(new VariableContextProvider(variables, tablesMap))
+    .register(new ErrorContextProvider(store));
 }

package/src/core/ai/context/providers/__tests__/__snapshots__/error.test.ts.snap

@@ -0,0 +1,3 @@
+// Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html
+
+exports[`ErrorContextProvider > formatContext > should format context for basic errors > basic-error-context 1`] = `"<error name="Cell 1" description="Invalid syntax"></error>"`;

package/src/core/ai/context/providers/__tests__/__snapshots__/tables.test.ts.snap

@@ -1,38 +1,34 @@
 // Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html
 
 exports[`TableContextProvider > formatContext > should format context for basic table > basic-table-context 1`] = `
-"Table: products
-Source: memory
+"<data name="products" source="memory">
 Shape: 100 rows, 3 columns
 Columns:
   - id: integer
   - name: string
-  - active: boolean"
+  - active: boolean</data>"
 `;
 
 exports[`TableContextProvider > formatContext > should format context for remote database table > remote-table-context 1`] = `
-"Table: remote_table
-Source: postgresql://localhost:5432/mydb
+"<data name="remote_table" source="postgresql://localhost:5432/mydb">
 Shape: 100 rows, 3 columns
 Columns:
   - uuid: string
   - created_at: string
-  - metadata: string"
+  - metadata: string</data>"
 `;
 
 exports[`TableContextProvider > formatContext > should format context for table without columns > no-columns-table-context 1`] = `
-"Table: no_columns
-Source: memory
-Shape: 100 rows, 3 columns"
+"<data name="no_columns" source="memory">
+Shape: 100 rows, 3 columns</data>"
 `;
 
 exports[`TableContextProvider > formatContext > should format context for table without shape info > no-shape-table-context 1`] = `
-"Table: no_shape
-Source: memory
+"<data name="no_shape" source="memory">
 Columns:
   - id: integer
   - name: string
-  - active: boolean"
+  - active: boolean</data>"
 `;
 
 exports[`TableContextProvider > getItems > should handle dataframe tables with variable names > dataframe-table 1`] = `

package/src/core/ai/context/providers/__tests__/__snapshots__/variable.test.ts.snap

@@ -1,28 +1,12 @@
 // Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html
 
-exports[`VariableContextProvider > formatContext > should format context for basic variable > basic-variable-context 1`] = `
-"Variable: username
-Type: str
-Preview: "\\"alice\\"""
-`;
+exports[`VariableContextProvider > formatContext > should format context for basic variable > basic-variable-context 1`] = `"<variable name="username" dataType="str">"\\"alice\\""</variable>"`;
 
-exports[`VariableContextProvider > formatContext > should format context for dataframe variable > dataframe-variable-context 1`] = `
-"Variable: sales_df
-Type: pandas.DataFrame
-Preview: "<DataFrame shape: (1000, 8)>\\n   date    product  quantity  price\\n0  2023-01-01  Widget A    10    29.99\\n1  2023-01-02  Widget B     5    49.99\\n...""
-`;
+exports[`VariableContextProvider > formatContext > should format context for dataframe variable > dataframe-variable-context 1`] = `"<variable name="sales_df" dataType="pandas.DataFrame">"&lt;DataFrame shape: (1000, 8)&gt;\\n   date    product  quantity  price\\n0  2023-01-01  Widget A    10    29.99\\n1  2023-01-02  Widget B     5    49.99\\n..."</variable>"`;
 
-exports[`VariableContextProvider > formatContext > should format context for variable with complex value > complex-value-variable-context 1`] = `
-"Variable: complex_data
-Type: dict
-Preview: "{\\"users\\": [{\\"id\\": 1, \\"name\\": \\"Alice\\"}, {\\"id\\": 2, \\"name\\": \\"Bob\\"}], \\"total\\": 2}""
-`;
+exports[`VariableContextProvider > formatContext > should format context for variable with complex value > complex-value-variable-context 1`] = `"<variable name="complex_data" dataType="dict">"{\\"users\\": [{\\"id\\": 1, \\"name\\": \\"Alice\\"}, {\\"id\\": 2, \\"name\\": \\"Bob\\"}], \\"total\\": 2}"</variable>"`;
 
-exports[`VariableContextProvider > formatContext > should format context for variable without dataType > no-datatype-variable-context 1`] = `
-"Variable: mystery_var
-Type: unknown
-Preview: "some_value""
-`;
+exports[`VariableContextProvider > formatContext > should format context for variable without dataType > no-datatype-variable-context 1`] = `"<variable name="mystery_var" dataType="null">"some_value"</variable>"`;
 
 exports[`VariableContextProvider > getItems > should handle complex data types > complex-data-types 1`] = `
 [

package/src/core/ai/context/providers/__tests__/error.test.ts

@@ -0,0 +1,328 @@
+/* Copyright 2024 Marimo. All rights reserved. */
+/* eslint-disable @typescript-eslint/no-explicit-any */
+
+import { createStore } from "jotai";
+import { beforeEach, describe, expect, it } from "vitest";
+import { MockNotebook } from "@/__mocks__/notebook";
+import { notebookAtom } from "@/core/cells/cells";
+import { type CellId, CellId as CellIdClass } from "@/core/cells/ids";
+import { ErrorContextProvider } from "../error";
+
+describe("ErrorContextProvider", () => {
+  let provider: ErrorContextProvider;
+  let store: ReturnType<typeof createStore>;
+
+  beforeEach(() => {
+    store = createStore();
+    provider = new ErrorContextProvider(store);
+  });
+
+  const createMockNotebookWithErrors = (
+    errors: Parameters<typeof MockNotebook.notebookStateWithErrors>[0],
+  ) => {
+    const notebookState = MockNotebook.notebookStateWithErrors(errors);
+    store.set(notebookAtom, notebookState);
+  };
+
+  describe("provider properties", () => {
+    it("should have correct provider properties", () => {
+      expect(provider.title).toBe("Errors");
+      expect(provider.mentionPrefix).toBe("@");
+      expect(provider.contextType).toBe("error");
+    });
+  });
+
+  describe("getItems", () => {
+    it("should return empty array when no errors", () => {
+      const items = provider.getItems();
+      expect(items).toEqual([]);
+    });
+
+    it("should return error items when errors exist", () => {
+      const cellId1 = CellIdClass.create();
+      const cellId2 = CellIdClass.create();
+
+      createMockNotebookWithErrors([
+        {
+          cellId: cellId1,
+          cellName: "Cell 1",
+          errorData: [MockNotebook.errors.syntax("Invalid syntax")],
+        },
+        {
+          cellId: cellId2,
+          cellName: "Cell 2",
+          errorData: [MockNotebook.errors.exception("Runtime error")],
+        },
+      ]);
+
+      const items = provider.getItems();
+
+      expect(items).toHaveLength(1);
+      expect(items[0]).toMatchObject({
+        name: "Errors",
+        type: "error",
+        description: "All errors in the notebook",
+        data: {
+          type: "all-errors",
+          errors: [
+            {
+              cellId: cellId1,
+              cellName: "Cell 1",
+              errorData: [{ type: "syntax", msg: "Invalid syntax" }],
+            },
+            {
+              cellId: cellId2,
+              cellName: "Cell 2",
+              errorData: [{ type: "exception", msg: "Runtime error" }],
+            },
+          ],
+        },
+      });
+    });
+
+    it("should handle cells without names", () => {
+      const cellId = CellIdClass.create();
+
+      createMockNotebookWithErrors([
+        {
+          cellId,
+          cellName: "",
+          errorData: [MockNotebook.errors.syntax("Invalid syntax")],
+        },
+      ]);
+
+      const items = provider.getItems();
+      expect(items[0].data.errors[0].cellName).toBe("");
+    });
+  });
+
+  describe("formatCompletion", () => {
+    it("should format completion for all-errors type", () => {
+      const cellId1 = CellIdClass.create();
+      const cellId2 = CellIdClass.create();
+
+      createMockNotebookWithErrors([
+        {
+          cellId: cellId1,
+          cellName: "Cell 1",
+          errorData: [MockNotebook.errors.syntax("Invalid syntax")],
+        },
+        {
+          cellId: cellId2,
+          cellName: "Cell 2",
+          errorData: [MockNotebook.errors.exception("Runtime error")],
+        },
+      ]);
+
+      const items = provider.getItems();
+      const completion = provider.formatCompletion(items[0]);
+
+      expect(completion).toMatchObject({
+        label: "@Errors",
+        displayLabel: "Errors",
+        detail: "2 errors",
+        boost: 2, // Boosts.ERROR
+        type: "error",
+        apply: "@Errors",
+        section: "Errors",
+      });
+
+      // Test the info function
+      expect(completion.info).toBeDefined();
+      if (typeof completion.info === "function") {
+        const infoResult = completion.info(completion);
+        if (
+          infoResult &&
+          typeof infoResult === "object" &&
+          "dom" in infoResult
+        ) {
+          const infoElement = infoResult.dom as HTMLElement;
+          expect(infoElement.tagName).toBe("DIV");
+          expect(infoElement.textContent).toContain("Errors");
+          expect(infoElement.textContent).toContain("2 errors");
+        } else if (infoResult && "tagName" in (infoResult as any)) {
+          const infoElement = infoResult as HTMLElement;
+          expect(infoElement.tagName).toBe("DIV");
+          expect(infoElement.textContent).toContain("Errors");
+          expect(infoElement.textContent).toContain("2 errors");
+        }
+      }
+    });
+
+    it("should handle single error correctly", () => {
+      const cellId = CellIdClass.create();
+
+      createMockNotebookWithErrors([
+        {
+          cellId,
+          cellName: "Cell 1",
+          errorData: [MockNotebook.errors.syntax("Invalid syntax")],
+        },
+      ]);
+
+      const items = provider.getItems();
+      const completion = provider.formatCompletion(items[0]);
+      expect(completion.detail).toBe("1 error");
+    });
+
+    it("should handle fallback for unknown error types", () => {
+      const item = {
+        uri: "error://unknown",
+        name: "Unknown Error",
+        type: "error" as const,
+        data: {
+          type: "unknown" as any,
+          errors: [],
+        },
+        description: "Unknown error type",
+      };
+
+      const completion = provider.formatCompletion(item);
+      expect(completion).toMatchObject({
+        label: "Error",
+        displayLabel: "Error",
+        boost: 2, // Boosts.ERROR
+      });
+    });
+  });
+
+  describe("formatContext", () => {
+    it("should format context for basic errors", () => {
+      const cellId = CellIdClass.create();
+
+      createMockNotebookWithErrors([
+        {
+          cellId,
+          cellName: "Cell 1",
+          errorData: [MockNotebook.errors.syntax("Invalid syntax")],
+        },
+      ]);
+
+      const items = provider.getItems();
+      const context = provider.formatContext(items[0]);
+      expect(context).toMatchSnapshot("basic-error-context");
+    });
+
+    it("should format context for multiple error types", () => {
+      const cellId1 = "cell-1" as CellId;
+      const cellId2 = "cell-2" as CellId;
+
+      createMockNotebookWithErrors([
+        {
+          cellId: cellId1,
+          cellName: "Cell 1",
+          errorData: [
+            MockNotebook.errors.syntax("Invalid syntax"),
+            MockNotebook.errors.exception("Runtime error"),
+          ],
+        },
+        {
+          cellId: cellId2,
+          cellName: "",
+          errorData: [
+            MockNotebook.errors.cycle(),
+            MockNotebook.errors.multipleDefs("variable_x"),
+          ],
+        },
+      ]);
+
+      const items = provider.getItems();
+      const context = provider.formatContext(items[0]);
+
+      // Check for expected content instead of exact snapshot match due to random cell IDs
+      expect(context).toMatchInlineSnapshot(`
+        "<error name="Cell 1" description="Invalid syntax
+        Runtime error"></error>
+
+        <error name="Cell cell-2" description="This cell is in a cycle
+        The variable 'variable_x' was defined by another cell"></error>"
+      `);
+    });
+
+    it("should handle cells without names", () => {
+      const cellId = CellIdClass.create();
+
+      createMockNotebookWithErrors([
+        {
+          cellId,
+          cellName: "",
+          errorData: [MockNotebook.errors.syntax("Invalid syntax")],
+        },
+      ]);
+
+      const items = provider.getItems();
+      const context = provider.formatContext(items[0]);
+      expect(context).toContain(`Cell ${cellId}`);
+    });
+  });
+
+  describe("error description handling", () => {
+    it("should handle all error types correctly", () => {
+      // Note: ancestor-* error types are intentionally filtered out by cellErrorsAtom
+      const errorTypes = [
+        {
+          error: MockNotebook.errors.setupRefs(),
+          expected: "The setup cell cannot have references",
+        },
+        {
+          error: MockNotebook.errors.cycle(),
+          expected: "This cell is in a cycle",
+        },
+        {
+          error: MockNotebook.errors.multipleDefs("var_x"),
+          expected: "The variable 'var_x' was defined by another cell",
+        },
+        {
+          error: MockNotebook.errors.importStar("Import star error"),
+          expected: "Import star error",
+        },
+        {
+          error: MockNotebook.errors.exception("Runtime exception"),
+          expected: "Runtime exception",
+        },
+        {
+          error: MockNotebook.errors.strictException("Strict exception", "ref"),
+          expected: "Strict exception",
+        },
+        {
+          error: MockNotebook.errors.interruption(),
+          expected: "This cell was interrupted and needs to be re-run",
+        },
+        {
+          error: MockNotebook.errors.syntax("Syntax error"),
+          expected: "Syntax error",
+        },
+        {
+          error: MockNotebook.errors.unknown("Unknown error"),
+          expected: "Unknown error",
+        },
+      ];
+
+      for (const { error, expected } of errorTypes) {
+        // Create a fresh store and provider for each error type
+        const testStore = createStore();
+        const testProvider = new ErrorContextProvider(testStore);
+
+        const { notebookState } = MockNotebook.cellWithErrors("Cell 1", [
+          error,
+        ]);
+        testStore.set(notebookAtom, notebookState);
+
+        const items = testProvider.getItems();
+        expect(items).toHaveLength(1);
+        const context = testProvider.formatContext(items[0]);
+        expect(context).toContain(expected);
+      }
+    });
+
+    it("should handle unknown error type gracefully", () => {
+      const error = { type: "unknown-type" as any };
+      const { notebookState } = MockNotebook.cellWithErrors("Cell 1", [error]);
+      store.set(notebookAtom, notebookState);
+
+      const items = provider.getItems();
+      const context = provider.formatContext(items[0]);
+      expect(context).toContain("Unknown error");
+    });
+  });
+});

package/src/core/ai/context/providers/common.ts

@@ -4,4 +4,5 @@ export const Boosts = {
   LOCAL_TABLE: 5,
   REMOTE_TABLE: 4,
   VARIABLE: 3,
+  ERROR: 2,
 } as const;