claude_rust_tools/infrastructure/
notebook_edit_tool.rs1use claude_rust_errors::{AppError, AppResult};
2use claude_rust_types::{PermissionLevel, Tool};
3use serde_json::{Value, json};
4
5pub struct NotebookEditTool;
6
7impl NotebookEditTool {
8 pub fn new() -> Self {
9 Self
10 }
11}
12
13#[async_trait::async_trait]
14impl Tool for NotebookEditTool {
15 fn name(&self) -> &str {
16 "notebook_edit"
17 }
18
19 fn description(&self) -> &str {
20 "Edit a Jupyter notebook (.ipynb) file by inserting, replacing, or deleting cells."
21 }
22
23 fn input_schema(&self) -> Value {
24 json!({
25 "type": "object",
26 "properties": {
27 "notebook_path": {
28 "type": "string",
29 "description": "Path to the .ipynb notebook file"
30 },
31 "operation": {
32 "type": "string",
33 "description": "Operation to perform: insert_cell, replace_cell, or delete_cell",
34 "enum": ["insert_cell", "replace_cell", "delete_cell"]
35 },
36 "cell_index": {
37 "type": "integer",
38 "description": "Index of the cell to operate on (0-based)"
39 },
40 "cell_type": {
41 "type": "string",
42 "description": "Type of cell (for insert_cell): code or markdown",
43 "enum": ["code", "markdown"]
44 },
45 "source": {
46 "type": "string",
47 "description": "Source content for the cell (for insert_cell and replace_cell)"
48 }
49 },
50 "required": ["notebook_path", "operation", "cell_index"]
51 })
52 }
53
54 fn permission_level(&self) -> PermissionLevel {
55 PermissionLevel::Dangerous
56 }
57
58 fn is_destructive(&self, _input: &Value) -> bool { true }
59
60 fn get_path(&self, input: &Value) -> Option<String> {
61 input.get("notebook_path").and_then(|v| v.as_str()).map(|s| s.to_string())
62 }
63
64 async fn execute(&self, input: Value) -> AppResult<String> {
65 let notebook_path = input
66 .get("notebook_path")
67 .and_then(|v| v.as_str())
68 .ok_or_else(|| AppError::Tool("missing 'notebook_path' field".into()))?;
69
70 let operation = input
71 .get("operation")
72 .and_then(|v| v.as_str())
73 .ok_or_else(|| AppError::Tool("missing 'operation' field".into()))?;
74
75 let cell_index = input
76 .get("cell_index")
77 .and_then(|v| v.as_u64())
78 .ok_or_else(|| AppError::Tool("missing 'cell_index' field".into()))? as usize;
79
80 tracing::info!(notebook_path, operation, cell_index, "editing notebook");
81
82 let content = tokio::fs::read_to_string(notebook_path)
83 .await
84 .map_err(|e| AppError::Tool(format!("cannot read '{notebook_path}': {e}")))?;
85
86 let mut notebook: Value = serde_json::from_str(&content)
87 .map_err(|e| AppError::Tool(format!("invalid notebook JSON: {e}")))?;
88
89 let cells = notebook
90 .get_mut("cells")
91 .and_then(|v| v.as_array_mut())
92 .ok_or_else(|| AppError::Tool("notebook has no 'cells' array".into()))?;
93
94 let result = match operation {
95 "insert_cell" => {
96 let cell_type = input
97 .get("cell_type")
98 .and_then(|v| v.as_str())
99 .unwrap_or("code");
100
101 let source = input
102 .get("source")
103 .and_then(|v| v.as_str())
104 .unwrap_or("");
105
106 let source_lines: Vec<Value> = source
107 .lines()
108 .enumerate()
109 .map(|(i, line)| {
110 let total = source.lines().count();
111 if i < total - 1 {
112 Value::String(format!("{line}\n"))
113 } else {
114 Value::String(line.to_string())
115 }
116 })
117 .collect();
118
119 let source_lines = if source_lines.is_empty() {
120 vec![Value::String(String::new())]
121 } else {
122 source_lines
123 };
124
125 let new_cell = if cell_type == "code" {
126 json!({
127 "cell_type": "code",
128 "execution_count": null,
129 "metadata": {},
130 "outputs": [],
131 "source": source_lines
132 })
133 } else {
134 json!({
135 "cell_type": "markdown",
136 "metadata": {},
137 "source": source_lines
138 })
139 };
140
141 let idx = cell_index.min(cells.len());
142 cells.insert(idx, new_cell);
143 format!("Inserted {cell_type} cell at index {idx}.")
144 }
145 "replace_cell" => {
146 if cell_index >= cells.len() {
147 return Err(AppError::Tool(format!(
148 "cell_index {cell_index} out of range (notebook has {} cells)",
149 cells.len()
150 )));
151 }
152
153 let source = input
154 .get("source")
155 .and_then(|v| v.as_str())
156 .unwrap_or("");
157
158 let source_lines: Vec<Value> = source
159 .lines()
160 .enumerate()
161 .map(|(i, line)| {
162 let total = source.lines().count();
163 if i < total - 1 {
164 Value::String(format!("{line}\n"))
165 } else {
166 Value::String(line.to_string())
167 }
168 })
169 .collect();
170
171 let source_lines = if source_lines.is_empty() {
172 vec![Value::String(String::new())]
173 } else {
174 source_lines
175 };
176
177 cells[cell_index]["source"] = Value::Array(source_lines);
178 format!("Replaced source of cell at index {cell_index}.")
179 }
180 "delete_cell" => {
181 if cell_index >= cells.len() {
182 return Err(AppError::Tool(format!(
183 "cell_index {cell_index} out of range (notebook has {} cells)",
184 cells.len()
185 )));
186 }
187
188 cells.remove(cell_index);
189 format!("Deleted cell at index {cell_index}.")
190 }
191 _ => {
192 return Err(AppError::Tool(format!(
193 "unknown operation '{operation}': expected insert_cell, replace_cell, or delete_cell"
194 )));
195 }
196 };
197
198 let output = serde_json::to_string_pretty(¬ebook)
199 .map_err(|e| AppError::Tool(format!("failed to serialize notebook: {e}")))?;
200
201 tokio::fs::write(notebook_path, output)
202 .await
203 .map_err(|e| AppError::Tool(format!("cannot write '{notebook_path}': {e}")))?;
204
205 Ok(result)
206 }
207}