@jupyterlite/ai 0.9.1 → 0.10.0

package/lib/tools/notebook.js

@@ -1,7 +1,7 @@
 import { CodeCell, MarkdownCell } from '@jupyterlab/cells';
 import { DocumentWidget } from '@jupyterlab/docregistry';
 import { NotebookPanel } from '@jupyterlab/notebook';
-import { tool } from '@openai/agents';
+import { tool } from 'ai';
 import { z } from 'zod';
 /**
  * Find a kernel name that matches the specified language
@@ -64,9 +64,9 @@ async function getNotebookWidget(notebookPath, docManager, notebookTracker) {
  */
 export function createNotebookCreationTool(docManager, kernelSpecManager) {
     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()
@@ -82,42 +82,40 @@ export function createNotebookCreationTool(docManager, kernelSpecManager) {
             const kernel = await findKernelByLanguage(kernelSpecManager, input.language);
             const { name } = input;
             if (!name) {
-                throw new 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'
-                });
-                // 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');
-                }
-                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
+                    success: false,
+                    error: 'A name must be provided to create a notebook'
                 };
             }
-            catch (error) {
+            // TODO: handle cwd / path?
+            const fileName = name.endsWith('.ipynb') ? name : `${name}.ipynb`;
+            // Create untitled notebook first
+            const notebookModel = await docManager.newUntitled({
+                type: 'notebook'
+            });
+            // 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)) {
                 return {
                     success: false,
-                    error: `Failed to create notebook: ${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
+            };
         }
     });
 }
@@ -126,9 +124,9 @@ export function createNotebookCreationTool(docManager, kernelSpecManager) {
  */
 export function createAddCellTool(docManager, notebookTracker) {
     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()
@@ -149,67 +147,59 @@ export function createAddCellTool(docManager, notebookTracker) {
                 .default('below')
                 .describe('Position relative to current cell')
         }),
-        async execute({ 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 notebook = currentWidget.content;
-                const model = notebook.model;
-                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() === '';
-                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 } : {}
-                };
-                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);
-                        }
-                    }
-                }
+        execute: async ({ notebookPath, content, cellType = 'code', position = 'below' }) => {
+            const currentWidget = await getNotebookWidget(notebookPath, docManager, notebookTracker);
+            if (!currentWidget) {
                 return {
-                    success: true,
-                    message: `${cellType} cell added successfully`,
-                    content: content || '',
-                    cellType,
-                    position
+                    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 {
                     success: false,
-                    error: `Failed to add ${cellType} cell: ${error.message}`
+                    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() === '';
+            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 } : {}
+            };
+            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
+            };
         }
     });
 }
@@ -218,9 +208,9 @@ export function createAddCellTool(docManager, notebookTracker) {
  */
 export function createGetNotebookInfoTool(docManager, notebookTracker) {
     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()
@@ -229,44 +219,36 @@ export function createGetNotebookInfoTool(docManager, notebookTracker) {
         }),
         execute: async (input) => {
             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.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
+            });
         }
     });
 }
@@ -275,9 +257,9 @@ export function createGetNotebookInfoTool(docManager, notebookTracker) {
  */
 export function createGetCellInfoTool(docManager, notebookTracker) {
     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()
@@ -292,59 +274,51 @@ export function createGetCellInfoTool(docManager, notebookTracker) {
         execute: async (input) => {
             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();
-                // Get outputs for code cells
-                let outputs = [];
-                if (cellType === 'code') {
-                    const rawOutputs = sharedModel.toJSON().outputs;
-                    outputs = Array.isArray(rawOutputs) ? rawOutputs : [];
-                }
+            const currentWidget = await getNotebookWidget(notebookPath, docManager, notebookTracker);
+            if (!currentWidget) {
                 return JSON.stringify({
-                    success: true,
-                    cellId: cell.id,
-                    cellIndex,
-                    cellType,
-                    source,
-                    outputs,
-                    executionCount: cellType === 'code' ? cell.executionCount : null
+                    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 cell info: ${error.message}`
+                    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();
+            // Get outputs for code cells
+            let outputs = [];
+            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.executionCount : null
+            });
         }
     });
 }
@@ -353,9 +327,9 @@ export function createGetCellInfoTool(docManager, notebookTracker) {
  */
 export function createSetCellContentTool(docManager, notebookTracker, diffManager) {
     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()
@@ -375,108 +349,100 @@ export function createSetCellContentTool(docManager, notebookTracker, diffManage
         }),
         execute: async (input) => {
             const { notebookPath, cellId, cellIndex, content } = input;
-            try {
-                const notebookWidget = await getNotebookWidget(notebookPath, docManager, notebookTracker);
-                if (!notebookWidget) {
+            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 model = notebook.model;
+            if (!model) {
+                return JSON.stringify({
+                    success: false,
+                    error: 'No notebook model available'
+                });
+            }
+            // Determine target cell index
+            let targetCellIndex;
+            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: notebookPath
-                            ? `Failed to open notebook at path: ${notebookPath}`
-                            : 'No active notebook and no notebook path provided'
+                        error: `Cell with ID '${cellId}' not found in notebook`
                     });
                 }
-                const notebook = notebookWidget.content;
-                const targetNotebookPath = notebookWidget.context.path;
-                const model = notebook.model;
-                if (!model) {
+            }
+            else if (cellIndex !== undefined && cellIndex !== null) {
+                // Use provided cell index
+                if (cellIndex < 0 || cellIndex >= model.cells.length) {
                     return JSON.stringify({
                         success: false,
-                        error: 'No notebook model available'
+                        error: `Invalid cell index: ${cellIndex}. Notebook has ${model.cells.length} cells.`
                     });
                 }
-                // Determine target cell index
-                let targetCellIndex;
-                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'
-                        });
-                    }
-                }
-                // 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`
-                    });
-                }
-                const sharedModel = targetCell.sharedModel;
-                // Get previous content and type
-                const previousContent = sharedModel.getSource();
-                const previousCellType = targetCell.type;
-                const retrievedCellId = targetCell.id;
-                sharedModel.setSource(content);
-                // Show the cell diff using the diff manager if available
-                if (diffManager) {
-                    await diffManager.showCellDiff({
-                        original: previousContent,
-                        modified: content,
-                        cellId: retrievedCellId,
-                        notebookPath: targetNotebookPath
+                        error: 'No active cell or invalid active cell index'
                     });
                 }
-                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
-                });
             }
-            catch (error) {
+            // Get the target cell
+            const targetCell = model.cells.get(targetCellIndex);
+            if (!targetCell) {
                 return JSON.stringify({
                     success: false,
-                    error: `Failed to replace cell content: ${error.message}`
+                    error: `Cell at index ${targetCellIndex} not found`
+                });
+            }
+            const sharedModel = targetCell.sharedModel;
+            // Get previous content and type
+            const previousContent = sharedModel.getSource();
+            const previousCellType = targetCell.type;
+            const retrievedCellId = targetCell.id;
+            sharedModel.setSource(content);
+            // Show the cell diff using the diff manager if available
+            if (diffManager) {
+                await diffManager.showCellDiff({
+                    original: previousContent,
+                    modified: content,
+                    cellId: retrievedCellId,
+                    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
+            });
         }
     });
 }
@@ -485,9 +451,9 @@ export function createSetCellContentTool(docManager, notebookTracker, diffManage
  */
 export function createRunCellTool(docManager, notebookTracker) {
     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()
@@ -502,78 +468,61 @@ export function createRunCellTool(docManager, notebookTracker) {
         needsApproval: true,
         execute: async (input) => {
             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;
-                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.`
-                    });
-                }
-                // 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`
-                    });
-                }
-                // 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;
-                        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.message}`,
-                        cellIndex
-                    });
-                }
+            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'
+                });
             }
-            catch (error) {
+            if (cellIndex < 0 || cellIndex >= model.cells.length) {
                 return JSON.stringify({
                     success: false,
-                    error: `Failed to run cell: ${error.message}`
+                    error: `Invalid cell index: ${cellIndex}. Notebook has ${model.cells.length} cells.`
+                });
+            }
+            // 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`
+                });
+            }
+            // 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;
+                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
                 });
             }
         }
@@ -584,9 +533,9 @@ export function createRunCellTool(docManager, notebookTracker) {
  */
 export function createDeleteCellTool(docManager, notebookTracker) {
     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()
@@ -596,53 +545,45 @@ export function createDeleteCellTool(docManager, notebookTracker) {
         }),
         execute: async (input) => {
             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.`
-                    });
-                }
-                // Validate cell exists
-                const targetCell = model.cells.get(cellIndex);
-                if (!targetCell) {
-                    return JSON.stringify({
-                        success: false,
-                        error: `Cell at index ${cellIndex} not found`
-                    });
-                }
-                // Delete cell using shared model (non-disruptive)
-                model.sharedModel.deleteCell(cellIndex);
+            const currentWidget = await getNotebookWidget(notebookPath, docManager, notebookTracker);
+            if (!currentWidget) {
                 return JSON.stringify({
-                    success: true,
-                    message: `Cell ${cellIndex} deleted successfully`,
-                    cellIndex,
-                    remainingCells: model.cells.length
+                    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'
                 });
             }
-            catch (error) {
+            if (cellIndex < 0 || cellIndex >= model.cells.length) {
                 return JSON.stringify({
                     success: false,
-                    error: `Failed to delete cell: ${error.message}`
+                    error: `Invalid cell index: ${cellIndex}. Notebook has ${model.cells.length} cells.`
                 });
             }
+            // Validate cell exists
+            const targetCell = model.cells.get(cellIndex);
+            if (!targetCell) {
+                return JSON.stringify({
+                    success: false,
+                    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
+            });
         }
     });
 }
@@ -651,9 +592,9 @@ export function createDeleteCellTool(docManager, notebookTracker) {
  */
 export function createExecuteActiveCellTool(docManager, notebookTracker) {
     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()
@@ -671,66 +612,58 @@ export function createExecuteActiveCellTool(docManager, notebookTracker) {
         }),
         execute: async (input) => {
             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);
-                }
-                // 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'
-                    });
-                }
-                // Execute using shared model approach (non-disruptive)
-                const sessionCtx = currentWidget.sessionContext;
-                await CodeCell.execute(cellWidget, sessionCtx, {
-                    recordTiming,
-                    deletedCells: model.deletedCells
+            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 codeModel = cellWidget.model;
+            }
+            const notebook = currentWidget.content;
+            const model = notebook.model;
+            const activeCellIndex = notebook.activeCellIndex;
+            if (!model || activeCellIndex === -1) {
                 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: 'No notebook model or active cell available'
                 });
             }
-            catch (error) {
+            const activeCell = model.cells.get(activeCellIndex);
+            if (!activeCell) {
                 return JSON.stringify({
                     success: false,
-                    error: `Failed to execute code: ${error.message}`
+                    error: 'Active cell not found'
                 });
             }
+            // 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: '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;
+            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()
+            });
         }
     });
 }
@@ -739,9 +672,9 @@ export function createExecuteActiveCellTool(docManager, notebookTracker) {
  */
 export function createSaveNotebookTool(docManager, notebookTracker) {
     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()
@@ -750,30 +683,22 @@ export function createSaveNotebookTool(docManager, notebookTracker) {
         }),
         execute: async (input) => {
             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.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
+            });
         }
     });
 }