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

package/src/components/editor/Cell.tsx

@@ -2,8 +2,7 @@
 import { closeCompletion, completionStatus } from "@codemirror/autocomplete";
 import type { EditorView } from "@codemirror/view";
 import clsx from "clsx";
-import { useAtom, useAtomValue, useSetAtom } from "jotai";
-import { ScopeProvider } from "jotai-scope";
+import { useAtomValue, useSetAtom } from "jotai";
 import {
   HelpCircleIcon,
   MoreHorizontalIcon,
@@ -78,7 +77,10 @@ import { useRunCell } from "./cell/useRunCells";
 import { HideCodeButton } from "./code/readonly-python-code";
 import { cellDomProps } from "./common";
 import { useCellNavigationProps } from "./navigation/navigation";
-import { temporarilyShownCodeAtom } from "./navigation/state";
+import {
+  useTemporarilyShownCode,
+  useTemporarilyShownCodeActions,
+} from "./navigation/state";
 import { OutputArea } from "./Output";
 import { ConsoleOutput } from "./output/ConsoleOutput";
 import { CellDragHandle, SortableCell } from "./SortableCell";
@@ -157,9 +159,8 @@ function useCellHiddenLogic({
   editorView: React.RefObject<EditorView | null>;
   editorViewParentRef: React.RefObject<HTMLDivElement | null>;
 }) {
-  const [temporarilyVisible, setTemporarilyVisible] = useAtom(
-    temporarilyShownCodeAtom,
-  );
+  const temporarilyVisible = useTemporarilyShownCode(cellId);
+  const temporarilyShownCodeActions = useTemporarilyShownCodeActions();
   const isUntouched = useAtomValue(
     useMemo(() => createUntouchedCellAtom(cellId), [cellId]),
   );
@@ -179,7 +180,7 @@ function useCellHiddenLogic({
 
     // Default to true
     const focus = opts?.focus ?? true;
-    setTemporarilyVisible(true);
+    temporarilyShownCodeActions.add(cellId);
 
     if (focus) {
       editorView.current?.focus();
@@ -297,16 +298,14 @@ const CellComponent = (props: CellProps) => {
 
   if (mode === "edit") {
     return (
-      <ScopeProvider atoms={[temporarilyShownCodeAtom]}>
-        <EditableCellComponent
-          {...props}
-          cellId={cellId}
-          editorView={editorView}
-          setEditorView={(ev) => {
-            editorView.current = ev;
-          }}
-        />
-      </ScopeProvider>
+      <EditableCellComponent
+        {...props}
+        cellId={cellId}
+        editorView={editorView}
+        setEditorView={(ev) => {
+          editorView.current = ev;
+        }}
+      />
     );
   }
 

package/src/components/editor/ai/__tests__/completion-utils.test.ts

@@ -52,41 +52,40 @@ describe("getAICompletionBody", () => {
     const input = "Use @data://dataset1 and @data://dataset2 for analysis";
     const result = getAICompletionBody({ input });
 
-    expect(result).toEqual({
-      includeOtherCode: "// Some other code",
-      context: {
-        schema: [
-          {
-            name: "dataset1",
-            columns: [
-              { name: "col1", type: "number" },
-              { name: "col2", type: "string" },
-            ],
-          },
-          {
-            name: "dataset2",
-            columns: [
-              { name: "col3", type: "boolean" },
-              { name: "col4", type: "date" },
-            ],
-          },
-        ],
-        variables: [],
-      },
-    });
+    expect(result).toMatchInlineSnapshot(`
+      {
+        "context": {
+          "plainText": "<data name="dataset1" source="unknown">
+      Columns:
+        - col1: number
+        - col2: string</data>
+
+      <data name="dataset2" source="unknown">
+      Columns:
+        - col3: boolean
+        - col4: date</data>",
+          "schema": [],
+          "variables": [],
+        },
+        "includeOtherCode": "// Some other code",
+      }
+    `);
   });
 
   it("should handle input with no mentioned datasets", () => {
     const input = "Perform some analysis without mentioning @data://datasets";
     const result = getAICompletionBody({ input });
 
-    expect(result).toEqual({
-      includeOtherCode: "// Some other code",
-      context: {
-        schema: [],
-        variables: [],
-      },
-    });
+    expect(result).toMatchInlineSnapshot(`
+      {
+        "context": {
+          "plainText": "",
+          "schema": [],
+          "variables": [],
+        },
+        "includeOtherCode": "// Some other code",
+      }
+    `);
   });
 
   it("should handle input with non-existent datasets", () => {
@@ -106,21 +105,19 @@ describe("getAICompletionBody", () => {
       "Use @data://existingDataset and @data://nonExistentDataset for analysis";
     const result = getAICompletionBody({ input });
 
-    expect(result).toEqual({
-      includeOtherCode: "// Some other code",
-      context: {
-        schema: [
-          {
-            name: "existingDataset",
-            columns: [
-              { name: "col1", type: "number" },
-              { name: "col2", type: "string" },
-            ],
-          },
-        ],
-        variables: [],
-      },
-    });
+    expect(result).toMatchInlineSnapshot(`
+      {
+        "context": {
+          "plainText": "<data name="existingDataset" source="unknown">
+      Columns:
+        - col1: number
+        - col2: string</data>",
+          "schema": [],
+          "variables": [],
+        },
+        "includeOtherCode": "// Some other code",
+      }
+    `);
   });
 
   it("should handle dataset names with dots", () => {
@@ -144,25 +141,23 @@ describe("getAICompletionBody", () => {
       "Use @data://dataset.with.dots and @data://regular_dataset for analysis";
     const result = getAICompletionBody({ input });
 
-    expect(result).toEqual({
-      includeOtherCode: "// Some other code",
-      context: {
-        schema: [
-          {
-            name: "dataset.with.dots",
-            columns: [
-              { name: "col1", type: "number" },
-              { name: "col2", type: "string" },
-            ],
-          },
-          {
-            name: "regular_dataset",
-            columns: [{ name: "col3", type: "boolean" }],
-          },
-        ],
-        variables: [],
-      },
-    });
+    expect(result).toMatchInlineSnapshot(`
+      {
+        "context": {
+          "plainText": "<data name="dataset.with.dots" source="unknown">
+      Columns:
+        - col1: number
+        - col2: string</data>
+
+      <data name="regular_dataset" source="unknown">
+      Columns:
+        - col3: boolean</data>",
+          "schema": [],
+          "variables": [],
+        },
+        "includeOtherCode": "// Some other code",
+      }
+    `);
   });
 
   it("should handle connections", () => {
@@ -200,18 +195,18 @@ describe("getAICompletionBody", () => {
     const input = "Use @data://table1 for analysis";
     const result = getAICompletionBody({ input });
 
-    expect(result).toEqual({
-      includeOtherCode: "// Some other code",
-      context: {
-        schema: [
-          {
-            name: "table1",
-            columns: [{ name: "col1", type: "number" }],
-          },
-        ],
-        variables: [],
-      },
-    });
+    expect(result).toMatchInlineSnapshot(`
+      {
+        "context": {
+          "plainText": "<data name="table1" source="unknown">
+      Columns:
+        - col1: number</data>",
+          "schema": [],
+          "variables": [],
+        },
+        "includeOtherCode": "// Some other code",
+      }
+    `);
   });
 
   it("should return the correct completion body with mentioned variables", () => {
@@ -237,24 +232,18 @@ describe("getAICompletionBody", () => {
     const input = "Use @variable://var1 and @variable://var2 for analysis";
     const result = getAICompletionBody({ input });
 
-    expect(result).toEqual({
-      includeOtherCode: "// Some other code",
-      context: {
-        schema: [],
-        variables: [
-          {
-            name: "var1",
-            valueType: "string",
-            previewValue: "string value",
-          },
-          {
-            name: "var2",
-            valueType: "number",
-            previewValue: "42",
-          },
-        ],
-      },
-    });
+    expect(result).toMatchInlineSnapshot(`
+      {
+        "context": {
+          "plainText": "<variable name="var1" dataType="string">"string value"</variable>
+
+      <variable name="var2" dataType="number">"42"</variable>",
+          "schema": [],
+          "variables": [],
+        },
+        "includeOtherCode": "// Some other code",
+      }
+    `);
   });
 
   it("should handle input with both datasets and variables", () => {
@@ -284,27 +273,21 @@ describe("getAICompletionBody", () => {
     const input = "Use @data://dataset1 and @variable://var1 for analysis";
     const result = getAICompletionBody({ input });
 
-    expect(result).toEqual({
-      includeOtherCode: "// Some other code",
-      context: {
-        schema: [
-          {
-            name: "dataset1",
-            columns: [
-              { name: "col1", type: "number" },
-              { name: "col2", type: "string" },
-            ],
-          },
-        ],
-        variables: [
-          {
-            name: "var1",
-            valueType: "string",
-            previewValue: "string value",
-          },
-        ],
-      },
-    });
+    expect(result).toMatchInlineSnapshot(`
+      {
+        "context": {
+          "plainText": "<data name="dataset1" source="unknown">
+      Columns:
+        - col1: number
+        - col2: string</data>
+
+      <variable name="var1" dataType="string">"string value"</variable>",
+          "schema": [],
+          "variables": [],
+        },
+        "includeOtherCode": "// Some other code",
+      }
+    `);
   });
 
   it("should handle non-existent variables", () => {
@@ -324,19 +307,16 @@ describe("getAICompletionBody", () => {
       "Use @variable://existingVar and @variable://nonExistentVar for analysis";
     const result = getAICompletionBody({ input });
 
-    expect(result).toEqual({
-      includeOtherCode: "// Some other code",
-      context: {
-        schema: [],
-        variables: [
-          {
-            name: "existingVar",
-            valueType: "string",
-            previewValue: "string value",
-          },
-        ],
-      },
-    });
+    expect(result).toMatchInlineSnapshot(`
+      {
+        "context": {
+          "plainText": "<variable name="existingVar" dataType="string">"string value"</variable>",
+          "schema": [],
+          "variables": [],
+        },
+        "includeOtherCode": "// Some other code",
+      }
+    `);
   });
 
   it("should prioritize datasets over variables when there's a name conflict", () => {
@@ -363,17 +343,17 @@ describe("getAICompletionBody", () => {
     const input = "Use @data://conflict for analysis";
     const result = getAICompletionBody({ input });
 
-    expect(result).toEqual({
-      includeOtherCode: "// Some other code",
-      context: {
-        schema: [
-          {
-            name: "conflict",
-            columns: [{ name: "col1", type: "number" }],
-          },
-        ],
-        variables: [],
-      },
-    });
+    expect(result).toMatchInlineSnapshot(`
+      {
+        "context": {
+          "plainText": "<data name="conflict" source="unknown">
+      Columns:
+        - col1: number</data>",
+          "schema": [],
+          "variables": [],
+        },
+        "includeOtherCode": "// Some other code",
+      }
+    `);
   });
 });

package/src/components/editor/ai/completion-utils.ts

@@ -6,13 +6,9 @@ import type {
   CompletionSource,
 } from "@codemirror/autocomplete";
 import { getAIContextRegistry } from "@/core/ai/context/context";
-import type { TableContextItem } from "@/core/ai/context/providers/tables";
-import type { VariableContextItem } from "@/core/ai/context/providers/variable";
 import { getCodes } from "@/core/codemirror/copilot/getCodes";
-import type { DataTable } from "@/core/kernel/messages";
 import type { AiCompletionRequest } from "@/core/network/types";
 import { store } from "@/core/state/jotai";
-import type { Variable } from "@/core/variables/types";
 import { Logger } from "@/utils/Logger";
 
 /**
@@ -23,26 +19,15 @@ export function getAICompletionBody({
 }: {
   input: string;
 }): Omit<AiCompletionRequest, "language" | "prompt" | "code"> {
-  const { datasets, variables } = extractDatasetsAndVariables(input);
-  Logger.debug("Included datasets", datasets);
-  Logger.debug("Included variables", variables);
+  const contextString = extractDatasetsAndVariables(input);
+  Logger.debug("Included context", contextString);
 
   return {
     includeOtherCode: getCodes(""),
     context: {
-      schema: datasets.map((dataset) => ({
-        name: dataset.name,
-        columns: dataset.columns.map((column) => ({
-          name: column.name,
-          type: column.type,
-          sampleValues: column.sample_values,
-        })),
-      })),
-      variables: variables.map((variable) => ({
-        name: variable.name,
-        valueType: variable.dataType ?? "",
-        previewValue: variable.value,
-      })),
+      plainText: contextString,
+      schema: [],
+      variables: [],
     },
   };
 }
@@ -52,21 +37,10 @@ export function getAICompletionBody({
  * References are with @<name> in the input.
  * Prioritizes datasets over variables if there's a name conflict.
  */
-function extractDatasetsAndVariables(input: string): {
-  datasets: DataTable[];
-  variables: Variable[];
-} {
+function extractDatasetsAndVariables(input: string): string {
   const registry = getAIContextRegistry(store);
   const contextIds = registry.parseAllContextIds(input);
-  const contextInfo = registry.getContextInfo(contextIds);
-  const datasets: DataTable[] = contextInfo
-    .filter((info): info is TableContextItem => info.type === "data")
-    .map((info) => info.data);
-  const variables: Variable[] = contextInfo
-    .filter((info): info is VariableContextItem => info.type === "variable")
-    .map((info) => info.data.variable);
-
-  return { datasets, variables };
+  return registry.formatContextForAI(contextIds);
 }
 
 /**

package/src/components/editor/navigation/navigation.ts

@@ -33,7 +33,7 @@ import {
   useCellSelectionActions,
   useIsCellSelected,
 } from "./selection";
-import { temporarilyShownCodeAtom } from "./state";
+import { useTemporarilyShownCodeActions } from "./state";
 import { handleVimKeybinding } from "./vim-bindings";
 
 interface HotkeyHandler {
@@ -85,7 +85,7 @@ function addSingleHandler(handler: HotkeyHandler["bulkHandle"]): HotkeyHandler {
 function useCellFocusProps(cellId: CellId) {
   const focusActions = useCellFocusActions();
   const actions = useCellActions();
-  const setTemporarilyShownCode = useSetAtom(temporarilyShownCodeAtom);
+  const temporarilyShownCodeActions = useTemporarilyShownCodeActions();
 
   // This occurs at the cell level and descedants.
   const { focusWithinProps } = useFocusWithin({
@@ -95,7 +95,7 @@ function useCellFocusProps(cellId: CellId) {
     },
     onBlurWithin: () => {
       // On blur, hide the code if it was temporarily shown.
-      setTemporarilyShownCode(false);
+      temporarilyShownCodeActions.remove(cellId);
       actions.markTouched({ cellId });
       focusActions.blurCell();
     },
@@ -132,7 +132,7 @@ export function useCellNavigationProps(
   const setAiCompletionCell = useSetAtom(aiCompletionCellAtom);
   const actions = useCellActions();
   const store = useStore();
-  const setTemporarilyShownCode = useSetAtom(temporarilyShownCodeAtom);
+  const temporarilyShownCodeActions = useTemporarilyShownCodeActions();
   const runCells = useRunCells();
   const keymapPreset = useAtomValue(keymapPresetAtom);
   const { copyCells, pasteAtCell } = useCellClipboard();
@@ -275,7 +275,7 @@ export function useCellNavigationProps(
         },
         // Enter will focus the cell editor.
         Enter: () => {
-          setTemporarilyShownCode(true);
+          temporarilyShownCodeActions.add(cellId);
           focusCellEditor(store, cellId);
           selectionActions.clear();
           return true;
@@ -608,7 +608,7 @@ export function useCellEditorNavigationProps(
   cellId: CellId,
   editorView: React.RefObject<EditorView | null>,
 ) {
-  const setTemporarilyShownCode = useSetAtom(temporarilyShownCodeAtom);
+  const temporarilyShownCodeActions = useTemporarilyShownCodeActions();
   const keymapPreset = useAtomValue(keymapPresetAtom);
   const hotkeys = useAtomValue(hotkeysAtom);
 
@@ -618,7 +618,7 @@ export function useCellEditorNavigationProps(
   }, [hotkeys]);
 
   const exitToCommandMode = () => {
-    setTemporarilyShownCode(false);
+    temporarilyShownCodeActions.remove(cellId);
     focusCell(cellId);
   };
 

package/src/components/editor/navigation/state.ts

@@ -1,5 +1,37 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 
-import { atom } from "jotai";
+import { atom, useAtomValue } from "jotai";
+import { useMemo } from "react";
+import type { CellId } from "@/core/cells/ids";
+import { createReducerAndAtoms } from "@/utils/createReducer";
 
-export const temporarilyShownCodeAtom = atom<boolean>(false);
+type TemporarilyShownCodeState = Set<CellId>;
+
+// This previously used jotai-scope, but unfortunately this bug creates unnecessary work:
+// https://github.com/jotaijs/jotai-scope/issues/25
+
+const {
+  valueAtom: temporarilyShownCodeAtom,
+  useActions: useTemporarilyShownCodeActions,
+} = createReducerAndAtoms(() => new Set<CellId>(), {
+  add: (state: TemporarilyShownCodeState, cellId: CellId) => {
+    const newState = new Set(state);
+    newState.add(cellId);
+    return newState;
+  },
+  remove: (state: TemporarilyShownCodeState, cellId: CellId) => {
+    const newState = new Set(state);
+    newState.delete(cellId);
+    return newState;
+  },
+});
+
+const createTemporarilyShownCodeAtom = (cellId: CellId) =>
+  atom((get) => get(temporarilyShownCodeAtom).has(cellId));
+
+export function useTemporarilyShownCode(cellId: CellId) {
+  const atom = useMemo(() => createTemporarilyShownCodeAtom(cellId), [cellId]);
+  return useAtomValue(atom);
+}
+
+export { temporarilyShownCodeAtom, useTemporarilyShownCodeActions };