@marimo-team/frontend 0.15.1-dev34 → 0.15.1-dev36

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

@@ -1,7 +1,15 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 
 import type { Completion } from "@codemirror/autocomplete";
+import { createStore } from "jotai";
 import { beforeEach, describe, expect, it } from "vitest";
+import { MockNotebook } from "@/__mocks__/notebook";
+import { notebookAtom } from "@/core/cells/cells";
+import { CellId as CellIdClass } from "@/core/cells/ids";
+import {
+  type ErrorContextItem,
+  ErrorContextProvider,
+} from "../providers/error";
 import {
   type AIContextItem,
   AIContextProvider,
@@ -513,4 +521,150 @@ describe("AIContextRegistry", () => {
       expect(provider).toBe(mockProvider);
     });
   });
+
+  describe("Integration with ErrorContextProvider", () => {
+    let errorProvider: ErrorContextProvider;
+    let store: ReturnType<typeof createStore>;
+    let registry: AIContextRegistry<ErrorContextItem>;
+
+    beforeEach(() => {
+      store = createStore();
+
+      // Create mock notebook with errors using the new MockNotebook utilities
+      const cellId1 = CellIdClass.create();
+      const cellId2 = CellIdClass.create();
+
+      const notebookState = MockNotebook.notebookStateWithErrors([
+        {
+          cellId: cellId1,
+          cellName: "Cell 1",
+          errorData: [
+            MockNotebook.errors.syntax("Invalid syntax"),
+            MockNotebook.errors.exception("Runtime error"),
+          ],
+        },
+        {
+          cellId: cellId2,
+          cellName: "Cell 2",
+          errorData: [MockNotebook.errors.cycle()],
+        },
+      ]);
+
+      store.set(notebookAtom, notebookState);
+      errorProvider = new ErrorContextProvider(store);
+      registry = new AIContextRegistry<ErrorContextItem>().register(
+        errorProvider,
+      );
+    });
+
+    it("should register and provide error context items", () => {
+      const provider = registry.getProvider("error");
+      expect(provider).toBe(errorProvider);
+
+      const items = registry.getAllItems();
+      expect(items).toHaveLength(1);
+      expect(items[0].type).toBe("error");
+      expect(items[0].name).toBe("Errors");
+    });
+
+    it("should parse error context IDs correctly", () => {
+      const contextIds = registry.parseAllContextIds(
+        "Use @error://all to analyze errors",
+      );
+      expect(contextIds).toEqual(["error://all"]);
+
+      const contextInfo = registry.getContextInfo(contextIds);
+      expect(contextInfo).toHaveLength(1);
+      expect(contextInfo[0].type).toBe("error");
+    });
+
+    it("should format error context for AI", () => {
+      const contextIds = ["error://all"] as ContextLocatorId[];
+      const formattedContext = registry.formatContextForAI(contextIds);
+
+      expect(formattedContext).toContain("<error");
+      expect(formattedContext).toContain("Cell 1");
+      expect(formattedContext).toContain("Cell 2");
+      expect(formattedContext).toContain("Invalid syntax");
+      expect(formattedContext).toContain("This cell is in a cycle");
+    });
+
+    it("should provide error completions", () => {
+      // Get the error provider and test its completion directly
+      const provider = registry.getProvider("error");
+      expect(provider).toBeDefined();
+
+      const items = provider!.getItems();
+      expect(items).toHaveLength(1);
+
+      const completion = provider!.formatCompletion(items[0]);
+      expect(completion.label).toBe("@Errors");
+      expect(completion.type).toBe("error");
+    });
+
+    it("should handle empty errors gracefully", () => {
+      // Create store with no errors using MockNotebook
+      const emptyStore = createStore();
+      const emptyNotebookState = MockNotebook.notebookStateWithErrors([]);
+      emptyStore.set(notebookAtom, emptyNotebookState);
+
+      const emptyErrorProvider = new ErrorContextProvider(emptyStore);
+      const emptyRegistry = new AIContextRegistry().register(
+        emptyErrorProvider,
+      );
+
+      const items = emptyRegistry.getAllItems();
+      expect(items).toHaveLength(0);
+    });
+
+    it("should work with mixed providers", () => {
+      const mixedRegistry = new AIContextRegistry<
+        ErrorContextItem | MockContextItem
+      >()
+        .register(errorProvider)
+        .register(mockProvider);
+
+      const items = mixedRegistry.getAllItems();
+      expect(items.length).toBeGreaterThan(1);
+
+      // Should have both error and mock items
+      const errorItems = items.filter((item) => item.type === "error");
+      const mockItems = items.filter((item) => item.type === "mock");
+
+      expect(errorItems).toHaveLength(1);
+      expect(mockItems).toHaveLength(3);
+    });
+
+    it("should parse mixed context IDs", () => {
+      const mixedRegistry = new AIContextRegistry<
+        ErrorContextItem | MockContextItem
+      >()
+        .register(errorProvider)
+        .register(mockProvider);
+
+      const input = "Use @error://all and @mock://item1 for analysis";
+      const contextIds = mixedRegistry.parseAllContextIds(input);
+
+      expect(contextIds).toContain("error://all");
+      expect(contextIds).toContain("mock://item1");
+      expect(contextIds).toHaveLength(2);
+    });
+
+    it("should format mixed context for AI", () => {
+      const mixedRegistry = new AIContextRegistry<
+        ErrorContextItem | MockContextItem
+      >()
+        .register(errorProvider)
+        .register(mockProvider);
+
+      const contextIds = ["error://all", "mock://item1"] as ContextLocatorId[];
+      const formattedContext = mixedRegistry.formatContextForAI(contextIds);
+
+      // Should contain both error and mock context
+      expect(formattedContext).toContain("<error");
+      expect(formattedContext).toContain("Mock:");
+      expect(formattedContext).toContain("Cell 1");
+      expect(formattedContext).toContain("Item 1");
+    });
+  });
 });

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`] = `
 [