@jupyterlite/ai 0.9.0 → 0.10.0

package/src/tools/notebook.ts

@@ -1,10 +1,15 @@
-import { CodeCell, ICodeCellModel, MarkdownCell } from '@jupyterlab/cells';
+import {
+  CodeCell,
+  CodeCellModel,
+  ICodeCellModel,
+  MarkdownCell
+} from '@jupyterlab/cells';
 import { IDocumentManager } from '@jupyterlab/docmanager';
 import { DocumentWidget } from '@jupyterlab/docregistry';
 import { INotebookTracker, NotebookPanel } from '@jupyterlab/notebook';
 import { KernelSpec } from '@jupyterlab/services';
 
-import { tool } from '@openai/agents';
+import { tool } from 'ai';
 
 import { z } from 'zod';
 
@@ -91,10 +96,10 @@ export function createNotebookCreationTool(
   kernelSpecManager: KernelSpec.IManager
 ): ITool {
   return tool({
-    name: 'create_notebook',
+    title: 'Create Notebook',
     description:
       'Create a new Jupyter notebook with a kernel for the specified programming language',
-    parameters: z.object({
+    inputSchema: z.object({
       language: z
         .string()
         .optional()
@@ -121,49 +126,48 @@ export function createNotebookCreationTool(
       const { name } = input;
 
       if (!name) {
-        throw new Error('A name must be provided to create a notebook');
+        return {
+          success: false,
+          error: 'A name must be provided to create a notebook'
+        };
       }
 
-      try {
-        // TODO: handle cwd / path?
-        const fileName = name.endsWith('.ipynb') ? name : `${name}.ipynb`;
-
-        // Create untitled notebook first
-        const notebookModel = await docManager.newUntitled({
-          type: 'notebook'
-        });
+      // TODO: handle cwd / path?
+      const fileName = name.endsWith('.ipynb') ? name : `${name}.ipynb`;
 
-        // Rename to desired filename
-        await docManager.services.contents.rename(notebookModel.path, fileName);
-
-        // Create widget with specific kernel
-        const notebook = docManager.createNew(fileName, 'default', {
-          name: kernel
-        });
-
-        if (!(notebook instanceof DocumentWidget)) {
-          throw new Error('Failed to create notebook widget');
-        }
+      // Create untitled notebook first
+      const notebookModel = await docManager.newUntitled({
+        type: 'notebook'
+      });
 
-        await notebook.context.ready;
-        await notebook.context.save();
+      // Rename to desired filename
+      await docManager.services.contents.rename(notebookModel.path, fileName);
 
-        docManager.openOrReveal(fileName);
+      // Create widget with specific kernel
+      const notebook = docManager.createNew(fileName, 'default', {
+        name: kernel
+      });
 
-        return {
-          success: true,
-          message: `Successfully created notebook ${fileName} with ${kernel} kernel${input.language ? ` for ${input.language}` : ''}`,
-          notebookPath: fileName,
-          notebookName: fileName,
-          kernel,
-          language: input.language
-        };
-      } catch (error) {
+      if (!(notebook instanceof DocumentWidget)) {
         return {
           success: false,
-          error: `Failed to create notebook: ${(error as Error).message}`
+          error: 'Failed to create notebook widget'
         };
       }
+
+      await notebook.context.ready;
+      await notebook.context.save();
+
+      docManager.openOrReveal(fileName);
+
+      return {
+        success: true,
+        message: `Successfully created notebook ${fileName} with ${kernel} kernel${input.language ? ` for ${input.language}` : ''}`,
+        notebookPath: fileName,
+        notebookName: fileName,
+        kernel,
+        language: input.language
+      };
     }
   });
 }
@@ -176,9 +180,9 @@ export function createAddCellTool(
   notebookTracker?: INotebookTracker
 ): ITool {
   return tool({
-    name: 'add_cell',
+    title: 'Add Cell',
     description: 'Add a cell to the current notebook with optional content',
-    parameters: z.object({
+    inputSchema: z.object({
       notebookPath: z
         .string()
         .optional()
@@ -201,83 +205,76 @@ export function createAddCellTool(
         .default('below')
         .describe('Position relative to current cell')
     }),
-    async execute({
+    execute: async ({
       notebookPath,
       content,
       cellType = 'code',
       position = 'below'
-    }) {
-      try {
-        const currentWidget = await getNotebookWidget(
-          notebookPath,
-          docManager,
-          notebookTracker
-        );
-        if (!currentWidget) {
-          return {
-            success: false,
-            error: notebookPath
-              ? `Failed to open notebook at path: ${notebookPath}`
-              : 'No active notebook and no notebook path provided'
-          };
-        }
+    }) => {
+      const currentWidget = await getNotebookWidget(
+        notebookPath,
+        docManager,
+        notebookTracker
+      );
+      if (!currentWidget) {
+        return {
+          success: false,
+          error: notebookPath
+            ? `Failed to open notebook at path: ${notebookPath}`
+            : 'No active notebook and no notebook path provided'
+        };
+      }
 
-        const notebook = currentWidget.content;
-        const model = notebook.model;
+      const notebook = currentWidget.content;
+      const model = notebook.model;
 
-        if (!model) {
-          return {
-            success: false,
-            error: 'No notebook model available'
-          };
-        }
+      if (!model) {
+        return {
+          success: false,
+          error: 'No notebook model available'
+        };
+      }
 
-        // Check if we should replace the first empty cell instead of adding
-        const shouldReplaceFirstCell =
-          model.cells.length === 1 &&
-          model.cells.get(0).sharedModel.getSource().trim() === '';
+      // Check if we should replace the first empty cell instead of adding
+      const shouldReplaceFirstCell =
+        model.cells.length === 1 &&
+        model.cells.get(0).sharedModel.getSource().trim() === '';
 
-        if (shouldReplaceFirstCell) {
-          // Replace the first empty cell by removing it and adding new one
-          model.sharedModel.deleteCell(0);
-        }
+      if (shouldReplaceFirstCell) {
+        // Replace the first empty cell by removing it and adding new one
+        model.sharedModel.deleteCell(0);
+      }
 
-        // Create the new cell using shared model
-        const newCellData = {
-          cell_type: cellType,
-          source: content || '',
-          metadata: cellType === 'code' ? { trusted: true } : {}
-        };
+      // Create the new cell using shared model
+      const newCellData = {
+        cell_type: cellType,
+        source: content || '',
+        metadata: cellType === 'code' ? { trusted: true } : {}
+      };
 
-        model.sharedModel.addCell(newCellData);
-
-        // Execute markdown cells after creation to render them
-        if (cellType === 'markdown' && content) {
-          const cellIndex = model.cells.length - 1;
-          const cellWidget = notebook.widgets[cellIndex];
-          if (cellWidget && cellWidget instanceof MarkdownCell) {
-            try {
-              await cellWidget.ready;
-              cellWidget.rendered = true;
-            } catch (error) {
-              console.warn('Failed to render markdown cell:', error);
-            }
+      model.sharedModel.addCell(newCellData);
+
+      // Execute markdown cells after creation to render them
+      if (cellType === 'markdown' && content) {
+        const cellIndex = model.cells.length - 1;
+        const cellWidget = notebook.widgets[cellIndex];
+        if (cellWidget && cellWidget instanceof MarkdownCell) {
+          try {
+            await cellWidget.ready;
+            cellWidget.rendered = true;
+          } catch (error) {
+            console.warn('Failed to render markdown cell:', error);
           }
         }
-
-        return {
-          success: true,
-          message: `${cellType} cell added successfully`,
-          content: content || '',
-          cellType,
-          position
-        };
-      } catch (error) {
-        return {
-          success: false,
-          error: `Failed to add ${cellType} cell: ${(error as Error).message}`
-        };
       }
+
+      return {
+        success: true,
+        message: `${cellType} cell added successfully`,
+        content: content || '',
+        cellType,
+        position
+      };
     }
   });
 }
@@ -290,10 +287,10 @@ export function createGetNotebookInfoTool(
   notebookTracker?: INotebookTracker
 ): ITool {
   return tool({
-    name: 'get_notebook_info',
+    title: 'Get Notebook Info',
     description:
       'Get information about a notebook including number of cells and active cell index',
-    parameters: z.object({
+    inputSchema: z.object({
       notebookPath: z
         .string()
         .optional()
@@ -305,51 +302,44 @@ export function createGetNotebookInfoTool(
     execute: async (input: { notebookPath?: string | null }) => {
       const { notebookPath } = input;
 
-      try {
-        const currentWidget = await getNotebookWidget(
-          notebookPath,
-          docManager,
-          notebookTracker
-        );
-        if (!currentWidget) {
-          return JSON.stringify({
-            success: false,
-            error: notebookPath
-              ? `Failed to open notebook at path: ${notebookPath}`
-              : 'No active notebook and no notebook path provided'
-          });
-        }
-
-        const notebook = currentWidget.content;
-        const model = notebook.model;
-
-        if (!model) {
-          return JSON.stringify({
-            success: false,
-            error: 'No notebook model available'
-          });
-        }
-
-        const cellCount = model.cells.length;
-        const activeCellIndex = notebook.activeCellIndex;
-        const activeCell = notebook.activeCell;
-        const activeCellType = activeCell?.model.type || 'unknown';
-
+      const currentWidget = await getNotebookWidget(
+        notebookPath,
+        docManager,
+        notebookTracker
+      );
+      if (!currentWidget) {
         return JSON.stringify({
-          success: true,
-          notebookName: currentWidget.title.label,
-          notebookPath: currentWidget.context.path,
-          cellCount,
-          activeCellIndex,
-          activeCellType,
-          isDirty: model.dirty
+          success: false,
+          error: notebookPath
+            ? `Failed to open notebook at path: ${notebookPath}`
+            : 'No active notebook and no notebook path provided'
         });
-      } catch (error) {
+      }
+
+      const notebook = currentWidget.content;
+      const model = notebook.model;
+
+      if (!model) {
         return JSON.stringify({
           success: false,
-          error: `Failed to get notebook info: ${(error as Error).message}`
+          error: 'No notebook model available'
         });
       }
+
+      const cellCount = model.cells.length;
+      const activeCellIndex = notebook.activeCellIndex;
+      const activeCell = notebook.activeCell;
+      const activeCellType = activeCell?.model.type || 'unknown';
+
+      return JSON.stringify({
+        success: true,
+        notebookName: currentWidget.title.label,
+        notebookPath: currentWidget.context.path,
+        cellCount,
+        activeCellIndex,
+        activeCellType,
+        isDirty: model.dirty
+      });
     }
   });
 }
@@ -362,10 +352,10 @@ export function createGetCellInfoTool(
   notebookTracker?: INotebookTracker
 ): ITool {
   return tool({
-    name: 'get_cell_info',
+    title: 'Get Cell Info',
     description:
       'Get information about a specific cell including its type, source content, and outputs',
-    parameters: z.object({
+    inputSchema: z.object({
       notebookPath: z
         .string()
         .optional()
@@ -387,70 +377,64 @@ export function createGetCellInfoTool(
     }) => {
       const { notebookPath } = input;
       let { cellIndex } = input;
-      try {
-        const currentWidget = await getNotebookWidget(
-          notebookPath,
-          docManager,
-          notebookTracker
-        );
-        if (!currentWidget) {
-          return JSON.stringify({
-            success: false,
-            error: notebookPath
-              ? `Failed to open notebook at path: ${notebookPath}`
-              : 'No active notebook and no notebook path provided'
-          });
-        }
-
-        const notebook = currentWidget.content;
-        const model = notebook.model;
-
-        if (!model) {
-          return JSON.stringify({
-            success: false,
-            error: 'No notebook model available'
-          });
-        }
 
-        if (cellIndex === undefined || cellIndex === null) {
-          cellIndex = notebook.activeCellIndex;
-        }
-
-        if (cellIndex < 0 || cellIndex >= model.cells.length) {
-          return JSON.stringify({
-            success: false,
-            error: `Invalid cell index: ${cellIndex}. Notebook has ${model.cells.length} cells.`
-          });
-        }
-
-        const cell = model.cells.get(cellIndex);
-        const cellType = cell.type;
-        const sharedModel = cell.sharedModel;
-        const source = sharedModel.getSource();
+      const currentWidget = await getNotebookWidget(
+        notebookPath,
+        docManager,
+        notebookTracker
+      );
+      if (!currentWidget) {
+        return JSON.stringify({
+          success: false,
+          error: notebookPath
+            ? `Failed to open notebook at path: ${notebookPath}`
+            : 'No active notebook and no notebook path provided'
+        });
+      }
 
-        // Get outputs for code cells
-        let outputs: any[] = [];
-        if (cellType === 'code') {
-          const rawOutputs = sharedModel.toJSON().outputs;
-          outputs = Array.isArray(rawOutputs) ? rawOutputs : [];
-        }
+      const notebook = currentWidget.content;
+      const model = notebook.model;
 
+      if (!model) {
         return JSON.stringify({
-          success: true,
-          cellId: cell.id,
-          cellIndex,
-          cellType,
-          source,
-          outputs,
-          executionCount:
-            cellType === 'code' ? (cell as any).executionCount : null
+          success: false,
+          error: 'No notebook model available'
         });
-      } catch (error) {
+      }
+
+      if (cellIndex === undefined || cellIndex === null) {
+        cellIndex = notebook.activeCellIndex;
+      }
+
+      if (cellIndex < 0 || cellIndex >= model.cells.length) {
         return JSON.stringify({
           success: false,
-          error: `Failed to get cell info: ${(error as Error).message}`
+          error: `Invalid cell index: ${cellIndex}. Notebook has ${model.cells.length} cells.`
         });
       }
+
+      const cell = model.cells.get(cellIndex);
+      const cellType = cell.type;
+      const sharedModel = cell.sharedModel;
+      const source = sharedModel.getSource();
+
+      // Get outputs for code cells
+      let outputs: any[] = [];
+      if (cellType === 'code') {
+        const rawOutputs = sharedModel.toJSON().outputs;
+        outputs = Array.isArray(rawOutputs) ? rawOutputs : [];
+      }
+
+      return JSON.stringify({
+        success: true,
+        cellId: cell.id,
+        cellIndex,
+        cellType,
+        source,
+        outputs,
+        executionCount:
+          cellType === 'code' ? (cell as CodeCellModel).executionCount : null
+      });
     }
   });
 }
@@ -464,10 +448,10 @@ export function createSetCellContentTool(
   diffManager?: IDiffManager
 ): ITool {
   return tool({
-    name: 'set_cell_content',
+    title: 'Set Cell Content',
     description:
       'Set the content of a specific cell and return both the previous and new content',
-    parameters: z.object({
+    inputSchema: z.object({
       notebookPath: z
         .string()
         .optional()
@@ -499,120 +483,113 @@ export function createSetCellContentTool(
     }) => {
       const { notebookPath, cellId, cellIndex, content } = input;
 
-      try {
-        const notebookWidget = await getNotebookWidget(
-          notebookPath,
-          docManager,
-          notebookTracker
-        );
-        if (!notebookWidget) {
-          return JSON.stringify({
-            success: false,
-            error: notebookPath
-              ? `Failed to open notebook at path: ${notebookPath}`
-              : 'No active notebook and no notebook path provided'
-          });
-        }
+      const notebookWidget = await getNotebookWidget(
+        notebookPath,
+        docManager,
+        notebookTracker
+      );
+      if (!notebookWidget) {
+        return JSON.stringify({
+          success: false,
+          error: notebookPath
+            ? `Failed to open notebook at path: ${notebookPath}`
+            : 'No active notebook and no notebook path provided'
+        });
+      }
 
-        const notebook = notebookWidget.content;
-        const targetNotebookPath = notebookWidget.context.path;
+      const notebook = notebookWidget.content;
+      const targetNotebookPath = notebookWidget.context.path;
 
-        const model = notebook.model;
+      const model = notebook.model;
 
-        if (!model) {
+      if (!model) {
+        return JSON.stringify({
+          success: false,
+          error: 'No notebook model available'
+        });
+      }
+
+      // Determine target cell index
+      let targetCellIndex: number;
+      if (cellId !== undefined && cellId !== null) {
+        // Find cell by ID
+        targetCellIndex = -1;
+        for (let i = 0; i < model.cells.length; i++) {
+          if (model.cells.get(i).id === cellId) {
+            targetCellIndex = i;
+            break;
+          }
+        }
+        if (targetCellIndex === -1) {
           return JSON.stringify({
             success: false,
-            error: 'No notebook model available'
+            error: `Cell with ID '${cellId}' not found in notebook`
           });
         }
-
-        // Determine target cell index
-        let targetCellIndex: number;
-        if (cellId !== undefined && cellId !== null) {
-          // Find cell by ID
-          targetCellIndex = -1;
-          for (let i = 0; i < model.cells.length; i++) {
-            if (model.cells.get(i).id === cellId) {
-              targetCellIndex = i;
-              break;
-            }
-          }
-          if (targetCellIndex === -1) {
-            return JSON.stringify({
-              success: false,
-              error: `Cell with ID '${cellId}' not found in notebook`
-            });
-          }
-        } else if (cellIndex !== undefined && cellIndex !== null) {
-          // Use provided cell index
-          if (cellIndex < 0 || cellIndex >= model.cells.length) {
-            return JSON.stringify({
-              success: false,
-              error: `Invalid cell index: ${cellIndex}. Notebook has ${model.cells.length} cells.`
-            });
-          }
-          targetCellIndex = cellIndex;
-        } else {
-          // Use active cell
-          targetCellIndex = notebook.activeCellIndex;
-          if (targetCellIndex === -1 || targetCellIndex >= model.cells.length) {
-            return JSON.stringify({
-              success: false,
-              error: 'No active cell or invalid active cell index'
-            });
-          }
+      } else if (cellIndex !== undefined && cellIndex !== null) {
+        // Use provided cell index
+        if (cellIndex < 0 || cellIndex >= model.cells.length) {
+          return JSON.stringify({
+            success: false,
+            error: `Invalid cell index: ${cellIndex}. Notebook has ${model.cells.length} cells.`
+          });
         }
-
-        // Get the target cell
-        const targetCell = model.cells.get(targetCellIndex);
-        if (!targetCell) {
+        targetCellIndex = cellIndex;
+      } else {
+        // Use active cell
+        targetCellIndex = notebook.activeCellIndex;
+        if (targetCellIndex === -1 || targetCellIndex >= model.cells.length) {
           return JSON.stringify({
             success: false,
-            error: `Cell at index ${targetCellIndex} not found`
+            error: 'No active cell or invalid active cell index'
           });
         }
+      }
 
-        const sharedModel = targetCell.sharedModel;
+      // Get the target cell
+      const targetCell = model.cells.get(targetCellIndex);
+      if (!targetCell) {
+        return JSON.stringify({
+          success: false,
+          error: `Cell at index ${targetCellIndex} not found`
+        });
+      }
 
-        // Get previous content and type
-        const previousContent = sharedModel.getSource();
-        const previousCellType = targetCell.type;
-        const retrievedCellId = targetCell.id;
+      const sharedModel = targetCell.sharedModel;
 
-        sharedModel.setSource(content);
+      // Get previous content and type
+      const previousContent = sharedModel.getSource();
+      const previousCellType = targetCell.type;
+      const retrievedCellId = targetCell.id;
 
-        // Show the cell diff using the diff manager if available
-        if (diffManager) {
-          await diffManager.showCellDiff({
-            original: previousContent,
-            modified: content,
-            cellId: retrievedCellId,
-            notebookPath: targetNotebookPath
-          });
-        }
+      sharedModel.setSource(content);
 
-        return JSON.stringify({
-          success: true,
-          message:
-            cellId !== undefined && cellId !== null
-              ? `Cell with ID '${cellId}' content replaced successfully`
-              : cellIndex !== undefined && cellIndex !== null
-                ? `Cell ${targetCellIndex} content replaced successfully`
-                : 'Active cell content replaced successfully',
-          notebookPath: targetNotebookPath,
+      // Show the cell diff using the diff manager if available
+      if (diffManager) {
+        await diffManager.showCellDiff({
+          original: previousContent,
+          modified: content,
           cellId: retrievedCellId,
-          cellIndex: targetCellIndex,
-          previousContent,
-          previousCellType,
-          newContent: content,
-          wasActiveCell: cellId === undefined && cellIndex === undefined
-        });
-      } catch (error) {
-        return JSON.stringify({
-          success: false,
-          error: `Failed to replace cell content: ${(error as Error).message}`
+          notebookPath: targetNotebookPath
         });
       }
+
+      return JSON.stringify({
+        success: true,
+        message:
+          cellId !== undefined && cellId !== null
+            ? `Cell with ID '${cellId}' content replaced successfully`
+            : cellIndex !== undefined && cellIndex !== null
+              ? `Cell ${targetCellIndex} content replaced successfully`
+              : 'Active cell content replaced successfully',
+        notebookPath: targetNotebookPath,
+        cellId: retrievedCellId,
+        cellIndex: targetCellIndex,
+        previousContent,
+        previousCellType,
+        newContent: content,
+        wasActiveCell: cellId === undefined && cellIndex === undefined
+      });
     }
   });
 }
@@ -625,9 +602,9 @@ export function createRunCellTool(
   notebookTracker?: INotebookTracker
 ): ITool {
   return tool({
-    name: 'run_cell',
+    title: 'Run Cell',
     description: 'Run a specific cell in the notebook by index',
-    parameters: z.object({
+    inputSchema: z.object({
       notebookPath: z
         .string()
         .optional()
@@ -649,85 +626,70 @@ export function createRunCellTool(
     }) => {
       const { notebookPath, cellIndex, recordTiming = true } = input;
 
-      try {
-        const currentWidget = await getNotebookWidget(
-          notebookPath,
-          docManager,
-          notebookTracker
-        );
-        if (!currentWidget) {
-          return JSON.stringify({
-            success: false,
-            error: notebookPath
-              ? `Failed to open notebook at path: ${notebookPath}`
-              : 'No active notebook and no notebook path provided'
-          });
-        }
-
-        const notebook = currentWidget.content;
-        const model = notebook.model;
+      const currentWidget = await getNotebookWidget(
+        notebookPath,
+        docManager,
+        notebookTracker
+      );
+      if (!currentWidget) {
+        return JSON.stringify({
+          success: false,
+          error: notebookPath
+            ? `Failed to open notebook at path: ${notebookPath}`
+            : 'No active notebook and no notebook path provided'
+        });
+      }
 
-        if (!model) {
-          return JSON.stringify({
-            success: false,
-            error: 'No notebook model available'
-          });
-        }
+      const notebook = currentWidget.content;
+      const model = notebook.model;
 
-        if (cellIndex < 0 || cellIndex >= model.cells.length) {
-          return JSON.stringify({
-            success: false,
-            error: `Invalid cell index: ${cellIndex}. Notebook has ${model.cells.length} cells.`
-          });
-        }
+      if (!model) {
+        return JSON.stringify({
+          success: false,
+          error: 'No notebook model available'
+        });
+      }
 
-        // Get the target cell widget
-        const cellWidget = notebook.widgets[cellIndex];
-        if (!cellWidget) {
-          return JSON.stringify({
-            success: false,
-            error: `Cell widget at index ${cellIndex} not found`
-          });
-        }
+      if (cellIndex < 0 || cellIndex >= model.cells.length) {
+        return JSON.stringify({
+          success: false,
+          error: `Invalid cell index: ${cellIndex}. Notebook has ${model.cells.length} cells.`
+        });
+      }
 
-        // Execute using shared model approach (non-disruptive)
-        try {
-          if (cellWidget instanceof CodeCell) {
-            // Use direct CodeCell.execute() method
-            const sessionCtx = currentWidget.sessionContext;
-            await CodeCell.execute(cellWidget, sessionCtx, {
-              recordTiming,
-              deletedCells: model.deletedCells
-            });
-
-            const codeModel = cellWidget.model as ICodeCellModel;
-            return JSON.stringify({
-              success: true,
-              message: `Cell ${cellIndex} executed successfully`,
-              cellIndex,
-              executionCount: codeModel.executionCount,
-              hasOutput: codeModel.outputs.length > 0
-            });
-          } else {
-            // For non-code cells, just return success
-            return JSON.stringify({
-              success: true,
-              message: `Cell ${cellIndex} is not a code cell, no execution needed`,
-              cellIndex,
-              cellType: cellWidget.model.type
-            });
-          }
-        } catch (error) {
-          return JSON.stringify({
-            success: false,
-            error: `Failed to execute cell: ${(error as Error).message}`,
-            cellIndex
-          });
-        }
-      } catch (error) {
+      // Get the target cell widget
+      const cellWidget = notebook.widgets[cellIndex];
+      if (!cellWidget) {
         return JSON.stringify({
           success: false,
-          error: `Failed to run cell: ${(error as Error).message}`
+          error: `Cell widget at index ${cellIndex} not found`
+        });
+      }
+
+      // Execute using shared model approach (non-disruptive)
+      if (cellWidget instanceof CodeCell) {
+        // Use direct CodeCell.execute() method
+        const sessionCtx = currentWidget.sessionContext;
+        await CodeCell.execute(cellWidget, sessionCtx, {
+          recordTiming,
+          deletedCells: model.deletedCells
+        });
+
+        const codeModel = cellWidget.model as ICodeCellModel;
+        return JSON.stringify({
+          success: true,
+          message: `Cell ${cellIndex} executed successfully`,
+          cellIndex,
+          executionCount: codeModel.executionCount,
+          hasOutput: codeModel.outputs.length > 0
+        });
+      } else {
+        // For non-code cells, just return success
+        return JSON.stringify({
+          success: true,
+          message: `Cell ${cellIndex} is not a code cell, no execution needed`,
+          cellIndex,
+          cellType: cellWidget.model.type
         });
       }
     }
@@ -742,9 +704,9 @@ export function createDeleteCellTool(
   notebookTracker?: INotebookTracker
 ): ITool {
   return tool({
-    name: 'delete_cell',
+    title: 'Delete Cell',
     description: 'Delete a specific cell from the notebook by index',
-    parameters: z.object({
+    inputSchema: z.object({
       notebookPath: z
         .string()
         .optional()
@@ -760,62 +722,55 @@ export function createDeleteCellTool(
     }) => {
       const { notebookPath, cellIndex } = input;
 
-      try {
-        const currentWidget = await getNotebookWidget(
-          notebookPath,
-          docManager,
-          notebookTracker
-        );
-        if (!currentWidget) {
-          return JSON.stringify({
-            success: false,
-            error: notebookPath
-              ? `Failed to open notebook at path: ${notebookPath}`
-              : 'No active notebook and no notebook path provided'
-          });
-        }
-
-        const notebook = currentWidget.content;
-        const model = notebook.model;
-
-        if (!model) {
-          return JSON.stringify({
-            success: false,
-            error: 'No notebook model available'
-          });
-        }
-
-        if (cellIndex < 0 || cellIndex >= model.cells.length) {
-          return JSON.stringify({
-            success: false,
-            error: `Invalid cell index: ${cellIndex}. Notebook has ${model.cells.length} cells.`
-          });
-        }
+      const currentWidget = await getNotebookWidget(
+        notebookPath,
+        docManager,
+        notebookTracker
+      );
+      if (!currentWidget) {
+        return JSON.stringify({
+          success: false,
+          error: notebookPath
+            ? `Failed to open notebook at path: ${notebookPath}`
+            : 'No active notebook and no notebook path provided'
+        });
+      }
 
-        // Validate cell exists
-        const targetCell = model.cells.get(cellIndex);
-        if (!targetCell) {
-          return JSON.stringify({
-            success: false,
-            error: `Cell at index ${cellIndex} not found`
-          });
-        }
+      const notebook = currentWidget.content;
+      const model = notebook.model;
 
-        // Delete cell using shared model (non-disruptive)
-        model.sharedModel.deleteCell(cellIndex);
+      if (!model) {
+        return JSON.stringify({
+          success: false,
+          error: 'No notebook model available'
+        });
+      }
 
+      if (cellIndex < 0 || cellIndex >= model.cells.length) {
         return JSON.stringify({
-          success: true,
-          message: `Cell ${cellIndex} deleted successfully`,
-          cellIndex,
-          remainingCells: model.cells.length
+          success: false,
+          error: `Invalid cell index: ${cellIndex}. Notebook has ${model.cells.length} cells.`
         });
-      } catch (error) {
+      }
+
+      // Validate cell exists
+      const targetCell = model.cells.get(cellIndex);
+      if (!targetCell) {
         return JSON.stringify({
           success: false,
-          error: `Failed to delete cell: ${(error as Error).message}`
+          error: `Cell at index ${cellIndex} not found`
         });
       }
+
+      // Delete cell using shared model (non-disruptive)
+      model.sharedModel.deleteCell(cellIndex);
+
+      return JSON.stringify({
+        success: true,
+        message: `Cell ${cellIndex} deleted successfully`,
+        cellIndex,
+        remainingCells: model.cells.length
+      });
     }
   });
 }
@@ -828,10 +783,10 @@ export function createExecuteActiveCellTool(
   notebookTracker?: INotebookTracker
 ): ITool {
   return tool({
-    name: 'execute_active_cell',
+    title: 'Execute Active Cell',
     description:
       'Execute the currently active cell in the notebook without disrupting user focus',
-    parameters: z.object({
+    inputSchema: z.object({
       notebookPath: z
         .string()
         .optional()
@@ -856,76 +811,69 @@ export function createExecuteActiveCellTool(
     }) => {
       const { notebookPath, code, recordTiming = true } = input;
 
-      try {
-        const currentWidget = await getNotebookWidget(
-          notebookPath,
-          docManager,
-          notebookTracker
-        );
-        if (!currentWidget) {
-          return JSON.stringify({
-            success: false,
-            error: notebookPath
-              ? `Failed to open notebook at path: ${notebookPath}`
-              : 'No active notebook and no notebook path provided'
-          });
-        }
-
-        const notebook = currentWidget.content;
-        const model = notebook.model;
-        const activeCellIndex = notebook.activeCellIndex;
-
-        if (!model || activeCellIndex === -1) {
-          return JSON.stringify({
-            success: false,
-            error: 'No notebook model or active cell available'
-          });
-        }
-
-        const activeCell = model.cells.get(activeCellIndex);
-        if (!activeCell) {
-          return JSON.stringify({
-            success: false,
-            error: 'Active cell not found'
-          });
-        }
-
-        // Set code content if provided
-        if (code) {
-          activeCell.sharedModel.setSource(code);
-        }
+      const currentWidget = await getNotebookWidget(
+        notebookPath,
+        docManager,
+        notebookTracker
+      );
+      if (!currentWidget) {
+        return JSON.stringify({
+          success: false,
+          error: notebookPath
+            ? `Failed to open notebook at path: ${notebookPath}`
+            : 'No active notebook and no notebook path provided'
+        });
+      }
 
-        // Get the cell widget for execution
-        const cellWidget = notebook.widgets[activeCellIndex];
-        if (!cellWidget || !(cellWidget instanceof CodeCell)) {
-          return JSON.stringify({
-            success: false,
-            error: 'Active cell is not a code cell'
-          });
-        }
+      const notebook = currentWidget.content;
+      const model = notebook.model;
+      const activeCellIndex = notebook.activeCellIndex;
 
-        // Execute using shared model approach (non-disruptive)
-        const sessionCtx = currentWidget.sessionContext;
-        await CodeCell.execute(cellWidget, sessionCtx, {
-          recordTiming,
-          deletedCells: model.deletedCells
+      if (!model || activeCellIndex === -1) {
+        return JSON.stringify({
+          success: false,
+          error: 'No notebook model or active cell available'
         });
+      }
 
-        const codeModel = cellWidget.model as ICodeCellModel;
+      const activeCell = model.cells.get(activeCellIndex);
+      if (!activeCell) {
         return JSON.stringify({
-          success: true,
-          message: 'Code executed successfully in active cell',
-          cellIndex: activeCellIndex,
-          executionCount: codeModel.executionCount,
-          hasOutput: codeModel.outputs.length > 0,
-          code: code || activeCell.sharedModel.getSource()
+          success: false,
+          error: 'Active cell not found'
         });
-      } catch (error) {
+      }
+
+      // Set code content if provided
+      if (code) {
+        activeCell.sharedModel.setSource(code);
+      }
+
+      // Get the cell widget for execution
+      const cellWidget = notebook.widgets[activeCellIndex];
+      if (!cellWidget || !(cellWidget instanceof CodeCell)) {
         return JSON.stringify({
           success: false,
-          error: `Failed to execute code: ${(error as Error).message}`
+          error: 'Active cell is not a code cell'
         });
       }
+
+      // Execute using shared model approach (non-disruptive)
+      const sessionCtx = currentWidget.sessionContext;
+      await CodeCell.execute(cellWidget, sessionCtx, {
+        recordTiming,
+        deletedCells: model.deletedCells
+      });
+
+      const codeModel = cellWidget.model as ICodeCellModel;
+      return JSON.stringify({
+        success: true,
+        message: 'Code executed successfully in active cell',
+        cellIndex: activeCellIndex,
+        executionCount: codeModel.executionCount,
+        hasOutput: codeModel.outputs.length > 0,
+        code: code || activeCell.sharedModel.getSource()
+      });
     }
   });
 }
@@ -938,9 +886,9 @@ export function createSaveNotebookTool(
   notebookTracker?: INotebookTracker
 ): ITool {
   return tool({
-    name: 'save_notebook',
+    title: 'Save Notebook',
     description: 'Save a specific notebook to disk',
-    parameters: z.object({
+    inputSchema: z.object({
       notebookPath: z
         .string()
         .optional()
@@ -952,35 +900,28 @@ export function createSaveNotebookTool(
     execute: async (input: { notebookPath?: string | null }) => {
       const { notebookPath } = input;
 
-      try {
-        const currentWidget = await getNotebookWidget(
-          notebookPath,
-          docManager,
-          notebookTracker
-        );
-        if (!currentWidget) {
-          return JSON.stringify({
-            success: false,
-            error: notebookPath
-              ? `Failed to open notebook at path: ${notebookPath}`
-              : 'No active notebook and no notebook path provided'
-          });
-        }
-
-        await currentWidget.context.save();
-
-        return JSON.stringify({
-          success: true,
-          message: 'Notebook saved successfully',
-          notebookName: currentWidget.title.label,
-          notebookPath: currentWidget.context.path
-        });
-      } catch (error) {
+      const currentWidget = await getNotebookWidget(
+        notebookPath,
+        docManager,
+        notebookTracker
+      );
+      if (!currentWidget) {
         return JSON.stringify({
           success: false,
-          error: `Failed to save notebook: ${(error as Error).message}`
+          error: notebookPath
+            ? `Failed to open notebook at path: ${notebookPath}`
+            : 'No active notebook and no notebook path provided'
         });
       }
+
+      await currentWidget.context.save();
+
+      return JSON.stringify({
+        success: true,
+        message: 'Notebook saved successfully',
+        notebookName: currentWidget.title.label,
+        notebookPath: currentWidget.context.path
+      });
     }
   });
 }