diff --git a/src/core/tools/__tests__/jupyter-notebook-handler.test.ts b/src/core/tools/__tests__/jupyter-notebook-handler.test.ts new file mode 100644 index 0000000000..76e87aefcd --- /dev/null +++ b/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())", + ]) + }) + }) +}) diff --git a/src/core/tools/applyDiffTool.ts b/src/core/tools/applyDiffTool.ts index d4f7fd883f..7f9025586c 100644 --- a/src/core/tools/applyDiffTool.ts +++ b/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() diff --git a/src/core/tools/jupyter-notebook-handler.ts b/src/core/tools/jupyter-notebook-handler.ts new file mode 100644 index 0000000000..aa046e0962 --- /dev/null +++ b/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 { + 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() + + 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 { + 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)}`, + } + } +} diff --git a/src/core/tools/multiApplyDiffTool.ts b/src/core/tools/multiApplyDiffTool.ts index a80075e10f..a5aa502e62 100644 --- a/src/core/tools/multiApplyDiffTool.ts +++ b/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 diff --git a/src/core/tools/writeToFileTool.ts b/src/core/tools/writeToFileTool.ts index d4469e9099..664d6e96a1 100644 --- a/src/core/tools/writeToFileTool.ts +++ b/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) diff --git a/test-jupyter-fix.js b/test-jupyter-fix.js new file mode 100644 index 0000000000..9532d309e0 --- /dev/null +++ b/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()