clearer edit tool schema

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;