add maybe import marimo to Generate with AI

frontend/src/core/ai/staged-cells.ts

@@ -10,7 +10,12 @@ import { useDeleteCellCallback } from "@/components/editor/cell/useDeleteCell";
 import { CellId } from "@/core/cells/ids";
 import { createReducerAndAtoms } from "@/utils/createReducer";
 import { Logger } from "@/utils/Logger";
-import { cellHandleAtom, useCellActions } from "../cells/cells";
+import { maybeAddMarimoImport } from "../cells/add-missing-import";
+import {
+  type CreateNewCellAction,
+  cellHandleAtom,
+  useCellActions,
+} from "../cells/cells";
 import type { LanguageAdapterType } from "../codemirror/language/types";
 import { updateEditorCodeFromPython } from "../codemirror/language/utils";
 import type { JotaiStore } from "../state/jotai";
@@ -113,6 +118,8 @@ export function useStagedCells(store: JotaiStore) {
         cellCreationStream.current = new CellCreationStream(
           createStagedCell,
           updateStagedCell,
+          addStagedCell,
+          createNewCell,
         );
         break;
       case "text-delta":
@@ -168,13 +175,20 @@ class CellCreationStream {
 
   private onCreateCell: (code: string) => CellId;
   private onUpdateCell: (opts: UpdateStagedCellAction) => void;
+  private addStagedCell: (payload: { cellId: CellId }) => void;
+  private createNewCell: (opts: CreateNewCellAction) => void;
+  private hasMarimoImport = false;
 
   constructor(
     onCreateCell: (code: string) => CellId,
     onUpdateCell: (opts: UpdateStagedCellAction) => void,
+    addStagedCell: (payload: { cellId: CellId }) => void,
+    createNewCell: (opts: CreateNewCellAction) => void,
   ) {
     this.onCreateCell = onCreateCell;
     this.onUpdateCell = onUpdateCell;
+    this.addStagedCell = addStagedCell;
+    this.createNewCell = createNewCell;
   }
 
   stream(chunk: TextDeltaChunk) {
@@ -187,6 +201,7 @@ class CellCreationStream {
     // For each parsed cell, we either update an existing staged cell or create a new one.
     for (const [idx, cell] of completionCells.entries()) {
       if (idx < this.createdCells.length) {
+        this.addMarimoImport(cell.language);
         const existingCell = this.createdCells[idx];
         this.createdCells[idx] = { ...existingCell, cell };
         this.onUpdateCell({
@@ -201,6 +216,24 @@ class CellCreationStream {
     }
   }
 
+  /** Add a marimo import if the cell is SQL or Markdown and we haven't added it yet. */
+  private addMarimoImport(language: LanguageAdapterType) {
+    if (this.hasMarimoImport || language === "python") {
+      return;
+    }
+
+    const cellId = maybeAddMarimoImport({
+      autoInstantiate: false,
+      createNewCell: this.createNewCell,
+      fromCellId: this.createdCells[0]?.cellId,
+      before: true,
+    });
+    if (cellId) {
+      this.addStagedCell({ cellId });
+    }
+    this.hasMarimoImport = true;
+  }
+
   stop() {
     // Clear all state
     this.buffer = "";

frontend/src/core/cells/add-missing-import.ts

@@ -47,24 +47,36 @@ export function maybeAddMissingImport({
   return true;
 }
 
+/**
+ * Adds a marimo import to the notebook if not already present.
+ * @param autoInstantiate Whether to automatically run the cell.
+ * @param createNewCell The function to create a new cell.
+ * @param fromCellId The cell to add the import to.
+ * @param before Whether to add the import before or after the cell.
+ *
+ * Returns the ID of the new cell if added, otherwise null.
+ */
 export function maybeAddMarimoImport({
   autoInstantiate,
   createNewCell,
   fromCellId,
+  before,
 }: {
   autoInstantiate: boolean;
   createNewCell: CellActions["createNewCell"];
   fromCellId?: CellId | null;
-}): boolean {
+  before?: boolean;
+}): CellId | null {
   const client = getRequestClient();
-  return maybeAddMissingImport({
+  let newCellId: CellId | null = null;
+  const added = maybeAddMissingImport({
     moduleName: "marimo",
     variableName: "mo",
     onAddImport: (importStatement) => {
-      const newCellId = CellId.create();
+      newCellId = CellId.create();
       createNewCell({
         cellId: fromCellId ?? "__end__",
-        before: false,
+        before: before ?? false,
         code: importStatement,
         lastCodeRun: autoInstantiate ? importStatement : undefined,
         newCellId: newCellId,
@@ -78,6 +90,7 @@ export function maybeAddMarimoImport({
       }
     },
   });
+  return added ? newCellId : null;
 }
 
 export function maybeAddAltairImport({

frontend/src/core/cells/cells.ts

@@ -145,6 +145,31 @@ export function initialNotebookState(): NotebookState {
   });
 }
 
+export interface CreateNewCellAction {
+  /** The target cell ID to create a new cell relative to. Can be:
+   * - A CellId string for an existing cell
+   * - "__end__" to append at the end of the first column
+   * - {type: "__end__", columnId} to append at the end of a specific column
+   */
+  cellId: CellId | "__end__" | { type: "__end__"; columnId: CellColumnId };
+  /** Whether to insert before (true) or after (false) the target cell */
+  before: boolean;
+  /** Initial code content for the new cell */
+  code?: string;
+  /** The last executed code for the new cell */
+  lastCodeRun?: string;
+  /** Timestamp of the last execution */
+  lastExecutionTime?: number;
+  /** Optional custom ID for the new cell. Auto-generated if not provided */
+  newCellId?: CellId;
+  /** Whether to focus the new cell after creation */
+  autoFocus?: boolean;
+  /** If true, skip creation if code already exists */
+  skipIfCodeExists?: boolean;
+  /** Hide the code in the new cell. This will be initially shown until the cell is blurred for the first time. */
+  hideCode?: boolean;
+}
+
 /**
  * Actions and reducer for the notebook state.
  */
@@ -154,33 +179,7 @@ const {
   useActions,
   valueAtom: notebookAtom,
 } = createReducerAndAtoms(initialNotebookState, {
-  createNewCell: (
-    state,
-    action: {
-      /** The target cell ID to create a new cell relative to. Can be:
-       * - A CellId string for an existing cell
-       * - "__end__" to append at the end of the first column
-       * - {type: "__end__", columnId} to append at the end of a specific column
-       */
-      cellId: CellId | "__end__" | { type: "__end__"; columnId: CellColumnId };
-      /** Whether to insert before (true) or after (false) the target cell */
-      before: boolean;
-      /** Initial code content for the new cell */
-      code?: string;
-      /** The last executed code for the new cell */
-      lastCodeRun?: string;
-      /** Timestamp of the last execution */
-      lastExecutionTime?: number;
-      /** Optional custom ID for the new cell. Auto-generated if not provided */
-      newCellId?: CellId;
-      /** Whether to focus the new cell after creation */
-      autoFocus?: boolean;
-      /** If true, skip creation if code already exists */
-      skipIfCodeExists?: boolean;
-      /** Hide the code in the new cell. This will be initially shown until the cell is blurred for the first time. */
-      hideCode?: boolean;
-    },
-  ) => {
+  createNewCell: (state, action: CreateNewCellAction) => {
     const {
       cellId,
       before,