1use async_trait::async_trait;
2use bamboo_agent_core::{Tool, ToolCtx, ToolError, ToolOutcome, ToolResult};
3use serde::Deserialize;
4use serde_json::{json, Value};
5use std::path::Path;
6
7use super::file_change;
8
9#[derive(Debug, Deserialize)]
10#[serde(rename_all = "lowercase")]
11enum CellType {
12 Code,
13 Markdown,
14}
15
16impl CellType {
17 fn as_str(&self) -> &'static str {
18 match self {
19 CellType::Code => "code",
20 CellType::Markdown => "markdown",
21 }
22 }
23}
24
25#[derive(Debug, Deserialize, Default)]
26#[serde(rename_all = "lowercase")]
27enum EditMode {
28 #[default]
29 Replace,
30 Insert,
31 Delete,
32}
33
34#[derive(Debug, Deserialize)]
35struct NotebookEditArgs {
36 notebook_path: String,
37 #[serde(default)]
38 cell_id: Option<String>,
39 new_source: String,
40 #[serde(default)]
41 cell_type: Option<CellType>,
42 #[serde(default)]
43 edit_mode: Option<EditMode>,
44}
45
46pub struct NotebookEditTool;
47
48impl NotebookEditTool {
49 pub fn new() -> Self {
50 Self
51 }
52
53 fn source_to_lines(source: &str) -> Vec<Value> {
54 if source.is_empty() {
55 return vec![Value::String(String::new())];
56 }
57
58 source
59 .lines()
60 .map(|line| Value::String(format!("{}\n", line)))
61 .collect()
62 }
63
64 fn find_cell_index(cells: &[Value], cell_id: Option<&str>) -> Option<usize> {
65 let cell_id = cell_id?;
66
67 cells.iter().position(|cell| {
68 cell.get("id")
69 .and_then(|value| value.as_str())
70 .map(|value| value == cell_id)
71 .unwrap_or(false)
72 })
73 }
74}
75
76impl Default for NotebookEditTool {
77 fn default() -> Self {
78 Self::new()
79 }
80}
81
82#[async_trait]
83impl Tool for NotebookEditTool {
84 fn name(&self) -> &str {
85 "NotebookEdit"
86 }
87
88 fn description(&self) -> &str {
89 "Replace, insert, or delete a Jupyter notebook cell"
90 }
91
92 fn parameters_schema(&self) -> serde_json::Value {
93 json!({
94 "type": "object",
95 "properties": {
96 "notebook_path": {
97 "type": "string",
98 "description": "Absolute path to the notebook file"
99 },
100 "cell_id": {
101 "type": "string",
102 "description": "Cell ID to edit"
103 },
104 "new_source": {
105 "type": "string",
106 "description": "New source content for the cell"
107 },
108 "cell_type": {
109 "type": "string",
110 "enum": ["code", "markdown"],
111 "description": "Cell type when inserting"
112 },
113 "edit_mode": {
114 "type": "string",
115 "enum": ["replace", "insert", "delete"],
116 "description": "Edit mode"
117 }
118 },
119 "required": ["notebook_path", "new_source"],
120 "additionalProperties": false
121 })
122 }
123
124 async fn invoke(
125 &self,
126 args: serde_json::Value,
127 _ctx: ToolCtx,
128 ) -> Result<ToolOutcome, ToolError> {
129 let parsed: NotebookEditArgs = serde_json::from_value(args).map_err(|e| {
130 ToolError::InvalidArguments(format!("Invalid NotebookEdit args: {}", e))
131 })?;
132
133 let path = Path::new(parsed.notebook_path.trim());
134 if !path.is_absolute() {
135 return Err(ToolError::InvalidArguments(
136 "notebook_path must be absolute".to_string(),
137 ));
138 }
139
140 let content = tokio::fs::read_to_string(path)
141 .await
142 .map_err(|e| ToolError::Execution(format!("Failed to read notebook: {}", e)))?;
143 let checkpoint = file_change::create_checkpoint(path, Some(content.as_bytes())).await?;
144
145 let mut notebook: Value = serde_json::from_str(&content)
146 .map_err(|e| ToolError::Execution(format!("Invalid notebook JSON: {}", e)))?;
147
148 let cells = notebook
149 .get_mut("cells")
150 .and_then(|value| value.as_array_mut())
151 .ok_or_else(|| ToolError::Execution("Notebook missing 'cells' array".to_string()))?;
152
153 let edit_mode = parsed.edit_mode.unwrap_or_default();
154 let cell_id = parsed
155 .cell_id
156 .as_deref()
157 .map(str::trim)
158 .filter(|value| !value.is_empty());
159 let target_index = Self::find_cell_index(cells, cell_id);
160
161 match edit_mode {
162 EditMode::Replace => {
163 if cell_id.is_none() {
164 return Err(ToolError::InvalidArguments(
165 "cell_id is required when edit_mode=replace".to_string(),
166 ));
167 }
168 let idx = target_index.ok_or_else(|| {
169 ToolError::Execution("Target cell not found for replace".to_string())
170 })?;
171 if let Some(cell) = cells.get_mut(idx) {
172 cell["source"] = Value::Array(Self::source_to_lines(&parsed.new_source));
173 }
174 }
175 EditMode::Insert => {
176 let cell_type = parsed.cell_type.ok_or_else(|| {
177 ToolError::InvalidArguments(
178 "cell_type is required when edit_mode=insert".to_string(),
179 )
180 })?;
181 let new_cell = json!({
182 "id": uuid::Uuid::new_v4().to_string(),
183 "cell_type": cell_type.as_str(),
184 "metadata": {},
185 "source": Self::source_to_lines(&parsed.new_source),
186 "outputs": [],
187 "execution_count": serde_json::Value::Null,
188 });
189
190 if let Some(cell_id) = cell_id {
191 let idx = target_index.ok_or_else(|| {
192 ToolError::Execution(format!(
193 "Target cell '{}' not found for insert",
194 cell_id
195 ))
196 })?;
197 cells.insert(idx + 1, new_cell);
198 } else {
199 cells.push(new_cell);
200 }
201 }
202 EditMode::Delete => {
203 if cell_id.is_none() {
204 return Err(ToolError::InvalidArguments(
205 "cell_id is required when edit_mode=delete".to_string(),
206 ));
207 }
208 let idx = target_index.ok_or_else(|| {
209 ToolError::Execution("Target cell not found for delete".to_string())
210 })?;
211 cells.remove(idx);
212 }
213 }
214
215 let updated = serde_json::to_string_pretty(¬ebook)
216 .map_err(|e| ToolError::Execution(format!("Failed to serialize notebook: {}", e)))?;
217
218 file_change::atomic_write_text(path, &updated).await?;
219
220 let payload = file_change::build_file_change_payload(
221 "NotebookEdit",
222 path,
223 format!("Notebook updated: {}", parsed.notebook_path),
224 checkpoint,
225 &content,
226 &updated,
227 );
228
229 Ok(ToolOutcome::Completed(ToolResult {
230 success: true,
231 result: payload,
232 display_preference: Some("Default".to_string()),
233 images: Vec::new(),
234 }))
235 }
236}
237
238#[cfg(test)]
239mod tests {
240 use super::*;
241
242 fn sample_notebook() -> serde_json::Value {
243 json!({
244 "cells": [
245 {
246 "id": "cell-a",
247 "cell_type": "code",
248 "metadata": {},
249 "source": ["print('a')\n"],
250 "outputs": [],
251 "execution_count": 1
252 },
253 {
254 "id": "cell-b",
255 "cell_type": "markdown",
256 "metadata": {},
257 "source": ["# title\n"]
258 }
259 ],
260 "metadata": {},
261 "nbformat": 4,
262 "nbformat_minor": 5
263 })
264 }
265
266 async fn write_notebook(path: &Path) {
267 tokio::fs::write(
268 path,
269 serde_json::to_string_pretty(&sample_notebook()).unwrap(),
270 )
271 .await
272 .unwrap();
273 }
274
275 #[tokio::test]
276 async fn replace_requires_cell_id() {
277 let file = tempfile::NamedTempFile::new().unwrap();
278 write_notebook(file.path()).await;
279
280 let tool = NotebookEditTool::new();
281 let result = tool
282 .invoke(
283 json!({
284 "notebook_path": file.path(),
285 "edit_mode": "replace",
286 "new_source": "updated"
287 }),
288 ToolCtx::none("t"),
289 )
290 .await;
291
292 assert!(matches!(result, Err(ToolError::InvalidArguments(msg)) if msg.contains("cell_id")));
293 }
294
295 #[tokio::test]
296 async fn delete_requires_cell_id() {
297 let file = tempfile::NamedTempFile::new().unwrap();
298 write_notebook(file.path()).await;
299
300 let tool = NotebookEditTool::new();
301 let result = tool
302 .invoke(
303 json!({
304 "notebook_path": file.path(),
305 "edit_mode": "delete",
306 "new_source": ""
307 }),
308 ToolCtx::none("t"),
309 )
310 .await;
311
312 assert!(matches!(result, Err(ToolError::InvalidArguments(msg)) if msg.contains("cell_id")));
313 }
314
315 #[tokio::test]
316 async fn insert_without_cell_id_appends_cell() {
317 let file = tempfile::NamedTempFile::new().unwrap();
318 write_notebook(file.path()).await;
319
320 let tool = NotebookEditTool::new();
321 let out = tool
322 .invoke(
323 json!({
324 "notebook_path": file.path(),
325 "edit_mode": "insert",
326 "cell_type": "markdown",
327 "new_source": "appended cell"
328 }),
329 ToolCtx::none("t"),
330 )
331 .await
332 .unwrap();
333 let ToolOutcome::Completed(result) = out else {
334 panic!("expected Completed")
335 };
336 assert!(result.success);
337
338 let updated: Value =
339 serde_json::from_str(&tokio::fs::read_to_string(file.path()).await.unwrap()).unwrap();
340 let cells = updated["cells"].as_array().unwrap();
341 assert_eq!(cells.len(), 3);
342 let last = cells.last().unwrap();
343 assert_eq!(last["cell_type"], "markdown");
344 let source = last["source"].as_array().unwrap();
345 assert_eq!(source[0], "appended cell\n");
346 }
347
348 #[tokio::test]
349 async fn insert_with_unknown_cell_id_returns_error() {
350 let file = tempfile::NamedTempFile::new().unwrap();
351 write_notebook(file.path()).await;
352
353 let tool = NotebookEditTool::new();
354 let result = tool
355 .invoke(
356 json!({
357 "notebook_path": file.path(),
358 "edit_mode": "insert",
359 "cell_id": "does-not-exist",
360 "cell_type": "code",
361 "new_source": "print('x')"
362 }),
363 ToolCtx::none("t"),
364 )
365 .await;
366
367 assert!(matches!(result, Err(ToolError::Execution(msg)) if msg.contains("not found")));
368 }
369}