Fixes #4956: Add proper Jupyter notebook support for apply_diff and write_to_file tools

- Created jupyter-notebook-handler.ts with utilities for parsing and handling .ipynb files
- Modified applyDiffTool.ts to extract cell content for editing and reconstruct notebook JSON
- Modified multiApplyDiffTool.ts to handle Jupyter notebooks in batch operations
- Modified writeToFileTool.ts to detect and properly handle notebook content vs raw JSON
- Added comprehensive test suite for all notebook operations
- Prevents notebook corruption by maintaining proper JSON structure
- Supports both extracted content editing and direct JSON editing workflows

src/core/tools/__tests__/jupyter-notebook-handler.test.ts

@@ -0,0 +1,370 @@
+import { describe, it, expect, beforeEach, afterEach } from "vitest"
+import fs from "fs/promises"
+import path from "path"
+import {
+	isJupyterNotebook,
+	parseJupyterNotebook,
+	applyChangesToNotebook,
+	writeJupyterNotebook,
+	validateJupyterNotebookJson,
+} from "../jupyter-notebook-handler"
+
+describe("Jupyter Notebook Handler", () => {
+	const testDir = path.join(__dirname, "test-notebooks")
+	const testNotebookPath = path.join(testDir, "test.ipynb")
+	const testTextPath = path.join(testDir, "test.txt")
+
+	beforeEach(async () => {
+		await fs.mkdir(testDir, { recursive: true })
+	})
+
+	afterEach(async () => {
+		try {
+			await fs.rm(testDir, { recursive: true, force: true })
+		} catch (error) {
+			// Ignore cleanup errors
+		}
+	})
+
+	describe("isJupyterNotebook", () => {
+		it("should return true for .ipynb files", () => {
+			expect(isJupyterNotebook("test.ipynb")).toBe(true)
+			expect(isJupyterNotebook("/path/to/notebook.ipynb")).toBe(true)
+			expect(isJupyterNotebook("NOTEBOOK.IPYNB")).toBe(true)
+		})
+
+		it("should return false for non-.ipynb files", () => {
+			expect(isJupyterNotebook("test.py")).toBe(false)
+			expect(isJupyterNotebook("test.txt")).toBe(false)
+			expect(isJupyterNotebook("test")).toBe(false)
+			expect(isJupyterNotebook("test.ipynb.backup")).toBe(false)
+		})
+	})
+
+	describe("validateJupyterNotebookJson", () => {
+		it("should validate correct notebook JSON", () => {
+			const validNotebook = JSON.stringify({
+				cells: [],
+				metadata: {},
+				nbformat: 4,
+				nbformat_minor: 2,
+			})
+
+			const result = validateJupyterNotebookJson(validNotebook)
+			expect(result.valid).toBe(true)
+			expect(result.error).toBeUndefined()
+		})
+
+		it("should reject invalid JSON", () => {
+			const result = validateJupyterNotebookJson("invalid json")
+			expect(result.valid).toBe(false)
+			expect(result.error).toContain("Invalid JSON")
+		})
+
+		it("should reject JSON without cells", () => {
+			const invalidNotebook = JSON.stringify({
+				metadata: {},
+				nbformat: 4,
+				nbformat_minor: 2,
+			})
+
+			const result = validateJupyterNotebookJson(invalidNotebook)
+			expect(result.valid).toBe(false)
+			expect(result.error).toContain("Missing or invalid 'cells' array")
+		})
+
+		it("should reject JSON without nbformat", () => {
+			const invalidNotebook = JSON.stringify({
+				cells: [],
+				metadata: {},
+				nbformat_minor: 2,
+			})
+
+			const result = validateJupyterNotebookJson(invalidNotebook)
+			expect(result.valid).toBe(false)
+			expect(result.error).toContain("Missing or invalid 'nbformat'")
+		})
+	})
+
+	describe("parseJupyterNotebook", () => {
+		it("should return isNotebook false for non-notebook files", async () => {
+			await fs.writeFile(testTextPath, "Hello world")
+			const result = await parseJupyterNotebook(testTextPath)
+			expect(result.isNotebook).toBe(false)
+		})
+
+		it("should parse a simple notebook with code and markdown cells", async () => {
+			const notebook = {
+				cells: [
+					{
+						cell_type: "markdown",
+						source: ["# Hello World\n", "This is a markdown cell."],
+					},
+					{
+						cell_type: "code",
+						source: ["print('Hello, World!')\n", "x = 42"],
+					},
+					{
+						cell_type: "raw",
+						source: ["This is raw text"],
+					},
+				],
+				metadata: {},
+				nbformat: 4,
+				nbformat_minor: 2,
+			}
+
+			await fs.writeFile(testNotebookPath, JSON.stringify(notebook, null, 2))
+			const result = await parseJupyterNotebook(testNotebookPath)
+
+			expect(result.isNotebook).toBe(true)
+			expect(result.originalJson).toEqual(notebook)
+			expect(result.extractedContent).toBe(
+				"# Hello World\nThis is a markdown cell.\nprint('Hello, World!')\nx = 42",
+			)
+			expect(result.cellBoundaries).toHaveLength(2)
+			expect(result.cellBoundaries![0]).toEqual({
+				cellIndex: 0,
+				startLine: 1,
+				endLine: 2,
+				cellType: "markdown",
+			})
+			expect(result.cellBoundaries![1]).toEqual({
+				cellIndex: 1,
+				startLine: 3,
+				endLine: 4,
+				cellType: "code",
+			})
+		})
+
+		it("should handle empty cells", async () => {
+			const notebook = {
+				cells: [
+					{
+						cell_type: "code",
+						source: [],
+					},
+					{
+						cell_type: "markdown",
+						source: ["# Title"],
+					},
+				],
+				metadata: {},
+				nbformat: 4,
+				nbformat_minor: 2,
+			}
+
+			await fs.writeFile(testNotebookPath, JSON.stringify(notebook, null, 2))
+			const result = await parseJupyterNotebook(testNotebookPath)
+
+			expect(result.isNotebook).toBe(true)
+			expect(result.extractedContent).toBe("# Title")
+			expect(result.cellBoundaries).toHaveLength(1)
+			expect(result.cellBoundaries![0]).toEqual({
+				cellIndex: 1,
+				startLine: 1,
+				endLine: 1,
+				cellType: "markdown",
+			})
+		})
+
+		it("should throw error for invalid JSON", async () => {
+			await fs.writeFile(testNotebookPath, "invalid json")
+			await expect(parseJupyterNotebook(testNotebookPath)).rejects.toThrow("Failed to parse Jupyter notebook")
+		})
+	})
+
+	describe("applyChangesToNotebook", () => {
+		it("should apply changes to notebook cells", () => {
+			const originalNotebook = {
+				cells: [
+					{
+						cell_type: "markdown" as const,
+						source: ["# Old Title\n", "Old content."],
+					},
+					{
+						cell_type: "code" as const,
+						source: ["print('old')\n", "x = 1"],
+					},
+				],
+				metadata: {},
+				nbformat: 4,
+				nbformat_minor: 2,
+			}
+
+			const cellBoundaries = [
+				{
+					cellIndex: 0,
+					startLine: 1,
+					endLine: 2,
+					cellType: "markdown",
+				},
+				{
+					cellIndex: 1,
+					startLine: 3,
+					endLine: 4,
+					cellType: "code",
+				},
+			]
+
+			const newExtractedContent = "# New Title\nNew content.\nprint('new')\nx = 2"
+
+			const result = applyChangesToNotebook(originalNotebook, newExtractedContent, cellBoundaries)
+
+			expect(result.cells[0].source).toEqual(["# New Title\n", "New content."])
+			expect(result.cells[1].source).toEqual(["print('new')\n", "x = 2"])
+		})
+
+		it("should handle single-line cells", () => {
+			const originalNotebook = {
+				cells: [
+					{
+						cell_type: "code" as const,
+						source: ["print('hello')"],
+					},
+				],
+				metadata: {},
+				nbformat: 4,
+				nbformat_minor: 2,
+			}
+
+			const cellBoundaries = [
+				{
+					cellIndex: 0,
+					startLine: 1,
+					endLine: 1,
+					cellType: "code",
+				},
+			]
+
+			const newExtractedContent = "print('world')"
+
+			const result = applyChangesToNotebook(originalNotebook, newExtractedContent, cellBoundaries)
+
+			expect(result.cells[0].source).toEqual(["print('world')"])
+		})
+
+		it("should preserve cells not in boundaries", () => {
+			const originalNotebook = {
+				cells: [
+					{
+						cell_type: "markdown" as const,
+						source: ["# Title"],
+					},
+					{
+						cell_type: "raw" as const,
+						source: ["Raw content"],
+					},
+					{
+						cell_type: "code" as const,
+						source: ["print('code')"],
+					},
+				],
+				metadata: {},
+				nbformat: 4,
+				nbformat_minor: 2,
+			}
+
+			const cellBoundaries = [
+				{
+					cellIndex: 0,
+					startLine: 1,
+					endLine: 1,
+					cellType: "markdown",
+				},
+				{
+					cellIndex: 2,
+					startLine: 2,
+					endLine: 2,
+					cellType: "code",
+				},
+			]
+
+			const newExtractedContent = "# New Title\nprint('new code')"
+
+			const result = applyChangesToNotebook(originalNotebook, newExtractedContent, cellBoundaries)
+
+			expect(result.cells[0].source).toEqual(["# New Title"])
+			expect(result.cells[1].source).toEqual(["Raw content"]) // Unchanged
+			expect(result.cells[2].source).toEqual(["print('new code')"])
+		})
+	})
+
+	describe("writeJupyterNotebook", () => {
+		it("should write notebook with proper formatting", async () => {
+			const notebook = {
+				cells: [
+					{
+						cell_type: "code" as const,
+						source: ["print('test')"],
+					},
+				],
+				metadata: {},
+				nbformat: 4,
+				nbformat_minor: 2,
+			}
+
+			await writeJupyterNotebook(testNotebookPath, notebook)
+
+			const writtenContent = await fs.readFile(testNotebookPath, "utf8")
+			const parsedContent = JSON.parse(writtenContent)
+
+			expect(parsedContent).toEqual(notebook)
+			// Check that it's properly formatted (indented)
+			expect(writtenContent).toContain('  "cells":')
+		})
+	})
+
+	describe("integration test", () => {
+		it("should handle full parse -> modify -> apply cycle", async () => {
+			const originalNotebook = {
+				cells: [
+					{
+						cell_type: "markdown" as const,
+						source: ["# Data Analysis\n", "Let's analyze some data."],
+					},
+					{
+						cell_type: "code" as const,
+						source: ["import pandas as pd\n", "df = pd.read_csv('data.csv')\n", "print(df.head())"],
+					},
+				],
+				metadata: {},
+				nbformat: 4,
+				nbformat_minor: 2,
+			}
+
+			// Write original notebook
+			await fs.writeFile(testNotebookPath, JSON.stringify(originalNotebook, null, 2))
+
+			// Parse it
+			const parseResult = await parseJupyterNotebook(testNotebookPath)
+			expect(parseResult.isNotebook).toBe(true)
+
+			// Modify the extracted content
+			const modifiedContent =
+				"# Advanced Data Analysis\nLet's do advanced analysis.\nimport pandas as pd\nimport numpy as np\ndf = pd.read_csv('data.csv')\nprint(df.describe())"
+
+			// Apply changes back
+			const updatedNotebook = applyChangesToNotebook(
+				parseResult.originalJson!,
+				modifiedContent,
+				parseResult.cellBoundaries!,
+			)
+
+			// Write it back
+			await writeJupyterNotebook(testNotebookPath, updatedNotebook)
+
+			// Verify the result
+			const finalContent = await fs.readFile(testNotebookPath, "utf8")
+			const finalNotebook = JSON.parse(finalContent)
+
+			expect(finalNotebook.cells[0].source).toEqual(["# Advanced Data Analysis\n", "Let's do advanced analysis."])
+			expect(finalNotebook.cells[1].source).toEqual([
+				"import pandas as pd\n",
+				"import numpy as np\n",
+				"df = pd.read_csv('data.csv')\n",
+				"print(df.describe())",
+			])
+		})
+	})
+})

src/core/tools/applyDiffTool.ts

@@ -11,6 +11,12 @@ import { formatResponse } from "../prompts/responses"
 import { fileExistsAtPath } from "../../utils/fs"
 import { RecordSource } from "../context-tracking/FileContextTrackerTypes"
 import { unescapeHtmlEntities } from "../../utils/text-normalization"
+import {
+	isJupyterNotebook,
+	parseJupyterNotebook,
+	applyChangesToNotebook,
+	writeJupyterNotebook,
+} from "./jupyter-notebook-handler"
 
 export async function applyDiffToolLegacy(
 	cline: Task,
@@ -86,7 +92,33 @@ export async function applyDiffToolLegacy(
 				return
 			}
 
-			let originalContent: string | null = await fs.readFile(absolutePath, "utf-8")
+			let originalContent: string | null
+			let isNotebook = false
+			let notebookData: any = null
+
+			// Handle Jupyter notebooks specially
+			if (isJupyterNotebook(absolutePath)) {
+				try {
+					const parseResult = await parseJupyterNotebook(absolutePath)
+					if (parseResult.isNotebook && parseResult.extractedContent !== undefined) {
+						originalContent = parseResult.extractedContent
+						isNotebook = true
+						notebookData = {
+							originalJson: parseResult.originalJson,
+							cellBoundaries: parseResult.cellBoundaries,
+						}
+					} else {
+						// Fallback to raw file content if parsing fails
+						originalContent = await fs.readFile(absolutePath, "utf-8")
+					}
+				} catch (error) {
+					// If notebook parsing fails, treat as regular file but warn
+					console.warn(`Failed to parse Jupyter notebook ${absolutePath}, treating as regular file:`, error)
+					originalContent = await fs.readFile(absolutePath, "utf-8")
+				}
+			} else {
+				originalContent = await fs.readFile(absolutePath, "utf-8")
+			}
 
 			// Apply the diff to the original content
 			const diffResult = (await cline.diffStrategy?.applyDiff(
@@ -165,6 +197,31 @@ export async function applyDiffToolLegacy(
 				return
 			}
 
+			// Handle saving for Jupyter notebooks
+			if (isNotebook && notebookData && diffResult.content) {
+				try {
+					// Apply changes back to the notebook structure
+					const updatedNotebook = applyChangesToNotebook(
+						notebookData.originalJson,
+						diffResult.content,
+						notebookData.cellBoundaries,
+					)
+
+					// Write the updated notebook
+					await writeJupyterNotebook(absolutePath, updatedNotebook)
+
+					// Update diff view with the notebook JSON for display
+					await cline.diffViewProvider.update(JSON.stringify(updatedNotebook, null, 2), true)
+				} catch (error) {
+					const errorMsg = `Failed to save Jupyter notebook: ${error instanceof Error ? error.message : String(error)}`
+					cline.consecutiveMistakeCount++
+					cline.recordToolError("apply_diff", errorMsg)
+					pushToolResult(errorMsg)
+					await cline.diffViewProvider.reset()
+					return
+				}
+			}
+
 			// Call saveChanges to update the DiffViewProvider properties
 			await cline.diffViewProvider.saveChanges()
 

src/core/tools/jupyter-notebook-handler.ts

@@ -0,0 +1,171 @@
+import fs from "fs/promises"
+import path from "path"
+
+/**
+ * Jupyter notebook cell interface
+ */
+interface JupyterCell {
+	cell_type: "code" | "markdown" | "raw"
+	source: string[]
+	metadata?: any
+	outputs?: any[]
+	execution_count?: number | null
+}
+
+/**
+ * Jupyter notebook interface
+ */
+interface JupyterNotebook {
+	cells: JupyterCell[]
+	metadata: any
+	nbformat: number
+	nbformat_minor: number
+}
+
+/**
+ * Result of parsing a Jupyter notebook for editing
+ */
+interface NotebookParseResult {
+	isNotebook: boolean
+	originalJson?: JupyterNotebook
+	extractedContent?: string
+	cellBoundaries?: Array<{
+		cellIndex: number
+		startLine: number
+		endLine: number
+		cellType: string
+	}>
+}
+
+/**
+ * Checks if a file is a Jupyter notebook based on its extension
+ */
+export function isJupyterNotebook(filePath: string): boolean {
+	return path.extname(filePath).toLowerCase() === ".ipynb"
+}
+
+/**
+ * Parses a Jupyter notebook file and extracts content in a format suitable for editing
+ */
+export async function parseJupyterNotebook(filePath: string): Promise<NotebookParseResult> {
+	if (!isJupyterNotebook(filePath)) {
+		return { isNotebook: false }
+	}
+
+	try {
+		const data = await fs.readFile(filePath, "utf8")
+		const notebook: JupyterNotebook = JSON.parse(data)
+
+		let extractedContent = ""
+		const cellBoundaries: Array<{
+			cellIndex: number
+			startLine: number
+			endLine: number
+			cellType: string
+		}> = []
+
+		let currentLine = 1
+
+		for (let i = 0; i < notebook.cells.length; i++) {
+			const cell = notebook.cells[i]
+			if ((cell.cell_type === "markdown" || cell.cell_type === "code") && cell.source) {
+				const cellContent = cell.source.join("\n")
+				const startLine = currentLine
+				const lines = cellContent.split("\n")
+				const endLine = currentLine + lines.length - 1
+
+				cellBoundaries.push({
+					cellIndex: i,
+					startLine,
+					endLine,
+					cellType: cell.cell_type,
+				})
+
+				extractedContent += cellContent + "\n"
+				currentLine = endLine + 1
+			}
+		}
+
+		return {
+			isNotebook: true,
+			originalJson: notebook,
+			extractedContent: extractedContent.trimEnd(),
+			cellBoundaries,
+		}
+	} catch (error) {
+		throw new Error(`Failed to parse Jupyter notebook: ${error instanceof Error ? error.message : String(error)}`)
+	}
+}
+
+/**
+ * Applies changes to extracted content back to the original notebook structure
+ */
+export function applyChangesToNotebook(
+	originalNotebook: JupyterNotebook,
+	newExtractedContent: string,
+	cellBoundaries: Array<{
+		cellIndex: number
+		startLine: number
+		endLine: number
+		cellType: string
+	}>,
+): JupyterNotebook {
+	const newNotebook: JupyterNotebook = JSON.parse(JSON.stringify(originalNotebook))
+	const newLines = newExtractedContent.split("\n")
+
+	// Clear all existing cell sources for cells that were in the boundaries
+	const processedCellIndices = new Set<number>()
+
+	for (const boundary of cellBoundaries) {
+		processedCellIndices.add(boundary.cellIndex)
+		// Extract the lines for this cell (1-based to 0-based conversion)
+		const cellLines = newLines.slice(boundary.startLine - 1, boundary.endLine)
+
+		// Update the cell source
+		if (newNotebook.cells[boundary.cellIndex]) {
+			newNotebook.cells[boundary.cellIndex].source = cellLines.map((line, index) => {
+				// Add newline to all lines except the last one in the cell
+				return index === cellLines.length - 1 ? line : line + "\n"
+			})
+		}
+	}
+
+	return newNotebook
+}
+
+/**
+ * Writes a Jupyter notebook back to disk with proper formatting
+ */
+export async function writeJupyterNotebook(filePath: string, notebook: JupyterNotebook): Promise<void> {
+	const jsonContent = JSON.stringify(notebook, null, 2)
+	await fs.writeFile(filePath, jsonContent, "utf8")
+}
+
+/**
+ * Validates that a string is valid JSON for a Jupyter notebook
+ */
+export function validateJupyterNotebookJson(content: string): { valid: boolean; error?: string } {
+	try {
+		const parsed = JSON.parse(content)
+
+		// Basic validation for Jupyter notebook structure
+		if (!parsed.cells || !Array.isArray(parsed.cells)) {
+			return { valid: false, error: "Missing or invalid 'cells' array" }
+		}
+
+		if (typeof parsed.nbformat !== "number") {
+			return { valid: false, error: "Missing or invalid 'nbformat'" }
+		}
+
+		if (typeof parsed.nbformat_minor !== "number") {
+			return { valid: false, error: "Missing or invalid 'nbformat_minor'" }
+		}
+
+		return { valid: true }
+	} catch (error) {
+		return {
+			valid: false,
+			error: `Invalid JSON: ${error instanceof Error ? error.message : String(error)}`,
+		}
+	}
+}

src/core/tools/multiApplyDiffTool.ts

@@ -14,6 +14,12 @@ import { unescapeHtmlEntities } from "../../utils/text-normalization"
 import { parseXml } from "../../utils/xml"
 import { EXPERIMENT_IDS, experiments } from "../../shared/experiments"
 import { applyDiffToolLegacy } from "./applyDiffTool"
+import {
+	isJupyterNotebook,
+	parseJupyterNotebook,
+	applyChangesToNotebook,
+	writeJupyterNotebook,
+} from "./jupyter-notebook-handler"
 
 interface DiffOperation {
 	path: string
@@ -401,7 +407,37 @@ Original error: ${errorMessage}`
 			const fileExists = opResult.fileExists!
 
 			try {
-				let originalContent: string | null = await fs.readFile(absolutePath, "utf-8")
+				// Handle Jupyter notebooks specially
+				let originalContent: string | null
+				let isNotebook = false
+				let notebookData: any = null
+
+				if (isJupyterNotebook(absolutePath)) {
+					try {
+						const parseResult = await parseJupyterNotebook(absolutePath)
+						if (parseResult.isNotebook && parseResult.extractedContent !== undefined) {
+							originalContent = parseResult.extractedContent
+							isNotebook = true
+							notebookData = {
+								originalJson: parseResult.originalJson,
+								cellBoundaries: parseResult.cellBoundaries,
+							}
+						} else {
+							// Fallback to raw file content if parsing fails
+							originalContent = await fs.readFile(absolutePath, "utf-8")
+						}
+					} catch (error) {
+						// If notebook parsing fails, treat as regular file but warn
+						console.warn(
+							`Failed to parse Jupyter notebook ${absolutePath}, treating as regular file:`,
+							error,
+						)
+						originalContent = await fs.readFile(absolutePath, "utf-8")
+					}
+				} else {
+					originalContent = await fs.readFile(absolutePath, "utf-8")
+				}
+
 				let successCount = 0
 				let formattedError = ""
 
@@ -419,6 +455,30 @@ Original error: ${errorMessage}`
 					error: "No diff strategy available - please ensure a valid diff strategy is configured",
 				}
 
+				// Handle saving for Jupyter notebooks
+				if (diffResult.success && isNotebook && notebookData && diffResult.content) {
+					try {
+						// Apply changes back to the notebook structure
+						const updatedNotebook = applyChangesToNotebook(
+							notebookData.originalJson,
+							diffResult.content,
+							notebookData.cellBoundaries,
+						)
+
+						// Write the updated notebook
+						await writeJupyterNotebook(absolutePath, updatedNotebook)
+
+						// Update the diff result content to show the notebook JSON for display
+						diffResult.content = JSON.stringify(updatedNotebook, null, 2)
+					} catch (error) {
+						const errorMsg = `Failed to save Jupyter notebook: ${error instanceof Error ? error.message : String(error)}`
+						cline.consecutiveMistakeCount++
+						cline.recordToolError("apply_diff", errorMsg)
+						results.push(errorMsg)
+						continue
+					}
+				}
+
 				// Release the original content from memory as it's no longer needed
 				originalContent = null
 

src/core/tools/writeToFileTool.ts

@@ -13,6 +13,13 @@ import { getReadablePath } from "../../utils/path"
 import { isPathOutsideWorkspace } from "../../utils/pathUtils"
 import { detectCodeOmission } from "../../integrations/editor/detect-omission"
 import { unescapeHtmlEntities } from "../../utils/text-normalization"
+import {
+	isJupyterNotebook,
+	parseJupyterNotebook,
+	applyChangesToNotebook,
+	writeJupyterNotebook,
+	validateJupyterNotebookJson,
+} from "./jupyter-notebook-handler"
 
 export async function writeToFileTool(
 	cline: Task,
@@ -148,6 +155,68 @@ export async function writeToFileTool(
 
 			cline.consecutiveMistakeCount = 0
 
+			// Handle Jupyter notebooks specially
+			const absolutePath = path.resolve(cline.cwd, relPath)
+			let isNotebook = false
+			let notebookData: any = null
+			let processedContent = newContent
+
+			if (isJupyterNotebook(absolutePath)) {
+				// Check if the content is raw JSON (user is trying to edit the notebook structure directly)
+				const jsonValidation = validateJupyterNotebookJson(newContent)
+
+				if (jsonValidation.valid) {
+					// Content is valid notebook JSON, use it directly
+					processedContent = newContent
+					isNotebook = true
+				} else {
+					// Content is not valid JSON, treat as extracted content and convert back to notebook
+					try {
+						if (fileExists) {
+							const parseResult = await parseJupyterNotebook(absolutePath)
+							if (parseResult.isNotebook && parseResult.originalJson && parseResult.cellBoundaries) {
+								// Apply the extracted content changes back to the notebook structure
+								const updatedNotebook = applyChangesToNotebook(
+									parseResult.originalJson,
+									newContent,
+									parseResult.cellBoundaries,
+								)
+								processedContent = JSON.stringify(updatedNotebook, null, 2)
+								isNotebook = true
+								notebookData = {
+									originalJson: parseResult.originalJson,
+									cellBoundaries: parseResult.cellBoundaries,
+									wasExtractedContent: true,
+								}
+							}
+						} else {
+							// New notebook file - provide guidance
+							cline.consecutiveMistakeCount++
+							cline.recordToolError("write_to_file")
+
+							pushToolResult(
+								formatResponse.toolError(
+									`Cannot create new Jupyter notebook from extracted content. For new .ipynb files, please provide valid JSON in Jupyter notebook format, or use a different file extension for plain text content.`,
+								),
+							)
+							await cline.diffViewProvider.revertChanges()
+							return
+						}
+					} catch (error) {
+						cline.consecutiveMistakeCount++
+						cline.recordToolError("write_to_file")
+
+						pushToolResult(
+							formatResponse.toolError(
+								`Failed to process Jupyter notebook content: ${error instanceof Error ? error.message : String(error)}. Please ensure the content is either valid Jupyter notebook JSON or use apply_diff for targeted changes.`,
+							),
+						)
+						await cline.diffViewProvider.revertChanges()
+						return
+					}
+				}
+			}
+
 			// if isEditingFile false, that means we have the full contents of the file already.
 			// it's important to note how cline function works, you can't make the assumption that the block.partial conditional will always be called since it may immediately get complete, non-partial data. So cline part of the logic will always be called.
 			// in other words, you must always repeat the block.partial logic here
@@ -159,49 +228,65 @@ export async function writeToFileTool(
 			}
 
 			await cline.diffViewProvider.update(
-				everyLineHasLineNumbers(newContent) ? stripLineNumbers(newContent) : newContent,
+				everyLineHasLineNumbers(processedContent) ? stripLineNumbers(processedContent) : processedContent,
 				true,
 			)
 
 			await delay(300) // wait for diff view to update
 			cline.diffViewProvider.scrollToFirstDiff()
 
-			// Check for code omissions before proceeding
-			if (detectCodeOmission(cline.diffViewProvider.originalContent || "", newContent, predictedLineCount)) {
-				if (cline.diffStrategy) {
-					await cline.diffViewProvider.revertChanges()
-
-					pushToolResult(
-						formatResponse.toolError(
-							`Content appears to be truncated (file has ${
-								newContent.split("\n").length
-							} lines but was predicted to have ${predictedLineCount} lines), and found comments indicating omitted code (e.g., '// rest of code unchanged', '/* previous code */'). Please provide the complete file content without any omissions if possible, or otherwise use the 'apply_diff' tool to apply the diff to the original file.`,
-						),
+			// Check for code omissions before proceeding (skip for notebooks with extracted content)
+			if (!isNotebook || !notebookData?.wasExtractedContent) {
+				if (
+					detectCodeOmission(
+						cline.diffViewProvider.originalContent || "",
+						processedContent,
+						predictedLineCount,
 					)
-					return
-				} else {
-					vscode.window
-						.showWarningMessage(
-							"Potential code truncation detected. cline happens when the AI reaches its max output limit.",
-							"Follow cline guide to fix the issue",
+				) {
+					if (cline.diffStrategy) {
+						await cline.diffViewProvider.revertChanges()
+
+						pushToolResult(
+							formatResponse.toolError(
+								`Content appears to be truncated (file has ${
+									processedContent.split("\n").length
+								} lines but was predicted to have ${predictedLineCount} lines), and found comments indicating omitted code (e.g., '// rest of code unchanged', '/* previous code */'). Please provide the complete file content without any omissions if possible, or otherwise use the 'apply_diff' tool to apply the diff to the original file.`,
+							),
 						)
-						.then((selection) => {
-							if (selection === "Follow cline guide to fix the issue") {
-								vscode.env.openExternal(
-									vscode.Uri.parse(
-										"https://github.com/cline/cline/wiki/Troubleshooting-%E2%80%90-Cline-Deleting-Code-with-%22Rest-of-Code-Here%22-Comments",
-									),
-								)
-							}
-						})
+						return
+					} else {
+						vscode.window
+							.showWarningMessage(
+								"Potential code truncation detected. cline happens when the AI reaches its max output limit.",
+								"Follow cline guide to fix the issue",
+							)
+							.then((selection) => {
+								if (selection === "Follow cline guide to fix the issue") {
+									vscode.env.openExternal(
+										vscode.Uri.parse(
+											"https://github.com/cline/cline/wiki/Troubleshooting-%E2%80%90-Cline-Deleting-Code-with-%22Rest-of-Code-Here%22-Comments",
+										),
+									)
+								}
+							})
+					}
 				}
 			}
 
 			const completeMessage = JSON.stringify({
 				...sharedMessageProps,
-				content: fileExists ? undefined : newContent,
+				content: fileExists
+					? undefined
+					: isNotebook && notebookData?.wasExtractedContent
+						? newContent
+						: processedContent,
 				diff: fileExists
-					? formatResponse.createPrettyPatch(relPath, cline.diffViewProvider.originalContent, newContent)
+					? formatResponse.createPrettyPatch(
+							relPath,
+							cline.diffViewProvider.originalContent,
+							processedContent,
+						)
 					: undefined,
 			} satisfies ClineSayTool)
 

test-jupyter-fix.js

@@ -0,0 +1,41 @@
+const fs = require("fs").promises
+const path = require("path")
+
+// Simple test to verify the Jupyter notebook handler works
+async function testJupyterHandler() {
+	try {
+		// Import the handler functions
+		const {
+			isJupyterNotebook,
+			parseJupyterNotebook,
+			applyChangesToNotebook,
+			writeJupyterNotebook,
+			validateJupyterNotebookJson,
+		} = require("./src/core/tools/jupyter-notebook-handler.ts")
+
+		console.log("✓ Jupyter notebook handler imported successfully")
+
+		// Test isJupyterNotebook
+		console.log("Testing isJupyterNotebook...")
+		console.log("test.ipynb:", isJupyterNotebook("test.ipynb")) // should be true
+		console.log("test.py:", isJupyterNotebook("test.py")) // should be false
+
+		// Test validateJupyterNotebookJson
+		console.log("Testing validateJupyterNotebookJson...")
+		const validNotebook = JSON.stringify({
+			cells: [],
+			metadata: {},
+			nbformat: 4,
+			nbformat_minor: 2,
+		})
+		const validation = validateJupyterNotebookJson(validNotebook)
+		console.log("Valid notebook validation:", validation)
+
+		console.log("✓ All basic tests passed")
+	} catch (error) {
+		console.error("✗ Test failed:", error.message)
+		process.exit(1)
+	}
+}
+
+testJupyterHandler()