diff --git a/src/core/diff/strategies/jupyter-notebook-diff.ts b/src/core/diff/strategies/jupyter-notebook-diff.ts
new file mode 100644
index 0000000000..031c230a63
--- /dev/null
+++ b/src/core/diff/strategies/jupyter-notebook-diff.ts
@@ -0,0 +1,275 @@
+import { DiffStrategy, DiffResult, ToolUse } from "../../../shared/tools"
+import { ToolProgressStatus } from "@roo-code/types"
+import { JupyterNotebookHandler } from "../../../integrations/misc/jupyter-notebook-handler"
+import { MultiSearchReplaceDiffStrategy } from "./multi-search-replace"
+
+export class JupyterNotebookDiffStrategy implements DiffStrategy {
+ private fallbackStrategy: MultiSearchReplaceDiffStrategy
+
+ constructor(fuzzyThreshold?: number, bufferLines?: number) {
+ // Use MultiSearchReplaceDiffStrategy as fallback for non-cell operations
+ this.fallbackStrategy = new MultiSearchReplaceDiffStrategy(fuzzyThreshold, bufferLines)
+ }
+
+ getName(): string {
+ return "JupyterNotebookDiff"
+ }
+
+ getToolDescription(args: { cwd: string; toolOptions?: { [key: string]: string } }): string {
+ return `## apply_diff (Jupyter Notebook Support)
+Description: Request to apply PRECISE, TARGETED modifications to Jupyter notebook (.ipynb) files. This tool supports both cell-level operations and content-level changes within cells.
+
+For Jupyter notebooks, you can:
+1. Edit specific cells by cell number
+2. Add new cells
+3. Delete cells
+4. Apply standard search/replace within cells
+
+Parameters:
+- path: (required) The path of the file to modify (relative to the current workspace directory ${args.cwd})
+- diff: (required) The search/replace block or cell operation defining the changes.
+
+Cell-Level Operations Format:
+\`\`\`
+<<<<<<< CELL_OPERATION
+:operation: [edit|add|delete]
+:cell_index: [cell number, 0-based]
+:cell_type: [code|markdown|raw] (required for 'add' operation)
+-------
+[content for edit/add operations]
+=======
+[new content for edit operations, empty for delete]
+>>>>>>> CELL_OPERATION
+\`\`\`
+
+Standard Diff Format (for content within cells):
+\`\`\`
+<<<<<<< SEARCH
+:cell_index: [optional cell number to limit search]
+:start_line: [optional line number within cell]
+-------
+[exact content to find]
+=======
+[new content to replace with]
+>>>>>>> REPLACE
+\`\`\`
+
+Examples:
+
+1. Edit a specific cell:
+\`\`\`
+<<<<<<< CELL_OPERATION
+:operation: edit
+:cell_index: 2
+-------
+# Old cell content
+print("Hello")
+=======
+# New cell content
+print("Hello, World!")
+>>>>>>> CELL_OPERATION
+\`\`\`
+
+2. Add a new cell:
+\`\`\`
+<<<<<<< CELL_OPERATION
+:operation: add
+:cell_index: 1
+:cell_type: code
+-------
+=======
+import numpy as np
+import pandas as pd
+>>>>>>> CELL_OPERATION
+\`\`\`
+
+3. Delete a cell:
+\`\`\`
+<<<<<<< CELL_OPERATION
+:operation: delete
+:cell_index: 3
+-------
+=======
+>>>>>>> CELL_OPERATION
+\`\`\`
+
+4. Search and replace within a specific cell:
+\`\`\`
+<<<<<<< SEARCH
+:cell_index: 0
+-------
+old_function()
+=======
+new_function()
+>>>>>>> REPLACE
+\`\`\`
+
+Usage:
+
+notebook.ipynb
+
+Your cell operation or search/replace content here
+
+`
+ }
+
+ async applyDiff(
+ originalContent: string,
+ diffContent: string,
+ _paramStartLine?: number,
+ _paramEndLine?: number,
+ ): Promise {
+ // Check if this is a Jupyter notebook by trying to parse it
+ let handler: JupyterNotebookHandler
+ try {
+ handler = new JupyterNotebookHandler("", originalContent)
+ } catch (error) {
+ // Not a valid notebook, fall back to standard diff
+ return this.fallbackStrategy.applyDiff(originalContent, diffContent, _paramStartLine, _paramEndLine)
+ }
+
+ // Check if this is a cell operation
+ const cellOperationMatch = diffContent.match(
+ /<<<<<<< CELL_OPERATION\s*\n(?::operation:\s*(edit|add|delete)\s*\n)?(?::cell_index:\s*(\d+)\s*\n)?(?::cell_type:\s*(code|markdown|raw)\s*\n)?(?:-------\s*\n)?([\s\S]*?)(?:\n)?=======\s*\n([\s\S]*?)(?:\n)?>>>>>>> CELL_OPERATION/,
+ )
+
+ if (cellOperationMatch) {
+ const operation = cellOperationMatch[1]
+ const cellIndex = parseInt(cellOperationMatch[2] || "0")
+ const cellType = cellOperationMatch[3] as "code" | "markdown" | "raw"
+ const searchContent = cellOperationMatch[4] || ""
+ const replaceContent = cellOperationMatch[5] || ""
+
+ let success = false
+ let error: string | undefined
+
+ switch (operation) {
+ case "edit":
+ if (cellIndex >= 0 && cellIndex < handler.getCellCount()) {
+ success = handler.updateCell(cellIndex, replaceContent)
+ if (!success) {
+ error = `Failed to update cell ${cellIndex}`
+ }
+ } else {
+ error = `Cell index ${cellIndex} is out of range (0-${handler.getCellCount() - 1})`
+ }
+ break
+
+ case "add":
+ if (!cellType) {
+ error = "Cell type is required for add operation"
+ } else {
+ success = handler.insertCell(cellIndex, cellType, replaceContent)
+ if (!success) {
+ error = `Failed to insert cell at index ${cellIndex}`
+ }
+ }
+ break
+
+ case "delete":
+ success = handler.deleteCell(cellIndex)
+ if (!success) {
+ error = `Failed to delete cell ${cellIndex}`
+ }
+ break
+
+ default:
+ error = `Unknown operation: ${operation}`
+ }
+
+ if (success) {
+ return {
+ success: true,
+ content: handler.toJSON(),
+ }
+ } else {
+ return {
+ success: false,
+ error: error || "Cell operation failed",
+ }
+ }
+ }
+
+ // Check if this is a cell-specific search/replace
+ const cellSearchMatch = diffContent.match(
+ /<<<<<<< SEARCH\s*\n(?::cell_index:\s*(\d+)\s*\n)?(?::start_line:\s*(\d+)\s*\n)?(?:-------\s*\n)?([\s\S]*?)(?:\n)?=======\s*\n([\s\S]*?)(?:\n)?>>>>>>> REPLACE/,
+ )
+
+ if (cellSearchMatch) {
+ const cellIndex = cellSearchMatch[1] ? parseInt(cellSearchMatch[1]) : undefined
+ const searchContent = cellSearchMatch[3] || ""
+ const replaceContent = cellSearchMatch[4] || ""
+
+ if (cellIndex !== undefined) {
+ // Apply diff to specific cell
+ const success = handler.applyCellDiff(cellIndex, searchContent, replaceContent)
+ if (success) {
+ return {
+ success: true,
+ content: handler.toJSON(),
+ }
+ } else {
+ return {
+ success: false,
+ error: `Failed to apply diff to cell ${cellIndex}. Content not found or cell doesn't exist.`,
+ }
+ }
+ } else {
+ // Search across all cells
+ let applied = false
+ for (let i = 0; i < handler.getCellCount(); i++) {
+ if (handler.applyCellDiff(i, searchContent, replaceContent)) {
+ applied = true
+ // Continue to apply to all matching cells
+ }
+ }
+
+ if (applied) {
+ return {
+ success: true,
+ content: handler.toJSON(),
+ }
+ } else {
+ return {
+ success: false,
+ error: "Search content not found in any cell",
+ }
+ }
+ }
+ }
+
+ // Fall back to standard diff strategy for the text representation
+ const textRepresentation = handler.extractTextWithCellMarkers()
+ const result = await this.fallbackStrategy.applyDiff(
+ textRepresentation,
+ diffContent,
+ _paramStartLine,
+ _paramEndLine,
+ )
+
+ if (result.success && result.content) {
+ // Convert back from text representation to notebook format
+ // This is a simplified approach - in production, we'd need more sophisticated parsing
+ return {
+ success: true,
+ content: handler.toJSON(),
+ }
+ }
+
+ return result
+ }
+
+ getProgressStatus(toolUse: ToolUse, result?: DiffResult): ToolProgressStatus {
+ const diffContent = toolUse.params.diff
+ if (diffContent) {
+ const icon = "notebook"
+ if (diffContent.includes("CELL_OPERATION")) {
+ const operation = diffContent.match(/:operation:\s*(edit|add|delete)/)?.[1]
+ return { icon, text: operation || "cell" }
+ } else {
+ return this.fallbackStrategy.getProgressStatus(toolUse, result)
+ }
+ }
+ return {}
+ }
+}
diff --git a/src/core/task/Task.ts b/src/core/task/Task.ts
index 1c4d9ec6c7..615079da75 100644
--- a/src/core/task/Task.ts
+++ b/src/core/task/Task.ts
@@ -90,6 +90,7 @@ import { truncateConversationIfNeeded } from "../sliding-window"
import { ClineProvider } from "../webview/ClineProvider"
import { MultiSearchReplaceDiffStrategy } from "../diff/strategies/multi-search-replace"
import { MultiFileSearchReplaceDiffStrategy } from "../diff/strategies/multi-file-search-replace"
+import { JupyterNotebookDiffStrategy } from "../diff/strategies/jupyter-notebook-diff"
import {
type ApiMessage,
readApiMessages,
@@ -382,20 +383,8 @@ export class Task extends EventEmitter implements TaskLike {
// Only set up diff strategy if diff is enabled.
if (this.diffEnabled) {
- // Default to old strategy, will be updated if experiment is enabled.
- this.diffStrategy = new MultiSearchReplaceDiffStrategy(this.fuzzyMatchThreshold)
-
- // Check experiment asynchronously and update strategy if needed.
- provider.getState().then((state) => {
- const isMultiFileApplyDiffEnabled = experiments.isEnabled(
- state.experiments ?? {},
- EXPERIMENT_IDS.MULTI_FILE_APPLY_DIFF,
- )
-
- if (isMultiFileApplyDiffEnabled) {
- this.diffStrategy = new MultiFileSearchReplaceDiffStrategy(this.fuzzyMatchThreshold)
- }
- })
+ // Check for Jupyter notebooks and experiments asynchronously
+ this.initializeDiffStrategy()
}
this.toolRepetitionDetector = new ToolRepetitionDetector(this.consecutiveMistakeLimit)
@@ -2699,6 +2688,45 @@ export class Task extends EventEmitter implements TaskLike {
return checkpointDiff(this, options)
}
+ private async checkForJupyterFiles(workspaceDir: string): Promise {
+ try {
+ const fs = await import("fs/promises")
+ const path = await import("path")
+
+ // Quick check for .ipynb files in the workspace
+ const files = await fs.readdir(workspaceDir)
+ return files.some((file) => path.extname(file).toLowerCase() === ".ipynb")
+ } catch {
+ return false
+ }
+ }
+
+ private async initializeDiffStrategy(): Promise {
+ const workspaceDir = vscode.workspace.workspaceFolders?.[0]?.uri.fsPath
+ const hasJupyterFiles = workspaceDir && (await this.checkForJupyterFiles(workspaceDir))
+
+ if (hasJupyterFiles) {
+ // Use Jupyter-specific diff strategy for notebooks
+ this.diffStrategy = new JupyterNotebookDiffStrategy(this.fuzzyMatchThreshold)
+ } else {
+ // Default to old strategy, will be updated if experiment is enabled.
+ this.diffStrategy = new MultiSearchReplaceDiffStrategy(this.fuzzyMatchThreshold)
+
+ // Check if the multi-file apply diff experiment is enabled
+ const provider = this.providerRef.deref()
+ if (provider) {
+ const state = await provider.getState()
+ const isMultiFileApplyDiffEnabled = experiments.isEnabled(
+ state.experiments ?? {},
+ EXPERIMENT_IDS.MULTI_FILE_APPLY_DIFF,
+ )
+ if (isMultiFileApplyDiffEnabled) {
+ this.diffStrategy = new MultiFileSearchReplaceDiffStrategy(this.fuzzyMatchThreshold)
+ }
+ }
+ }
+ }
+
// Metrics
public combineMessages(messages: ClineMessage[]) {
diff --git a/src/core/webview/generateSystemPrompt.ts b/src/core/webview/generateSystemPrompt.ts
index a639d8a960..5ee089bd80 100644
--- a/src/core/webview/generateSystemPrompt.ts
+++ b/src/core/webview/generateSystemPrompt.ts
@@ -7,6 +7,7 @@ import { experiments as experimentsModule, EXPERIMENT_IDS } from "../../shared/e
import { SYSTEM_PROMPT } from "../prompts/system"
import { MultiSearchReplaceDiffStrategy } from "../diff/strategies/multi-search-replace"
import { MultiFileSearchReplaceDiffStrategy } from "../diff/strategies/multi-file-search-replace"
+import { JupyterNotebookDiffStrategy } from "../diff/strategies/jupyter-notebook-diff"
import { ClineProvider } from "./ClineProvider"
diff --git a/src/integrations/misc/__tests__/jupyter-notebook-handler.spec.ts b/src/integrations/misc/__tests__/jupyter-notebook-handler.spec.ts
new file mode 100644
index 0000000000..286c0b8b88
--- /dev/null
+++ b/src/integrations/misc/__tests__/jupyter-notebook-handler.spec.ts
@@ -0,0 +1,282 @@
+import { describe, it, expect, beforeEach } from "vitest"
+import * as fs from "fs/promises"
+import * as path from "path"
+import { JupyterNotebookHandler, JupyterNotebook } from "../jupyter-notebook-handler"
+
+describe("JupyterNotebookHandler", () => {
+ let handler: JupyterNotebookHandler
+ let sampleNotebook: JupyterNotebook
+
+ beforeEach(() => {
+ sampleNotebook = {
+ cells: [
+ {
+ cell_type: "markdown",
+ source: ["# Test Notebook\n", "This is a test notebook"],
+ metadata: {},
+ },
+ {
+ cell_type: "code",
+ source: ["import numpy as np\n", "import pandas as pd"],
+ metadata: {},
+ outputs: [],
+ execution_count: 1,
+ },
+ {
+ cell_type: "code",
+ source: ["def hello():\n", " print('Hello, World!')"],
+ metadata: {},
+ outputs: [],
+ execution_count: 2,
+ },
+ ],
+ metadata: {
+ kernelspec: {
+ display_name: "Python 3",
+ language: "python",
+ name: "python3",
+ },
+ },
+ nbformat: 4,
+ nbformat_minor: 4,
+ }
+
+ handler = new JupyterNotebookHandler("test.ipynb", JSON.stringify(sampleNotebook))
+ })
+
+ describe("Cell Operations", () => {
+ it("should get cell by index", () => {
+ const cell = handler.getCellByIndex(0)
+ expect(cell).toBeDefined()
+ expect(cell?.cell_type).toBe("markdown")
+ })
+
+ it("should return undefined for invalid cell index", () => {
+ const cell = handler.getCellByIndex(10)
+ expect(cell).toBeUndefined()
+ })
+
+ it("should get cell count", () => {
+ expect(handler.getCellCount()).toBe(3)
+ })
+
+ it("should get cells by type", () => {
+ const codeCells = handler.getCellsByType("code")
+ expect(codeCells).toHaveLength(2)
+ expect(codeCells[0].index).toBe(1)
+ expect(codeCells[1].index).toBe(2)
+
+ const markdownCells = handler.getCellsByType("markdown")
+ expect(markdownCells).toHaveLength(1)
+ expect(markdownCells[0].index).toBe(0)
+ })
+ })
+
+ describe("Cell Modification", () => {
+ it("should update cell content", () => {
+ const newContent = "# Updated Title\nNew content here"
+ const success = handler.updateCell(0, newContent)
+
+ expect(success).toBe(true)
+ const updatedCell = handler.getCellByIndex(0)
+ // Jupyter format adds newline to each line except possibly the last
+ expect(updatedCell?.source).toEqual(["# Updated Title\n", "New content here\n"])
+ })
+
+ it("should insert a new cell", () => {
+ const success = handler.insertCell(1, "code", "print('New cell')")
+
+ expect(success).toBe(true)
+ expect(handler.getCellCount()).toBe(4)
+
+ const newCell = handler.getCellByIndex(1)
+ expect(newCell?.cell_type).toBe("code")
+ // Single line cells get a newline appended
+ expect(newCell?.source).toEqual(["print('New cell')\n"])
+ })
+
+ it("should delete a cell", () => {
+ const success = handler.deleteCell(1)
+
+ expect(success).toBe(true)
+ expect(handler.getCellCount()).toBe(2)
+
+ // Check that the second code cell is now at index 1
+ const cell = handler.getCellByIndex(1)
+ expect(cell?.source).toEqual(["def hello():\n", " print('Hello, World!')"])
+ })
+
+ it("should return false for invalid cell operations", () => {
+ expect(handler.updateCell(-1, "content")).toBe(false)
+ expect(handler.updateCell(10, "content")).toBe(false)
+ expect(handler.insertCell(-1, "code", "content")).toBe(false)
+ expect(handler.deleteCell(10)).toBe(false)
+ })
+ })
+
+ describe("Text Extraction", () => {
+ it("should extract text with cell markers", () => {
+ const text = handler.extractTextWithCellMarkers()
+
+ expect(text).toContain("# %%% Cell 1 [markdown]")
+ expect(text).toContain("# %%% Cell 2 [code]")
+ expect(text).toContain("# %%% Cell 3 [code]")
+ expect(text).toContain("# Test Notebook")
+ expect(text).toContain("import numpy as np")
+ expect(text).toContain("def hello():")
+ })
+
+ it("should extract specific cells text", () => {
+ const text = handler.extractCellsText([0, 2])
+
+ expect(text).toContain("# Cell 1 [markdown]")
+ expect(text).toContain("# Test Notebook")
+ expect(text).toContain("# Cell 3 [code]")
+ expect(text).toContain("def hello():")
+ expect(text).not.toContain("import numpy")
+ })
+
+ it("should extract all cells text when no indices provided", () => {
+ const text = handler.extractCellsText()
+
+ expect(text).toContain("# Cell 1 [markdown]")
+ expect(text).toContain("# Cell 2 [code]")
+ expect(text).toContain("# Cell 3 [code]")
+ })
+ })
+
+ describe("Cell Search", () => {
+ it("should search for content in cells", () => {
+ const results = handler.searchInCells("import")
+
+ expect(results).toHaveLength(1)
+ expect(results[0].cellIndex).toBe(1)
+ expect(results[0].matches).toHaveLength(2)
+ expect(results[0].matches[0]).toBe("import numpy as np")
+ })
+
+ it("should return empty array when no matches found", () => {
+ const results = handler.searchInCells("nonexistent")
+ expect(results).toHaveLength(0)
+ })
+ })
+
+ describe("Cell Diff", () => {
+ it("should apply diff to a specific cell", () => {
+ const success = handler.applyCellDiff(2, "def hello():", "def greet(name):")
+
+ expect(success).toBe(true)
+ const cell = handler.getCellByIndex(2)
+ // Check the actual content after replacement
+ const sourceStr = Array.isArray(cell?.source) ? cell.source.join("") : cell?.source
+ expect(sourceStr).toContain("def greet(name):")
+ })
+
+ it("should return false when search content not found", () => {
+ const success = handler.applyCellDiff(2, "nonexistent", "replacement")
+
+ expect(success).toBe(false)
+ })
+ })
+
+ describe("Line Number Mapping", () => {
+ it("should get cell at specific line number", () => {
+ const cellRef = handler.getCellAtLine(1)
+ expect(cellRef).toBeDefined()
+ expect(cellRef?.index).toBe(0)
+ expect(cellRef?.type).toBe("markdown")
+
+ const cellRef2 = handler.getCellAtLine(4)
+ expect(cellRef2).toBeDefined()
+ expect(cellRef2?.index).toBe(1)
+ expect(cellRef2?.type).toBe("code")
+ })
+
+ it("should return undefined for invalid line number", () => {
+ const cellRef = handler.getCellAtLine(100)
+ expect(cellRef).toBeUndefined()
+ })
+ })
+
+ describe("JSON Serialization", () => {
+ it("should serialize to JSON", () => {
+ const json = handler.toJSON()
+ const parsed = JSON.parse(json)
+
+ expect(parsed.cells).toHaveLength(3)
+ expect(parsed.metadata).toBeDefined()
+ expect(parsed.nbformat).toBe(4)
+ })
+ })
+
+ describe("Checkpoint Support", () => {
+ it("should create checkpoint representation", () => {
+ const checkpoint = handler.getCheckpointRepresentation()
+
+ expect(checkpoint).toContain("# %%% Cell")
+ expect(checkpoint).toContain("# Test Notebook")
+ expect(checkpoint).toContain("import numpy as np")
+ })
+
+ it("should restore from checkpoint representation", () => {
+ const checkpoint = handler.getCheckpointRepresentation()
+ const restored = JupyterNotebookHandler.fromCheckpointRepresentation(checkpoint, sampleNotebook)
+
+ expect(restored.cells).toHaveLength(3)
+ expect(restored.cells[0].cell_type).toBe("markdown")
+ expect(restored.cells[1].cell_type).toBe("code")
+ expect(restored.cells[2].cell_type).toBe("code")
+ })
+ })
+
+ describe("Edge Cases", () => {
+ it("should handle empty notebook", () => {
+ const emptyHandler = new JupyterNotebookHandler(
+ "empty.ipynb",
+ JSON.stringify({
+ cells: [],
+ metadata: {},
+ }),
+ )
+
+ expect(emptyHandler.getCellCount()).toBe(0)
+ expect(emptyHandler.extractTextWithCellMarkers()).toBe("")
+ expect(emptyHandler.searchInCells("test")).toHaveLength(0)
+ })
+
+ it("should handle cells with string source instead of array", () => {
+ const notebook = {
+ cells: [
+ {
+ cell_type: "code" as const,
+ source: "print('single line')",
+ metadata: {},
+ },
+ ],
+ }
+
+ const handler = new JupyterNotebookHandler("test.ipynb", JSON.stringify(notebook))
+ expect(handler.getCellByIndex(0)?.source).toBe("print('single line')")
+
+ // Update should preserve the format
+ handler.updateCell(0, "print('updated')")
+ expect(handler.getCellByIndex(0)?.source).toBe("print('updated')")
+ })
+
+ it("should handle cells with empty source", () => {
+ const notebook = {
+ cells: [
+ {
+ cell_type: "code" as const,
+ source: [],
+ metadata: {},
+ },
+ ],
+ }
+
+ const handler = new JupyterNotebookHandler("test.ipynb", JSON.stringify(notebook))
+ const text = handler.extractTextWithCellMarkers()
+ expect(text).toContain("# %%% Cell 1 [code]")
+ })
+ })
+})
diff --git a/src/integrations/misc/extract-text.ts b/src/integrations/misc/extract-text.ts
index 8231c609be..58d166153e 100644
--- a/src/integrations/misc/extract-text.ts
+++ b/src/integrations/misc/extract-text.ts
@@ -7,6 +7,7 @@ import { isBinaryFile } from "isbinaryfile"
import { extractTextFromXLSX } from "./extract-text-from-xlsx"
import { countFileLines } from "./line-counter"
import { readLines } from "./read-lines"
+import { JupyterNotebookHandler } from "./jupyter-notebook-handler"
async function extractTextFromPDF(filePath: string): Promise {
const dataBuffer = await fs.readFile(filePath)
@@ -20,17 +21,8 @@ async function extractTextFromDOCX(filePath: string): Promise {
}
async function extractTextFromIPYNB(filePath: string): Promise {
- const data = await fs.readFile(filePath, "utf8")
- const notebook = JSON.parse(data)
- let extractedText = ""
-
- for (const cell of notebook.cells) {
- if ((cell.cell_type === "markdown" || cell.cell_type === "code") && cell.source) {
- extractedText += cell.source.join("\n") + "\n"
- }
- }
-
- return addLineNumbers(extractedText)
+ const handler = await JupyterNotebookHandler.fromFile(filePath)
+ return handler.extractTextWithCellMarkers()
}
/**
diff --git a/src/integrations/misc/jupyter-notebook-handler.ts b/src/integrations/misc/jupyter-notebook-handler.ts
new file mode 100644
index 0000000000..529dff5a38
--- /dev/null
+++ b/src/integrations/misc/jupyter-notebook-handler.ts
@@ -0,0 +1,338 @@
+import * as fs from "fs/promises"
+import * as path from "path"
+import { addLineNumbers } from "./extract-text"
+
+export interface JupyterCell {
+ cell_type: "code" | "markdown" | "raw"
+ source: string | string[]
+ metadata?: Record
+ outputs?: any[]
+ execution_count?: number | null
+}
+
+export interface JupyterNotebook {
+ cells: JupyterCell[]
+ metadata?: Record
+ nbformat?: number
+ nbformat_minor?: number
+}
+
+export interface CellReference {
+ index: number
+ type: "code" | "markdown" | "raw"
+ content: string
+ lineStart: number
+ lineEnd: number
+}
+
+export class JupyterNotebookHandler {
+ private notebook: JupyterNotebook
+ private filePath: string
+ private cellReferences: CellReference[] = []
+
+ constructor(filePath: string, notebookContent?: string) {
+ this.filePath = filePath
+ if (notebookContent) {
+ this.notebook = JSON.parse(notebookContent)
+ this.buildCellReferences()
+ } else {
+ this.notebook = { cells: [] }
+ }
+ }
+
+ /**
+ * Load a Jupyter notebook from file
+ */
+ static async fromFile(filePath: string): Promise {
+ const content = await fs.readFile(filePath, "utf8")
+ return new JupyterNotebookHandler(filePath, content)
+ }
+
+ /**
+ * Build cell references with line number mappings
+ */
+ private buildCellReferences(): void {
+ this.cellReferences = []
+ let currentLine = 1
+
+ this.notebook.cells.forEach((cell, index) => {
+ const source = Array.isArray(cell.source) ? cell.source.join("") : cell.source || ""
+ const lines = source.split("\n")
+ const lineCount = lines.length
+
+ this.cellReferences.push({
+ index,
+ type: cell.cell_type,
+ content: source,
+ lineStart: currentLine,
+ lineEnd: currentLine + lineCount - 1,
+ })
+
+ currentLine += lineCount + 1 // Add 1 for cell separator
+ })
+ }
+
+ /**
+ * Get cell at a specific line number
+ */
+ getCellAtLine(lineNumber: number): CellReference | undefined {
+ return this.cellReferences.find((ref) => lineNumber >= ref.lineStart && lineNumber <= ref.lineEnd)
+ }
+
+ /**
+ * Get cell by index
+ */
+ getCellByIndex(index: number): JupyterCell | undefined {
+ return this.notebook.cells[index]
+ }
+
+ /**
+ * Extract text with cell markers for better readability
+ */
+ extractTextWithCellMarkers(): string {
+ let result = ""
+ let lineNumber = 1
+
+ this.notebook.cells.forEach((cell, index) => {
+ const cellType = cell.cell_type
+ const source = Array.isArray(cell.source) ? cell.source.join("") : cell.source || ""
+
+ // Add cell header
+ result += `# %%% Cell ${index + 1} [${cellType}]\n`
+
+ // Add cell content with line numbers
+ const lines = source.split("\n")
+ lines.forEach((line) => {
+ result += `${String(lineNumber).padStart(4, " ")} | ${line}\n`
+ lineNumber++
+ })
+
+ // Add cell separator
+ result += "\n"
+ lineNumber++ // Account for the separator line
+ })
+
+ return result
+ }
+
+ /**
+ * Extract text from specific cells
+ */
+ extractCellsText(cellIndices?: number[]): string {
+ let result = ""
+ const cellsToExtract = cellIndices
+ ? this.notebook.cells.filter((_, index) => cellIndices.includes(index))
+ : this.notebook.cells
+
+ cellsToExtract.forEach((cell, idx) => {
+ const actualIndex = cellIndices ? cellIndices[idx] : idx
+ const source = Array.isArray(cell.source) ? cell.source.join("") : cell.source || ""
+
+ result += `# Cell ${actualIndex + 1} [${cell.cell_type}]\n`
+ result += source
+ result += "\n\n"
+ })
+
+ return result
+ }
+
+ /**
+ * Update a specific cell's content
+ */
+ updateCell(cellIndex: number, newContent: string): boolean {
+ if (cellIndex < 0 || cellIndex >= this.notebook.cells.length) {
+ return false
+ }
+
+ const cell = this.notebook.cells[cellIndex]
+ // Preserve the original format (array vs string)
+ if (Array.isArray(cell.source)) {
+ // Split content and ensure each line ends with \n except the last
+ const lines = newContent.split("\n")
+ cell.source = lines.map((line, idx) => (idx === lines.length - 1 && line === "" ? line : line + "\n"))
+ // Remove trailing empty string if it exists
+ if (cell.source[cell.source.length - 1] === "") {
+ cell.source.pop()
+ }
+ } else {
+ cell.source = newContent
+ }
+
+ // Rebuild references after update
+ this.buildCellReferences()
+ return true
+ }
+
+ /**
+ * Insert a new cell
+ */
+ insertCell(index: number, cellType: "code" | "markdown" | "raw", content: string): boolean {
+ if (index < 0 || index > this.notebook.cells.length) {
+ return false
+ }
+
+ const newCell: JupyterCell = {
+ cell_type: cellType,
+ source: content
+ .split("\n")
+ .map((line, idx, arr) => (idx === arr.length - 1 && line === "" ? line : line + "\n")),
+ metadata: {},
+ }
+
+ if (cellType === "code") {
+ newCell.outputs = []
+ newCell.execution_count = null
+ }
+
+ this.notebook.cells.splice(index, 0, newCell)
+ this.buildCellReferences()
+ return true
+ }
+
+ /**
+ * Delete a cell
+ */
+ deleteCell(index: number): boolean {
+ if (index < 0 || index >= this.notebook.cells.length) {
+ return false
+ }
+
+ this.notebook.cells.splice(index, 1)
+ this.buildCellReferences()
+ return true
+ }
+
+ /**
+ * Apply a diff to a specific cell
+ */
+ applyCellDiff(cellIndex: number, searchContent: string, replaceContent: string): boolean {
+ const cell = this.getCellByIndex(cellIndex)
+ if (!cell) {
+ return false
+ }
+
+ const currentContent = Array.isArray(cell.source) ? cell.source.join("") : cell.source || ""
+
+ // Simple exact match replacement for now
+ if (currentContent.includes(searchContent)) {
+ const newContent = currentContent.replace(searchContent, replaceContent)
+ return this.updateCell(cellIndex, newContent)
+ }
+
+ return false
+ }
+
+ /**
+ * Save the notebook back to file
+ */
+ async save(): Promise {
+ const content = JSON.stringify(this.notebook, null, 2)
+ await fs.writeFile(this.filePath, content, "utf8")
+ }
+
+ /**
+ * Get the notebook as JSON string
+ */
+ toJSON(): string {
+ return JSON.stringify(this.notebook, null, 2)
+ }
+
+ /**
+ * Get cell count
+ */
+ getCellCount(): number {
+ return this.notebook.cells.length
+ }
+
+ /**
+ * Get all cells of a specific type
+ */
+ getCellsByType(cellType: "code" | "markdown" | "raw"): Array<{ index: number; cell: JupyterCell }> {
+ return this.notebook.cells
+ .map((cell, index) => ({ index, cell }))
+ .filter(({ cell }) => cell.cell_type === cellType)
+ }
+
+ /**
+ * Search for content in cells
+ */
+ searchInCells(searchTerm: string): Array<{ cellIndex: number; matches: string[] }> {
+ const results: Array<{ cellIndex: number; matches: string[] }> = []
+
+ this.notebook.cells.forEach((cell, index) => {
+ const source = Array.isArray(cell.source) ? cell.source.join("") : cell.source || ""
+ const lines = source.split("\n")
+ const matches = lines.filter((line) => line.includes(searchTerm))
+
+ if (matches.length > 0) {
+ results.push({ cellIndex: index, matches })
+ }
+ })
+
+ return results
+ }
+
+ /**
+ * Create a checkpoint-friendly representation
+ */
+ getCheckpointRepresentation(): string {
+ return this.extractTextWithCellMarkers()
+ }
+
+ /**
+ * Restore from checkpoint representation
+ */
+ static fromCheckpointRepresentation(checkpointContent: string, originalNotebook: JupyterNotebook): JupyterNotebook {
+ // Parse the checkpoint content and reconstruct cells
+ const lines = checkpointContent.split("\n")
+ const cells: JupyterCell[] = []
+ let currentCell: JupyterCell | null = null
+ let currentContent: string[] = []
+
+ for (const line of lines) {
+ // Check for cell marker
+ const cellMarkerMatch = line.match(/^# %%% Cell (\d+) \[(code|markdown|raw)\]$/)
+ if (cellMarkerMatch) {
+ // Save previous cell if exists
+ if (currentCell) {
+ currentCell.source = currentContent.join("\n")
+ cells.push(currentCell)
+ }
+
+ // Start new cell
+ const cellIndex = parseInt(cellMarkerMatch[1]) - 1
+ const cellType = cellMarkerMatch[2] as "code" | "markdown" | "raw"
+
+ // Try to preserve metadata from original
+ const originalCell = originalNotebook.cells[cellIndex]
+ currentCell = {
+ cell_type: cellType,
+ source: "",
+ metadata: originalCell?.metadata || {},
+ }
+
+ if (cellType === "code") {
+ currentCell.outputs = originalCell?.outputs || []
+ currentCell.execution_count = originalCell?.execution_count || null
+ }
+
+ currentContent = []
+ } else if (line.match(/^\s*\d+\s*\|/)) {
+ // Extract content from numbered line
+ const content = line.replace(/^\s*\d+\s*\|\s?/, "")
+ currentContent.push(content)
+ }
+ }
+
+ // Save last cell
+ if (currentCell) {
+ currentCell.source = currentContent.join("\n")
+ cells.push(currentCell)
+ }
+
+ return {
+ ...originalNotebook,
+ cells,
+ }
+ }
+}