From 08dd47528ad39d80ba0f3bb744120ff49f21b624 Mon Sep 17 00:00:00 2001
From: Shahmir Varqha <sham9871@gmail.com>
Date: Tue, 30 Sep 2025 00:25:28 +0800
Subject: [PATCH 1/4] add simple auto fix button for sql cell

---
 .../src/components/editor/errors/auto-fix.tsx |  8 +++++++-
 .../editor/output/MarimoErrorOutput.tsx       |  8 +++++++-
 .../src/core/errors/__tests__/errors.test.ts  | 12 ++++++++++++
 frontend/src/core/errors/errors.ts            | 19 +++++++++++++++++++
 4 files changed, 45 insertions(+), 2 deletions(-)

diff --git a/frontend/src/components/editor/errors/auto-fix.tsx b/frontend/src/components/editor/errors/auto-fix.tsx
index 1ffc01359b2..e8e69e82455 100644
--- a/frontend/src/components/editor/errors/auto-fix.tsx
+++ b/frontend/src/components/editor/errors/auto-fix.tsx
@@ -1,17 +1,21 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 
+import { useSetAtom } from "jotai";
 import { WrenchIcon } from "lucide-react";
 import { Button } from "@/components/ui/button";
 import { Tooltip } from "@/components/ui/tooltip";
+import { aiCompletionCellAtom } from "@/core/ai/state";
 import { notebookAtom, useCellActions } from "@/core/cells/cells";
 import type { CellId } from "@/core/cells/ids";
 import { getAutoFixes } from "@/core/errors/errors";
 import type { MarimoError } from "@/core/kernel/messages";
 import { store } from "@/core/state/jotai";
+import { cn } from "@/utils/cn";
 
 export const AutoFixButton = ({
   errors,
   cellId,
+  className,
 }: {
   errors: MarimoError[];
   cellId: CellId;
@@ -19,6 +23,7 @@ export const AutoFixButton = ({
 }) => {
   const { createNewCell } = useCellActions();
   const autoFixes = errors.flatMap((error) => getAutoFixes(error));
+  const setAiCompletionCell = useSetAtom(aiCompletionCellAtom);
 
   if (autoFixes.length === 0) {
     return null;
@@ -33,7 +38,7 @@ export const AutoFixButton = ({
       <Button
         size="xs"
         variant="outline"
-        className="my-2 font-normal"
+        className={cn("my-2 font-normal", className)}
         onClick={() => {
           const editorView =
             store.get(notebookAtom).cellHandles[cellId].current?.editorView;
@@ -48,6 +53,7 @@ export const AutoFixButton = ({
             },
             editor: editorView,
             cellId: cellId,
+            setAiCompletionCell,
           });
           // Focus the editor
           editorView?.focus();
diff --git a/frontend/src/components/editor/output/MarimoErrorOutput.tsx b/frontend/src/components/editor/output/MarimoErrorOutput.tsx
index f460b4cb905..c7bb719c698 100644
--- a/frontend/src/components/editor/output/MarimoErrorOutput.tsx
+++ b/frontend/src/components/editor/output/MarimoErrorOutput.tsx
@@ -533,7 +533,13 @@ export const MarimoErrorOutput = ({
               </div>
             );
           })}
-          {cellId && <AutoFixButton errors={sqlErrors} cellId={cellId} />}
+          {cellId && (
+            <AutoFixButton
+              errors={sqlErrors}
+              cellId={cellId}
+              className="mt-2.5"
+            />
+          )}
         </div>,
       );
     }
diff --git a/frontend/src/core/errors/__tests__/errors.test.ts b/frontend/src/core/errors/__tests__/errors.test.ts
index 3c1793e4014..810d6199058 100644
--- a/frontend/src/core/errors/__tests__/errors.test.ts
+++ b/frontend/src/core/errors/__tests__/errors.test.ts
@@ -42,6 +42,18 @@ describe("getAutoFixes", () => {
     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);
+    expect(fixes).toHaveLength(1);
+    expect(fixes[0].title).toBe("AI Fix: Fix the SQL error");
+  });
+
   it("returns no fixes for NameError with unknown import", () => {
     const error: MarimoError = {
       type: "exception",
diff --git a/frontend/src/core/errors/errors.ts b/frontend/src/core/errors/errors.ts
index c66051ffe5b..6aba43774b7 100644
--- a/frontend/src/core/errors/errors.ts
+++ b/frontend/src/core/errors/errors.ts
@@ -12,6 +12,10 @@ export interface AutoFix {
     addCodeBelow: (code: string) => void;
     editor: EditorView | undefined;
     cellId: CellId;
+    setAiCompletionCell?: (cell: {
+      cellId: CellId;
+      initialPrompt?: string;
+    }) => void;
   }) => Promise<void>;
 }
 
@@ -57,6 +61,21 @@ export function getAutoFixes(error: MarimoError): AutoFix[] {
     ];
   }
 
+  if (error.type === "sql-error") {
+    return [
+      {
+        title: "AI Fix: Fix the SQL error",
+        description: "Fix the SQL statement",
+        onFix: async (ctx) => {
+          ctx.setAiCompletionCell?.({
+            cellId: ctx.cellId,
+            initialPrompt: `Fix the SQL statement: ${error.msg}`,
+          });
+        },
+      },
+    ];
+  }
+
   return [];
 }
 

From 9f8bafb10805a70b4cffff20794b150fd790a47c Mon Sep 17 00:00:00 2001
From: Myles Scolnick <myles@marimo.io>
Date: Mon, 29 Sep 2025 13:16:37 -0400
Subject: [PATCH 2/4] Update frontend/src/core/errors/errors.ts

---
 frontend/src/core/errors/errors.ts | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/frontend/src/core/errors/errors.ts b/frontend/src/core/errors/errors.ts
index 6aba43774b7..147edebaad4 100644
--- a/frontend/src/core/errors/errors.ts
+++ b/frontend/src/core/errors/errors.ts
@@ -64,7 +64,7 @@ export function getAutoFixes(error: MarimoError): AutoFix[] {
   if (error.type === "sql-error") {
     return [
       {
-        title: "AI Fix: Fix the SQL error",
+        title: "Fix with AI",
         description: "Fix the SQL statement",
         onFix: async (ctx) => {
           ctx.setAiCompletionCell?.({

From 24c7d1bf127f05724aa0c75d798938f7a478f5e0 Mon Sep 17 00:00:00 2001
From: Myles Scolnick <myles@marimo.io>
Date: Mon, 29 Sep 2025 13:23:49 -0400
Subject: [PATCH 3/4] fix

---
 .../src/components/editor/errors/auto-fix.tsx  |  8 ++++++--
 .../src/core/errors/__tests__/errors.test.ts   | 18 ++++++++++++------
 frontend/src/core/errors/errors.ts             | 11 ++++++++++-
 3 files changed, 28 insertions(+), 9 deletions(-)

diff --git a/frontend/src/components/editor/errors/auto-fix.tsx b/frontend/src/components/editor/errors/auto-fix.tsx
index e8e69e82455..c5ed56e0789 100644
--- a/frontend/src/components/editor/errors/auto-fix.tsx
+++ b/frontend/src/components/editor/errors/auto-fix.tsx
@@ -1,12 +1,13 @@
 /* Copyright 2024 Marimo. All rights reserved. */
 
-import { useSetAtom } from "jotai";
+import { useAtomValue, useSetAtom } from "jotai";
 import { WrenchIcon } from "lucide-react";
 import { Button } from "@/components/ui/button";
 import { Tooltip } from "@/components/ui/tooltip";
 import { aiCompletionCellAtom } from "@/core/ai/state";
 import { notebookAtom, useCellActions } from "@/core/cells/cells";
 import type { CellId } from "@/core/cells/ids";
+import { aiEnabledAtom } from "@/core/config/config";
 import { getAutoFixes } from "@/core/errors/errors";
 import type { MarimoError } from "@/core/kernel/messages";
 import { store } from "@/core/state/jotai";
@@ -22,7 +23,10 @@ export const AutoFixButton = ({
   className?: string;
 }) => {
   const { createNewCell } = useCellActions();
-  const autoFixes = errors.flatMap((error) => getAutoFixes(error));
+  const aiEnabled = useAtomValue(aiEnabledAtom);
+  const autoFixes = errors.flatMap((error) =>
+    getAutoFixes(error, { aiEnabled }),
+  );
   const setAiCompletionCell = useSetAtom(aiCompletionCellAtom);
 
   if (autoFixes.length === 0) {
diff --git a/frontend/src/core/errors/__tests__/errors.test.ts b/frontend/src/core/errors/__tests__/errors.test.ts
index 810d6199058..720bb8c86d5 100644
--- a/frontend/src/core/errors/__tests__/errors.test.ts
+++ b/frontend/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,7 +41,7 @@ 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'");
   });
@@ -49,9 +53,11 @@ describe("getAutoFixes", () => {
       sql_statement: "SELECT * FROM table",
     };
 
-    const fixes = getAutoFixes(error);
+    const fixes = getAutoFixes(error, opts);
     expect(fixes).toHaveLength(1);
-    expect(fixes[0].title).toBe("AI Fix: Fix the SQL error");
+    expect(fixes[0].title).toBe("Fix with AI");
+
+    expect(getAutoFixes(error, { aiEnabled: false })).toHaveLength(0);
   });
 
   it("returns no fixes for NameError with unknown import", () => {
@@ -62,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", () => {
@@ -71,6 +77,6 @@ describe("getAutoFixes", () => {
       msg: "invalid syntax",
     };
 
-    expect(getAutoFixes(error)).toHaveLength(0);
+    expect(getAutoFixes(error, opts)).toHaveLength(0);
   });
 });
diff --git a/frontend/src/core/errors/errors.ts b/frontend/src/core/errors/errors.ts
index 147edebaad4..0cc30d1da88 100644
--- a/frontend/src/core/errors/errors.ts
+++ b/frontend/src/core/errors/errors.ts
@@ -19,7 +19,12 @@ export interface AutoFix {
   }) => Promise<void>;
 }
 
-export function getAutoFixes(error: MarimoError): AutoFix[] {
+export function getAutoFixes(
+  error: MarimoError,
+  opts: {
+    aiEnabled: boolean;
+  },
+): AutoFix[] {
   if (error.type === "multiple-defs") {
     return [
       {
@@ -62,6 +67,10 @@ 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",

From 8a45cdd73001729607d012fb752c7d782cda5bb1 Mon Sep 17 00:00:00 2001
From: Myles Scolnick <myles@marimo.io>
Date: Mon, 29 Sep 2025 13:31:48 -0400
Subject: [PATCH 4/4] fix ai prompts

---
 marimo/_server/ai/prompts.py                  |  4 +--
 .../ai/snapshots/chat_system_prompts.txt      | 14 ----------
 tests/_server/ai/snapshots/system_prompts.txt | 26 ++++++++-----------
 3 files changed, 12 insertions(+), 32 deletions(-)

diff --git a/marimo/_server/ai/prompts.py b/marimo/_server/ai/prompts.py
index 772a9cab034..23776e653a0 100644
--- a/marimo/_server/ai/prompts.py
+++ b/marimo/_server/ai/prompts.py
@@ -29,8 +29,6 @@
     "markdown": [],
     "sql": [
         "The SQL must use duckdb syntax.",
-        'SQL cells start with df = mo.sql(f"""<your query>""") for DuckDB, or df = mo.sql(f"""<your query>""", engine=engine) for other SQL engines.',
-        "This will automatically display the result in the UI. You do not need to return the dataframe in the cell.",
     ],
 }
 
@@ -194,7 +192,7 @@ def get_refactor_or_insert_notebook_cell_system_prompt(
     if support_multiple_cells:
         system_prompt += "\n\nAgain, just output code wrapped in cells. Each cell is wrapped in backticks with the appropriate language identifier (python, sql, markdown)."
     else:
-        system_prompt += "\n\nAgain, just output the code itself."
+        system_prompt += f"\n\nAgain, just output the code itself and make sure to return the code as just {language}."
 
     return system_prompt
 
diff --git a/tests/_server/ai/snapshots/chat_system_prompts.txt b/tests/_server/ai/snapshots/chat_system_prompts.txt
index 76c36fbd401..69f3377a28e 100644
--- a/tests/_server/ai/snapshots/chat_system_prompts.txt
+++ b/tests/_server/ai/snapshots/chat_system_prompts.txt
@@ -141,8 +141,6 @@ chart
 
 ## Rules for sql:
 1. The SQL must use duckdb syntax.
-2. SQL cells start with df = mo.sql(f"""<your query>""") for DuckDB, or df = mo.sql(f"""<your query>""", engine=engine) for other SQL engines.
-3. This will automatically display the result in the UI. You do not need to return the dataframe in the cell.
 
 ==================== with custom rules ====================
 
@@ -285,8 +283,6 @@ chart
 
 ## Rules for sql:
 1. The SQL must use duckdb syntax.
-2. SQL cells start with df = mo.sql(f"""<your query>""") for DuckDB, or df = mo.sql(f"""<your query>""", engine=engine) for other SQL engines.
-3. This will automatically display the result in the UI. You do not need to return the dataframe in the cell.
 
 ## Additional rules:
 Always be polite.
@@ -432,8 +428,6 @@ chart
 
 ## Rules for sql:
 1. The SQL must use duckdb syntax.
-2. SQL cells start with df = mo.sql(f"""<your query>""") for DuckDB, or df = mo.sql(f"""<your query>""", engine=engine) for other SQL engines.
-3. This will automatically display the result in the UI. You do not need to return the dataframe in the cell.
 
 ## Available variables from other cells:
 - variable: `var1`- variable: `var2`
@@ -579,8 +573,6 @@ chart
 
 ## Rules for sql:
 1. The SQL must use duckdb syntax.
-2. SQL cells start with df = mo.sql(f"""<your query>""") for DuckDB, or df = mo.sql(f"""<your query>""", engine=engine) for other SQL engines.
-3. This will automatically display the result in the UI. You do not need to return the dataframe in the cell.
 
 ## Available variables from other cells:
 - variable: `df`
@@ -732,8 +724,6 @@ chart
 
 ## Rules for sql:
 1. The SQL must use duckdb syntax.
-2. SQL cells start with df = mo.sql(f"""<your query>""") for DuckDB, or df = mo.sql(f"""<your query>""", engine=engine) for other SQL engines.
-3. This will automatically display the result in the UI. You do not need to return the dataframe in the cell.
 
 ## Available schema:
 - Table: df_1
@@ -893,8 +883,6 @@ chart
 
 ## Rules for sql:
 1. The SQL must use duckdb syntax.
-2. SQL cells start with df = mo.sql(f"""<your query>""") for DuckDB, or df = mo.sql(f"""<your query>""", engine=engine) for other SQL engines.
-3. This will automatically display the result in the UI. You do not need to return the dataframe in the cell.
 
 <code_from_other_cells>
 import pandas as pd
@@ -1042,8 +1030,6 @@ chart
 
 ## Rules for sql:
 1. The SQL must use duckdb syntax.
-2. SQL cells start with df = mo.sql(f"""<your query>""") for DuckDB, or df = mo.sql(f"""<your query>""", engine=engine) for other SQL engines.
-3. This will automatically display the result in the UI. You do not need to return the dataframe in the cell.
 
 ## Additional rules:
 Always be polite.
diff --git a/tests/_server/ai/snapshots/system_prompts.txt b/tests/_server/ai/snapshots/system_prompts.txt
index 396cb6f08a9..b924a95325e 100644
--- a/tests/_server/ai/snapshots/system_prompts.txt
+++ b/tests/_server/ai/snapshots/system_prompts.txt
@@ -22,7 +22,7 @@ Immediately start with the following format. Do NOT comment on the code, just ou
 6. If an import already exists, do not import it again.
 7. If a variable is already defined, use another name, or make it private by adding an underscore at the beginning.
 
-Again, just output the code itself.
+Again, just output the code itself and make sure to return the code as just python.
 
 ==================== markdown ====================
 
@@ -37,7 +37,7 @@ Immediately start with the following format. Do NOT comment on the code, just ou
 {CELL_CODE}
 ```
 
-Again, just output the code itself.
+Again, just output the code itself and make sure to return the code as just markdown.
 
 ==================== sql ====================
 
@@ -54,10 +54,8 @@ Immediately start with the following format. Do NOT comment on the code, just ou
 
 ## Rules for sql
 1. The SQL must use duckdb syntax.
-2. SQL cells start with df = mo.sql(f"""<your query>""") for DuckDB, or df = mo.sql(f"""<your query>""", engine=engine) for other SQL engines.
-3. This will automatically display the result in the UI. You do not need to return the dataframe in the cell.
 
-Again, just output the code itself.
+Again, just output the code itself and make sure to return the code as just sql.
 
 ==================== idk ====================
 
@@ -72,7 +70,7 @@ Immediately start with the following format. Do NOT comment on the code, just ou
 {CELL_CODE}
 ```
 
-Again, just output the code itself.
+Again, just output the code itself and make sure to return the code as just idk.
 
 ==================== with custom rules ====================
 
@@ -99,7 +97,7 @@ Immediately start with the following format. Do NOT comment on the code, just ou
 ## Additional rules:
 Always use type hints.
 
-Again, just output the code itself.
+Again, just output the code itself and make sure to return the code as just python.
 
 ==================== with context ====================
 
@@ -133,7 +131,7 @@ Immediately start with the following format. Do NOT comment on the code, just ou
     - Sample values: Alice, Bob, Charlie
 
 
-Again, just output the code itself.
+Again, just output the code itself and make sure to return the code as just python.
 
 ==================== with is_insert=True ====================
 
@@ -164,7 +162,7 @@ Immediately start with the following format. Do NOT comment on the code, just ou
 6. If an import already exists, do not import it again.
 7. If a variable is already defined, use another name, or make it private by adding an underscore at the beginning.
 
-Again, just output the code itself.
+Again, just output the code itself and make sure to return the code as just python.
 
 ==================== with cell_code ====================
 
@@ -201,7 +199,7 @@ print('Hello, world!')
 6. If an import already exists, do not import it again.
 7. If a variable is already defined, use another name, or make it private by adding an underscore at the beginning.
 
-Again, just output the code itself.
+Again, just output the code itself and make sure to return the code as just python.
 
 ==================== with selected_text ====================
 
@@ -238,7 +236,7 @@ print('Hello, world!')
 6. If an import already exists, do not import it again.
 7. If a variable is already defined, use another name, or make it private by adding an underscore at the beginning.
 
-Again, just output the code itself.
+Again, just output the code itself and make sure to return the code as just python.
 
 ==================== with other_cell_codes ====================
 
@@ -279,7 +277,7 @@ import pandas as pd
 import numpy as np
 </code_from_other_cells>
 
-Again, just output the code itself.
+Again, just output the code itself and make sure to return the code as just python.
 
 ==================== with VariableContext objects ====================
 
@@ -312,7 +310,7 @@ Immediately start with the following format. Do NOT comment on the code, just ou
   - value_preview: <Model object>
 
 
-Again, just output the code itself.
+Again, just output the code itself and make sure to return the code as just python.
 
 ==================== with support_multiple_cells=True ====================
 
@@ -350,8 +348,6 @@ Separate logic into multiple cells to keep the code organized and readable.
 
 ## Rules for sql:
 1. The SQL must use duckdb syntax.
-2. SQL cells start with df = mo.sql(f"""<your query>""") for DuckDB, or df = mo.sql(f"""<your query>""", engine=engine) for other SQL engines.
-3. This will automatically display the result in the UI. You do not need to return the dataframe in the cell.
 
 ## Available variables from other cells:
 - variable: `df`