clearer edit tool schema (#7059)

…use it

## 📝 Summary

<!--
Provide a concise summary of what this pull request is addressing.

If this PR fixes any issues, list them here by number (e.g., Fixes
#123).
-->
Makes it easier for LLM's to use this tool

## 🔍 Description of Changes

<!--
Detail the specific changes made in this pull request. Explain the
problem addressed and how it was resolved. If applicable, provide before
and after comparisons, screenshots, or any relevant details to help
reviewers understand the changes easily.
-->

## 📋 Checklist

- [x] I have read the [contributor
guidelines](https://github.com/marimo-team/marimo/blob/main/CONTRIBUTING.md).
- [ ] For large changes, or changes that affect the public API: this
change was discussed or approved through an issue, on
[Discord](https://marimo.io/discord?ref=pr), or the community
[discussions](https://github.com/marimo-team/marimo/discussions) (Please
provide a link if applicable).
- [x] I have added tests for the changes made.
- [x] I have run the code and verified that it works as expected.

Author: Light2Dark

frontend/src/core/ai/model-registry.ts

@@ -13,19 +13,6 @@ import { once } from "@/utils/once";
 import type { ProviderId } from "./ids/ids";
 import { AiModelId, type QualifiedModelId, type ShortModelId } from "./ids/ids";
 
-export const PROVIDER_SORT_ORDER: ProviderId[] = [
-  // Sort by popular ones
-  "anthropic",
-  "openai",
-  "google",
-  "github",
-  "openrouter",
-  "deepseek",
-  "azure",
-  "bedrock",
-  "ollama",
-];
-
 export interface AiModel extends AiModelType {
   roles: Role[];
   model: ShortModelId;

frontend/src/core/ai/tools/__tests__/edit-notebook-tool.test.ts

@@ -249,7 +249,7 @@ describe("EditNotebookTool", () => {
         {
           edit: {
             type: "add_cell",
-            position: "end",
+            position: { type: "notebook_end" },
             code: newCode,
           },
         },
@@ -280,7 +280,7 @@ describe("EditNotebookTool", () => {
         {
           edit: {
             type: "add_cell",
-            position: { cellId: cellId2, before: true },
+            position: { type: "relative", cellId: cellId2, before: true },
             code: newCode,
           },
         },
@@ -310,7 +310,7 @@ describe("EditNotebookTool", () => {
         {
           edit: {
             type: "add_cell",
-            position: { cellId: cellId2, before: false },
+            position: { type: "relative", cellId: cellId2, before: false },
             code: newCode,
           },
         },
@@ -340,7 +340,7 @@ describe("EditNotebookTool", () => {
         {
           edit: {
             type: "add_cell",
-            position: { type: "end", columnIndex: 1 },
+            position: { type: "column_end", columnIndex: 1 },
             code: newCode,
           },
         },
@@ -367,7 +367,11 @@ describe("EditNotebookTool", () => {
           {
             edit: {
               type: "add_cell",
-              position: { cellId: "nonexistent" as CellId, before: true },
+              position: {
+                type: "relative",
+                cellId: "nonexistent" as CellId,
+                before: true,
+              },
               code: "y = 2",
             },
           },
@@ -390,7 +394,7 @@ describe("EditNotebookTool", () => {
             edit: {
               type: "add_cell",
               position: {
-                type: "end",
+                type: "column_end",
                 columnIndex: -1,
               },
               code: "y = 2",
@@ -540,7 +544,7 @@ describe("EditNotebookTool", () => {
         {
           edit: {
             type: "add_cell",
-            position: "end",
+            position: { type: "notebook_end" },
             code: "y = 2",
           },
         },

frontend/src/core/ai/tools/__tests__/registry.test.ts

@@ -55,9 +55,9 @@ describe("FrontendToolRegistry", () => {
     expect(typeof response.error).toBe("string");
 
     // Verify error message contains expected prefix
-    expect(response.error).toContain("Error invoking tool ToolExecutionError:");
-    expect(response.error).toContain('"code":"TOOL_ERROR"');
-    expect(response.error).toContain('"is_retryable":false');
+    expect(response.error).toMatchInlineSnapshot(
+      `"Error invoking tool ToolExecutionError: {"message":"Tool test_frontend_tool returned invalid input: ✖ Invalid input: expected string, received undefined\\n  → at name","code":"INVALID_ARGUMENTS","is_retryable":true,"suggested_fix":"Please check the arguments and try again."}"`,
+    );
   });
 
   it("returns tool schemas with expected shape and memoizes the result", () => {

frontend/src/core/ai/tools/base.ts

@@ -51,8 +51,8 @@ export class ToolExecutionError extends Error {
       message: this.message,
       code: this.code,
       is_retryable: this.isRetryable,
-      suggested_fix: this.suggestedFix,
-      meta: this.meta ?? {},
+      ...(this.suggestedFix && { suggested_fix: this.suggestedFix }),
+      ...(this.meta && { meta: this.meta }),
     });
     return `Error invoking tool ${this.name}: ${stringError}`;
   }

frontend/src/core/ai/tools/edit-notebook-tool.ts

@@ -31,12 +31,12 @@ const description: ToolDescription = {
   additionalInfo: `
   Args:
     edit (object): The editing operation to perform. Must be one of:
-    - update_cell: Update the code of an existing cell, pass CellId and the new code.
-    - add_cell: Add a new cell to the notebook. The position of the new cell is specified by the position argument.
-        Pass "end" to add the new cell at the end of the notebook. 
-        Pass { cellId: cellId, before: true } to add the new cell before the specified cell. And before: false if after the specified cell.
-        Pass { type: "end", columnIndex: number } to add the new cell at the end of a specified column index. The column index is 0-based.
-    - delete_cell: Delete an existing cell, pass CellId. For deleting cells, the user needs to accept the deletion to actually delete the cell, so you may still see the cell in the notebook on subsequent edits which is fine.
+    - update_cell: Update the code of an existing cell, pass cellId and the new code.
+    - add_cell: Add a new cell to the notebook. The position is specified by the position object with a "type" field:
+        { type: "notebook_end" } - Add at the end of the notebook
+        { type: "relative", cellId: "...", before: true } - Add before the specified cell (before: false for after)
+        { type: "column_end", columnIndex: 0 } - Add at the end of a specific column (0-based index)
+    - delete_cell: Delete an existing cell, pass cellId. For deleting cells, the user needs to accept the deletion to actually delete the cell, so you may still see the cell in the notebook on subsequent edits which is fine.
 
     For adding code, use the following guidelines:
     - Markdown cells: use mo.md(f"""{content}""") function to insert content.
@@ -47,9 +47,9 @@ const description: ToolDescription = {
 };
 
 type CellPosition =
-  | { cellId: CellId; before: boolean }
-  | { type: "end"; columnIndex: number }
-  | "end";
+  | { type: "relative"; cellId: CellId; before: boolean }
+  | { type: "column_end"; columnIndex: number }
+  | { type: "notebook_end" };
 
 const editNotebookSchema = z.object({
   edit: z.discriminatedUnion("type", [
@@ -60,16 +60,19 @@ const editNotebookSchema = z.object({
     }),
     z.object({
       type: z.literal("add_cell"),
-      position: z.union([
+      position: z.discriminatedUnion("type", [
         z.object({
+          type: z.literal("relative"),
           cellId: z.string() as unknown as z.ZodType<CellId>,
           before: z.boolean(),
         }),
         z.object({
-          type: z.literal("end"),
+          type: z.literal("column_end"),
           columnIndex: z.number(),
         }),
-        z.literal("end"),
+        z.object({
+          type: z.literal("notebook_end"),
+        }),
       ]) satisfies z.ZodType<CellPosition>,
       code: z.string(),
     }),
@@ -135,21 +138,25 @@ export class EditNotebookTool
       case "add_cell": {
         const { position, code } = edit;
 
-        // By default, add the new cell to the end of the notebook
         let notebookPosition: NotebookCellPosition = "__end__";
         let before = false;
         const newCellId = CellId.create();
+        const notebook = store.get(notebookAtom);
 
-        if (typeof position === "object") {
-          const notebook = store.get(notebookAtom);
-          if ("cellId" in position) {
+        switch (position.type) {
+          case "relative":
             this.validateCellIdExists(position.cellId, notebook);
             notebookPosition = position.cellId;
             before = position.before;
-          } else if ("columnIndex" in position) {
+            break;
+          case "column_end": {
             const columnId = this.getColumnId(position.columnIndex, notebook);
             notebookPosition = { type: "__end__", columnId };
+            break;
           }
+          case "notebook_end":
+            // Use default: notebookPosition = "__end__"
+            break;
         }
 
         createNewCell({

frontend/src/core/ai/tools/registry.ts

@@ -67,7 +67,12 @@ export class FrontendToolRegistry {
       const inputResponse = await inputSchema.safeParseAsync(rawArgs);
       if (inputResponse.error) {
         const strError = z.prettifyError(inputResponse.error);
-        throw new Error(`Tool ${toolName} returned invalid input: ${strError}`);
+        throw new ToolExecutionError(
+          `Tool ${toolName} returned invalid input: ${strError}`,
+          "INVALID_ARGUMENTS",
+          true,
+          "Please check the arguments and try again.",
+        );
       }
       const args = inputResponse.data;