| 1 | package builtin |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "encoding/json" |
| 6 | "fmt" |
| 7 | "strings" |
| 8 | |
| 9 | "reasonix/internal/diff" |
| 10 | "reasonix/internal/sandbox" |
| 11 | "reasonix/internal/tool" |
| 12 | ) |
| 13 | |
| 14 | func init() { tool.RegisterBuiltin(notebookEdit{}) } |
| 15 | |
| 16 | // notebookEdit edits a single cell of a Jupyter notebook (.ipynb). A notebook is |
| 17 | // JSON with a "cells" array; editing it with edit_file means matching escaped |
| 18 | // JSON by hand, which is fragile. This tool targets a cell by index (or id) and |
| 19 | // replaces, inserts, or deletes it, re-serialising so the JSON stays valid and |
| 20 | // unrelated cells, outputs, and top-level metadata are preserved. |
| 21 | // |
| 22 | // roots, when non-empty, confines the target to the workspace (see confine); |
| 23 | // guard rejects Reasonix session-data targets (see SessionDataGuard); the |
| 24 | // zero value registered at init is unconfined and is overridden per run by |
| 25 | // ConfineWriters. workDir, when non-empty, is the directory a relative path |
| 26 | // resolves against (see resolveIn). |
| 27 | type notebookEdit struct { |
| 28 | roots []string |
| 29 | rootSet *sandbox.WritableRootSet |
| 30 | guard SessionDataGuard |
| 31 | managed ManagedConfigPaths |
| 32 | workDir string |
| 33 | overlay FileOverlay |
| 34 | } |
| 35 | |
| 36 | func (notebookEdit) Name() string { return "notebook_edit" } |
| 37 | |
| 38 | func (notebookEdit) ReadOnly() bool { return false } |
| 39 | |
| 40 | func (n notebookEdit) DeclareWriteAccess(args json.RawMessage) (tool.WriteAccessDeclaration, error) { |
| 41 | return declareFilePathWriteAccess(n.workDir, args) |
| 42 | } |
| 43 | |
| 44 | func (notebookEdit) Description() string { |
| 45 | return "Edit one cell of a Jupyter notebook (.ipynb). Target a cell by 0-based " + |
| 46 | "cell_number (or cell_id). edit_mode: \"replace\" (default) swaps the cell's " + |
| 47 | "source; \"insert\" adds a new cell after cell_number (use -1 to prepend at the " + |
| 48 | "top), taking cell_type and new_source; \"delete\" removes the cell. cell_type is " + |
| 49 | "\"code\" or \"markdown\" (required for insert). Editing a code cell clears its " + |
| 50 | "outputs. Prefer this over edit_file for notebooks — it keeps the JSON valid." |
| 51 | } |
| 52 | |
| 53 | func (notebookEdit) Schema() json.RawMessage { |
| 54 | return json.RawMessage(`{ |
| 55 | "type": "object", |
| 56 | "properties": { |
| 57 | "path": {"type": "string", "description": "Path to the .ipynb notebook."}, |
| 58 | "cell_number": {"type": "integer", "description": "0-based index of the target cell. For insert, the new cell goes after this one (-1 prepends)."}, |
| 59 | "cell_id": {"type": "string", "description": "Target the cell by its id instead of cell_number (replace/delete)."}, |
| 60 | "new_source": {"type": "string", "description": "The cell's new source text (replace/insert)."}, |
| 61 | "cell_type": {"type": "string", "enum": ["code", "markdown"], "description": "Cell type for insert (and optional retype on replace)."}, |
| 62 | "edit_mode": {"type": "string", "enum": ["replace", "insert", "delete"], "description": "replace (default), insert, or delete."} |
| 63 | }, |
| 64 | "required": ["path"] |
| 65 | }`) |
| 66 | } |
| 67 | |
| 68 | type notebookArgs struct { |
| 69 | Path string `json:"path"` |
| 70 | CellNumber *int `json:"cell_number"` |
| 71 | CellID string `json:"cell_id"` |
| 72 | NewSource string `json:"new_source"` |
| 73 | CellType string `json:"cell_type"` |
| 74 | EditMode string `json:"edit_mode"` |
| 75 | } |
| 76 | |
| 77 | // notebook is the minimal .ipynb shape we touch. Unknown top-level keys |
| 78 | // (metadata, nbformat, …) and unknown per-cell keys are preserved verbatim via |
| 79 | // json.RawMessage round-tripping. |
| 80 | type notebook struct { |
| 81 | rest map[string]json.RawMessage |
| 82 | cells []map[string]json.RawMessage |
| 83 | } |
| 84 | |
| 85 | func (n notebookEdit) Execute(ctx context.Context, raw json.RawMessage) (string, error) { |
| 86 | a, err := parseNotebookArgs(raw) |
| 87 | if err != nil { |
| 88 | return "", err |
| 89 | } |
| 90 | a.Path = resolveIn(n.workDir, a.Path) |
| 91 | if err := confineWrite(ctx, effectiveWriteRoots(ctx, n.rootSet, n.roots), n.guard, n.managed, a.Path); err != nil { |
| 92 | return "", err |
| 93 | } |
| 94 | unlock := lockMutationPath(a.Path) |
| 95 | defer unlock() |
| 96 | src, err := readEditSource(ctx, n.overlay, a.Path) |
| 97 | if err != nil { |
| 98 | return "", fmt.Errorf("read %s: %w", a.Path, err) |
| 99 | } |
| 100 | if err := src.requireObserved(ctx, n.overlay, a.Path); err != nil { |
| 101 | return "", err |
| 102 | } |
| 103 | nb, err := parseNotebook([]byte(src.content)) |
| 104 | if err != nil { |
| 105 | return "", fmt.Errorf("%s: %w", a.Path, err) |
| 106 | } |
| 107 | |
| 108 | idx, summary, err := applyNotebookEdit(nb, a) |
| 109 | if err != nil { |
| 110 | return "", err |
| 111 | } |
| 112 | |
| 113 | out, err := nb.marshal() |
| 114 | if err != nil { |
| 115 | return "", err |
| 116 | } |
| 117 | if err := src.write(ctx, n.overlay, a.Path, string(out)); err != nil { |
| 118 | return "", fmt.Errorf("write %s: %w", a.Path, err) |
| 119 | } |
| 120 | return fmt.Sprintf("%s in %s (cell %d; %d cells total)", summary, a.Path, idx, len(nb.cells)), nil |
| 121 | } |
| 122 | |
| 123 | // Preview implements tool.Previewer so a checkpoint can snapshot the notebook's |
| 124 | // before/after for rewind. It mirrors Execute's transformation exactly but never |
| 125 | // writes — same arg parsing and targeting rules, so the previewed change equals |
| 126 | // what Execute would persist. |
| 127 | func (n notebookEdit) Preview(ctx context.Context, raw json.RawMessage) (diff.Change, error) { |
| 128 | a, err := parseNotebookArgs(raw) |
| 129 | if err != nil { |
| 130 | return diff.Change{}, err |
| 131 | } |
| 132 | a.Path = resolveIn(n.workDir, a.Path) |
| 133 | if err := confinePreview(effectiveWriteRoots(ctx, n.rootSet, n.roots), n.guard, n.managed, a.Path); err != nil { |
| 134 | return diff.Change{}, err |
| 135 | } |
| 136 | src, err := readEditSource(ctx, n.overlay, a.Path) |
| 137 | if err != nil { |
| 138 | return diff.Change{}, fmt.Errorf("read %s: %w", a.Path, err) |
| 139 | } |
| 140 | nb, err := parseNotebook([]byte(src.content)) |
| 141 | if err != nil { |
| 142 | return diff.Change{}, fmt.Errorf("%s: %w", a.Path, err) |
| 143 | } |
| 144 | if _, _, err := applyNotebookEdit(nb, a); err != nil { |
| 145 | return diff.Change{}, err |
| 146 | } |
| 147 | out, err := nb.marshal() |
| 148 | if err != nil { |
| 149 | return diff.Change{}, err |
| 150 | } |
| 151 | return diff.Build(a.Path, src.content, string(out), diff.Modify), nil |
| 152 | } |
| 153 | |
| 154 | func parseNotebookArgs(raw json.RawMessage) (notebookArgs, error) { |
| 155 | var a notebookArgs |
| 156 | if err := json.Unmarshal(raw, &a); err != nil { |
| 157 | return a, fmt.Errorf("invalid args: %w", err) |
| 158 | } |
| 159 | // Be forgiving about the source field: models reach for the write_file/edit_file |
| 160 | // vocabulary ("content"/"source"/"new_string"). Accept those as new_source when |
| 161 | // new_source itself wasn't given, so a near-miss call succeeds instead of looping. |
| 162 | if a.NewSource == "" { |
| 163 | var alias struct { |
| 164 | Content string `json:"content"` |
| 165 | Source string `json:"source"` |
| 166 | NewString string `json:"new_string"` |
| 167 | } |
| 168 | _ = json.Unmarshal(raw, &alias) |
| 169 | switch { |
| 170 | case alias.Content != "": |
| 171 | a.NewSource = alias.Content |
| 172 | case alias.Source != "": |
| 173 | a.NewSource = alias.Source |
| 174 | case alias.NewString != "": |
| 175 | a.NewSource = alias.NewString |
| 176 | } |
| 177 | } |
| 178 | if a.Path == "" { |
| 179 | return a, fmt.Errorf("path is required") |
| 180 | } |
| 181 | if a.EditMode == "" { |
| 182 | a.EditMode = "replace" |
| 183 | } |
| 184 | switch a.EditMode { |
| 185 | case "replace", "insert", "delete": |
| 186 | default: |
| 187 | return a, fmt.Errorf("edit_mode must be replace, insert, or delete (got %q)", a.EditMode) |
| 188 | } |
| 189 | return a, nil |
| 190 | } |
| 191 | |
| 192 | // applyNotebookEdit mutates nb.cells per the args and returns the affected index |
| 193 | // and a one-line summary. Cell targeting is by cell_id when set, else cell_number. |
| 194 | func applyNotebookEdit(nb *notebook, a notebookArgs) (int, string, error) { |
| 195 | if a.EditMode == "insert" { |
| 196 | if a.CellType == "" { |
| 197 | return 0, "", fmt.Errorf("cell_type is required for insert") |
| 198 | } |
| 199 | after := -1 |
| 200 | if a.CellNumber != nil { |
| 201 | after = *a.CellNumber |
| 202 | } |
| 203 | if after < -1 || after >= len(nb.cells) { |
| 204 | return 0, "", fmt.Errorf("cell_number %d out of range for insert (notebook has %d cells; use -1 to prepend)", after, len(nb.cells)) |
| 205 | } |
| 206 | cell := newCell(a.CellType, a.NewSource) |
| 207 | at := after + 1 // insert after `after`; -1 → prepend at 0 |
| 208 | nb.cells = append(nb.cells[:at], append([]map[string]json.RawMessage{cell}, nb.cells[at:]...)...) |
| 209 | return at, "inserted " + a.CellType + " cell", nil |
| 210 | } |
| 211 | |
| 212 | idx, err := nb.targetIndex(a) |
| 213 | if err != nil { |
| 214 | return 0, "", err |
| 215 | } |
| 216 | if a.EditMode == "delete" { |
| 217 | nb.cells = append(nb.cells[:idx], nb.cells[idx+1:]...) |
| 218 | return idx, "deleted cell", nil |
| 219 | } |
| 220 | // replace |
| 221 | setCellSource(nb.cells[idx], a.NewSource) |
| 222 | if a.CellType != "" { |
| 223 | nb.cells[idx]["cell_type"] = jsonString(a.CellType) |
| 224 | } |
| 225 | normalizeOutputs(nb.cells[idx], cellTypeOf(nb.cells[idx])) |
| 226 | return idx, "replaced cell source", nil |
| 227 | } |
| 228 | |
| 229 | // targetIndex resolves the cell to act on: cell_id wins when given, else |
| 230 | // cell_number (which must be in range). |
| 231 | func (nb *notebook) targetIndex(a notebookArgs) (int, error) { |
| 232 | if a.CellID != "" { |
| 233 | for i, c := range nb.cells { |
| 234 | if cellID(c) == a.CellID { |
| 235 | return i, nil |
| 236 | } |
| 237 | } |
| 238 | return 0, fmt.Errorf("no cell with id %q", a.CellID) |
| 239 | } |
| 240 | if a.CellNumber == nil { |
| 241 | // A one-cell notebook is unambiguous: default to cell 0 rather than forcing |
| 242 | // the caller to restate it. With more than one cell, require an explicit target. |
| 243 | if len(nb.cells) == 1 { |
| 244 | return 0, nil |
| 245 | } |
| 246 | return 0, fmt.Errorf("cell_number or cell_id is required for %s (notebook has %d cells; pass the 0-based cell_number)", a.EditMode, len(nb.cells)) |
| 247 | } |
| 248 | n := *a.CellNumber |
| 249 | if n < 0 || n >= len(nb.cells) { |
| 250 | return 0, fmt.Errorf("cell_number %d out of range (notebook has %d cells)", n, len(nb.cells)) |
| 251 | } |
| 252 | return n, nil |
| 253 | } |
| 254 | |
| 255 | // parseNotebook decodes just enough of the .ipynb to edit cells while preserving |
| 256 | // every other key (top-level and per-cell) verbatim for re-serialisation. |
| 257 | func parseNotebook(data []byte) (*notebook, error) { |
| 258 | var top map[string]json.RawMessage |
| 259 | if err := json.Unmarshal(data, &top); err != nil { |
| 260 | return nil, fmt.Errorf("not valid notebook JSON: %w", err) |
| 261 | } |
| 262 | rawCells, ok := top["cells"] |
| 263 | if !ok { |
| 264 | return nil, fmt.Errorf("no \"cells\" array — not a notebook") |
| 265 | } |
| 266 | var cells []map[string]json.RawMessage |
| 267 | if err := json.Unmarshal(rawCells, &cells); err != nil { |
| 268 | return nil, fmt.Errorf("cells is not an array of objects: %w", err) |
| 269 | } |
| 270 | return ¬ebook{rest: top, cells: cells}, nil |
| 271 | } |
| 272 | |
| 273 | // marshal re-serialises the notebook with the edited cells, pretty-printed with |
| 274 | // the one-space indent Jupyter uses, and a trailing newline. |
| 275 | func (nb *notebook) marshal() ([]byte, error) { |
| 276 | cellsJSON, err := json.Marshal(nb.cells) |
| 277 | if err != nil { |
| 278 | return nil, err |
| 279 | } |
| 280 | nb.rest["cells"] = cellsJSON |
| 281 | out, err := json.MarshalIndent(nb.rest, "", " ") |
| 282 | if err != nil { |
| 283 | return nil, err |
| 284 | } |
| 285 | return append(out, '\n'), nil |
| 286 | } |
| 287 | |
| 288 | // newCell builds a fresh cell. Jupyter stores source as an array of lines (each |
| 289 | // ending in \n except the last); outputs/execution_count exist only for code. |
| 290 | func newCell(cellType, source string) map[string]json.RawMessage { |
| 291 | c := map[string]json.RawMessage{ |
| 292 | "cell_type": jsonString(cellType), |
| 293 | "metadata": json.RawMessage(`{}`), |
| 294 | "source": sourceLines(source), |
| 295 | } |
| 296 | if cellType == "code" { |
| 297 | c["outputs"] = json.RawMessage(`[]`) |
| 298 | c["execution_count"] = json.RawMessage(`null`) |
| 299 | } |
| 300 | return c |
| 301 | } |
| 302 | |
| 303 | func setCellSource(cell map[string]json.RawMessage, source string) { |
| 304 | cell["source"] = sourceLines(source) |
| 305 | } |
| 306 | |
| 307 | // normalizeOutputs makes a cell's output fields match its (possibly just-retyped) |
| 308 | // type: a code cell's stale results are cleared; a markdown cell must not carry |
| 309 | // outputs/execution_count at all, so a code→markdown retype drops them. |
| 310 | func normalizeOutputs(cell map[string]json.RawMessage, cellType string) { |
| 311 | if cellType == "markdown" { |
| 312 | delete(cell, "outputs") |
| 313 | delete(cell, "execution_count") |
| 314 | return |
| 315 | } |
| 316 | cell["outputs"] = json.RawMessage(`[]`) |
| 317 | cell["execution_count"] = json.RawMessage(`null`) |
| 318 | } |
| 319 | |
| 320 | func cellTypeOf(cell map[string]json.RawMessage) string { |
| 321 | var t string |
| 322 | _ = json.Unmarshal(cell["cell_type"], &t) |
| 323 | return t |
| 324 | } |
| 325 | |
| 326 | func cellID(cell map[string]json.RawMessage) string { |
| 327 | raw, ok := cell["id"] |
| 328 | if !ok { |
| 329 | return "" |
| 330 | } |
| 331 | var id string |
| 332 | _ = json.Unmarshal(raw, &id) |
| 333 | return id |
| 334 | } |
| 335 | |
| 336 | // sourceLines encodes a string as Jupyter's line-array source form: split on |
| 337 | // newlines, keeping the \n on every line but the last (matching nbformat). |
| 338 | func sourceLines(s string) json.RawMessage { |
| 339 | if s == "" { |
| 340 | return json.RawMessage(`[]`) |
| 341 | } |
| 342 | parts := strings.SplitAfter(s, "\n") |
| 343 | // SplitAfter leaves a trailing "" when s ends in \n; drop it. |
| 344 | if parts[len(parts)-1] == "" { |
| 345 | parts = parts[:len(parts)-1] |
| 346 | } |
| 347 | b, _ := json.Marshal(parts) |
| 348 | return b |
| 349 | } |
| 350 | |
| 351 | func jsonString(s string) json.RawMessage { |
| 352 | b, _ := json.Marshal(s) |
| 353 | return b |
| 354 | } |
| 355 |