diff --git a/docs/developers/sdk-java.md b/docs/developers/sdk-java.md index d3eab983c..8e2a8dbfe 100644 --- a/docs/developers/sdk-java.md +++ b/docs/developers/sdk-java.md @@ -144,7 +144,7 @@ The SDK supports different permission modes for controlling tool execution: - **`default`**: Write tools are denied unless approved via `canUseTool` callback or in `allowedTools`. Read-only tools execute without confirmation. - **`plan`**: Blocks all write tools, instructing AI to present a plan first. -- **`auto-edit`**: Auto-approve edit tools (edit, write_file) while other tools require confirmation. +- **`auto-edit`**: Auto-approve edit tools (`edit`, `write_file`, `notebook_edit`) while other tools require confirmation. - **`yolo`**: All tools execute automatically without confirmation. ### Session Event Consumers and Assistant Content Consumers diff --git a/docs/developers/sdk-typescript.md b/docs/developers/sdk-typescript.md index 8ba590997..a3de0f2e5 100644 --- a/docs/developers/sdk-typescript.md +++ b/docs/developers/sdk-typescript.md @@ -158,7 +158,7 @@ The SDK supports different permission modes for controlling tool execution: - **`default`**: Write tools are denied unless approved via `canUseTool` callback or in `allowedTools`. Read-only tools execute without confirmation. - **`plan`**: Blocks all write tools, instructing AI to present a plan first. -- **`auto-edit`**: Auto-approve edit tools (edit, write_file) while other tools require confirmation. +- **`auto-edit`**: Auto-approve edit tools (`edit`, `write_file`, `notebook_edit`) while other tools require confirmation. - **`yolo`**: All tools execute automatically without confirmation. ### Permission Priority Chain diff --git a/docs/developers/tools/file-system.md b/docs/developers/tools/file-system.md index 118f5e0b6..d07fd805c 100644 --- a/docs/developers/tools/file-system.md +++ b/docs/developers/tools/file-system.md @@ -44,7 +44,72 @@ Qwen Code provides a comprehensive suite of tools for interacting with the local - For other binary files: A message like `Cannot display content of binary file: /path/to/data.bin`. - **Confirmation:** No. -## 3. `write_file` (WriteFile) +### Jupyter notebook reads + +For Jupyter notebooks (`.ipynb`), `read_file` parses the notebook JSON and returns a structured, model-readable notebook view instead of raw JSON. The rendered output includes the notebook language, ordered cells, cell IDs, source, and summarized outputs. + +Notebook cells can then be edited with `notebook_edit`. The model should use the cell IDs shown by `read_file` when targeting a cell. + +`offset` and `limit` are not supported for `.ipynb` files. Notebook reads are treated as structured full-file reads; if the rendered notebook output is internally truncated because it is too large, `notebook_edit` will reject cell-level edits and ask you to reduce outputs or split the notebook before editing. + +## 3. `notebook_edit` (NotebookEdit) + +`notebook_edit` edits Jupyter notebook (`.ipynb`) files safely at the cell level. Use it instead of `edit` or `write_file` when changing notebook cells. + +- **Tool name:** `notebook_edit` +- **Display name:** NotebookEdit +- **File:** `notebook-edit.ts` +- **Parameters:** + - `notebook_path` (string, required): The absolute path to the `.ipynb` file. + - `cell_id` (string, optional): The target cell ID shown by `read_file`. Required for `replace` and `delete`. For `insert`, the new cell is inserted after this cell; if omitted, the new cell is inserted at the beginning. + - `new_source` (string, optional): The new cell source for `replace` and `insert`. Not required for `delete`. + - `cell_type` (`code` or `markdown`, optional): The cell type for inserted cells, or the target type when replacing a cell. + - `edit_mode` (`replace`, `insert`, or `delete`, optional): The edit operation. Defaults to `replace`. +- **Behavior:** + - Requires the notebook to have been read first with `read_file` in the current session. + - Targets cells using the IDs rendered by `read_file`, including real notebook cell IDs and displayed `cell-N` fallback IDs. + - Rejects ambiguous rendered cell IDs instead of guessing. + - For code cells, clears stale outputs and resets `execution_count` when source changes. + - Preserves notebook JSON formatting, line endings, encoding, and BOM where possible. + - Invalidates the prior-read state after structural edits when displayed fallback IDs can shift, so the next notebook edit requires a fresh `read_file`. +- **Output (`llmContent`):** A success message describing the edited notebook cell and, for non-delete operations, the updated source. +- **Confirmation:** Yes. Shows a notebook JSON diff and asks for user approval before writing, unless the current permission mode or rules auto-approve edit tools. + +### `notebook_edit` examples + +Replace a code cell: + +``` +notebook_edit( + notebook_path="/path/to/analysis.ipynb", + cell_id="load-data", + new_source="result = 41 + 1\nprint(result)" +) +``` + +Insert a markdown cell after an existing cell: + +``` +notebook_edit( + notebook_path="/path/to/analysis.ipynb", + edit_mode="insert", + cell_id="summary", + cell_type="markdown", + new_source="## Findings\n\nThe cleaned data is ready for modeling." +) +``` + +Delete a cell: + +``` +notebook_edit( + notebook_path="/path/to/analysis.ipynb", + edit_mode="delete", + cell_id="old-experiment" +) +``` + +## 4. `write_file` (WriteFile) `write_file` writes content to a specified file. If the file exists, it will be overwritten. If the file doesn't exist, it (and any necessary parent directories) will be created. @@ -56,11 +121,12 @@ Qwen Code provides a comprehensive suite of tools for interacting with the local - `content` (string, required): The content to write into the file. - **Behavior:** - Writes the provided `content` to the `file_path`. + - Does not write raw Jupyter notebook JSON. Use `notebook_edit` for `.ipynb` cell edits. - Creates parent directories if they don't exist. - **Output (`llmContent`):** A success message, e.g., `Successfully overwrote file: /path/to/your/file.txt` or `Successfully created and wrote to new file: /path/to/new/file.txt`. - **Confirmation:** Yes. Shows a diff of changes and asks for user approval before writing. -## 4. `glob` (Glob) +## 5. `glob` (Glob) `glob` finds files matching specific glob patterns (e.g., `src/**/*.ts`, `*.md`), returning absolute paths sorted by modification time (newest first). @@ -78,7 +144,7 @@ Qwen Code provides a comprehensive suite of tools for interacting with the local - **Output (`llmContent`):** A message like: `Found 5 file(s) matching "*.ts" within /path/to/search/dir, sorted by modification time (newest first):\n---\n/path/to/file1.ts\n/path/to/subdir/file2.ts\n---\n[95 files truncated] ...` - **Confirmation:** No. -## 5. `grep_search` (Grep) +## 6. `grep_search` (Grep) `grep_search` searches for a regular expression pattern within the content of files in a specified directory. Can filter files by a glob pattern. Returns the lines containing matches, along with their file paths and line numbers. @@ -131,7 +197,7 @@ Search for a pattern with file filtering and custom result limiting: grep_search(pattern="function", glob="*.js", limit=10) ``` -## 6. `edit` (Edit) +## 7. `edit` (Edit) `edit` replaces text within a file. By default it requires `old_string` to match a single unique location; set `replace_all` to `true` when you intentionally want to change every occurrence. This tool is designed for precise, targeted changes and requires significant context around the `old_string` to ensure it modifies the correct location. @@ -148,6 +214,7 @@ grep_search(pattern="function", glob="*.js", limit=10) - `replace_all` (boolean, optional): Replace all occurrences of `old_string`. Defaults to `false`. - **Behavior:** + - Does not edit raw Jupyter notebook JSON. Use `notebook_edit` for `.ipynb` cell edits. - If `old_string` is empty and `file_path` does not exist, creates a new file with `new_string` as content. - If `old_string` is provided, it reads the `file_path` and attempts to find exactly one occurrence unless `replace_all` is true. - If the match is unique (or `replace_all` is true), it replaces the text with `new_string`. diff --git a/docs/users/configuration/settings.md b/docs/users/configuration/settings.md index f6b71a076..48d26e8ef 100644 --- a/docs/users/configuration/settings.md +++ b/docs/users/configuration/settings.md @@ -298,6 +298,8 @@ The first matching rule wins. Rules use the format `"ToolName"` or `"ToolName(sp | `Read`, `ReadFile` | `read_file` | Meta-category — see below | | `Edit`, `EditFile` | `edit` | Meta-category — see below | | `Write`, `WriteFile` | `write_file` | | +| `NotebookEdit` | `notebook_edit` | | +| `NotebookEditTool` | `notebook_edit` | | | `Grep`, `SearchFiles` | `grep_search` | | | `Glob`, `FindFiles` | `glob` | | | `ListFiles` | `list_directory` | | @@ -312,7 +314,7 @@ Some rule names automatically cover multiple tools: | Rule name | Tools covered | | --------- | ---------------------------------------------------- | | `Read` | `read_file`, `grep_search`, `glob`, `list_directory` | -| `Edit` | `edit`, `write_file` | +| `Edit` | `edit`, `write_file`, `notebook_edit` | > [!important] > `Read(/path/**)` matches **all four** read tools (file read, grep, glob, and directory listing). @@ -603,42 +605,42 @@ For sandbox image selection, precedence is: ### Command-Line Arguments Table -| Argument | Alias | Description | Possible Values | Notes | -| ---------------------------- | ----- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -------------------------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| `--model` | `-m` | Specifies the Qwen model to use for this session. | Model name | Example: `npm start -- --model qwen3-coder-plus` | -| `--prompt` | `-p` | Used to pass a prompt directly to the command. This invokes Qwen Code in a non-interactive mode. | Your prompt text | For scripting examples, use the `--output-format json` flag to get structured output. | -| `--prompt-interactive` | `-i` | Starts an interactive session with the provided prompt as the initial input. | Your prompt text | The prompt is processed within the interactive session, not before it. Cannot be used when piping input from stdin. Example: `qwen -i "explain this code"` | -| `--system-prompt` | | Overrides the built-in main session system prompt for this run. | Your prompt text | Loaded context files such as `QWEN.md` are still appended after this override. Can be combined with `--append-system-prompt`. | -| `--append-system-prompt` | | Appends extra instructions to the main session system prompt for this run. | Your prompt text | Applied after the built-in prompt and loaded context files. Can be combined with `--system-prompt`. See [Headless Mode](../features/headless) for examples. | -| `--output-format` | `-o` | Specifies the format of the CLI output for non-interactive mode. | `text`, `json`, `stream-json` | `text`: (Default) The standard human-readable output. `json`: A machine-readable JSON output emitted at the end of execution. `stream-json`: Streaming JSON messages emitted as they occur during execution. For structured output and scripting, use the `--output-format json` or `--output-format stream-json` flag. See [Headless Mode](../features/headless) for detailed information. | -| `--input-format` | | Specifies the format consumed from standard input. | `text`, `stream-json` | `text`: (Default) Standard text input from stdin or command-line arguments. `stream-json`: JSON message protocol via stdin for bidirectional communication. Requirement: `--input-format stream-json` requires `--output-format stream-json` to be set. When using `stream-json`, stdin is reserved for protocol messages. See [Headless Mode](../features/headless) for detailed information. | -| `--include-partial-messages` | | Include partial assistant messages when using `stream-json` output format. When enabled, emits stream events (message_start, content_block_delta, etc.) as they occur during streaming. | | Default: `false`. Requirement: Requires `--output-format stream-json` to be set. See [Headless Mode](../features/headless) for detailed information about stream events. | -| `--sandbox` | `-s` | Enables sandbox mode for this session. | | | -| `--sandbox-image` | | Sets the sandbox image URI. | | | -| `--debug` | `-d` | Enables debug mode for this session, providing more verbose output. | | | -| `--all-files` | `-a` | If set, recursively includes all files within the current directory as context for the prompt. | | | -| `--help` | `-h` | Displays help information about command-line arguments. | | | -| `--show-memory-usage` | | Displays the current memory usage. | | | -| `--yolo` | | Enables YOLO mode, which automatically approves all tool calls. | | | -| `--approval-mode` | | Sets the approval mode for tool calls. | `plan`, `default`, `auto-edit`, `yolo` | Supported modes: `plan`: Analyze only—do not modify files or execute commands. `default`: Require approval for file edits or shell commands (default behavior). `auto-edit`: Automatically approve edit tools (edit, write_file) while prompting for others. `yolo`: Automatically approve all tool calls (equivalent to `--yolo`). Cannot be used together with `--yolo`. Use `--approval-mode=yolo` instead of `--yolo` for the new unified approach. Example: `qwen --approval-mode auto-edit`
See more about [Approval Mode](../features/approval-mode). | -| `--allowed-tools` | | A comma-separated list of tool names that will bypass the confirmation dialog. | Tool names | Example: `qwen --allowed-tools "Shell(git status)"` | -| `--disabled-slash-commands` | | Slash command names to hide/disable (comma-separated or repeated). Unioned with the `slashCommands.disabled` setting and the `QWEN_DISABLED_SLASH_COMMANDS` environment variable. Matched case-insensitively against the final command name. | Command names | Example: `qwen --disabled-slash-commands "auth,mcp,extensions"` | -| `--telemetry` | | Enables [telemetry](/developers/development/telemetry). | | | -| `--telemetry-target` | | Sets the telemetry target. | | See [telemetry](/developers/development/telemetry) for more information. | -| `--telemetry-otlp-endpoint` | | Sets the OTLP endpoint for telemetry. | | See [telemetry](../../developers/development/telemetry) for more information. | -| `--telemetry-otlp-protocol` | | Sets the OTLP protocol for telemetry (`grpc` or `http`). | | Defaults to `grpc`. See [telemetry](../../developers/development/telemetry) for more information. | -| `--telemetry-log-prompts` | | Enables logging of prompts for telemetry. | | See [telemetry](../../developers/development/telemetry) for more information. | -| `--checkpointing` | | Enables [checkpointing](../features/checkpointing). | | | -| `--acp` | | Enables ACP mode (Agent Client Protocol). Useful for IDE/editor integrations like [Zed](../integration-zed). | | Stable. Replaces the deprecated `--experimental-acp` flag. | -| `--experimental-lsp` | | Enables experimental [LSP (Language Server Protocol)](../features/lsp) feature for code intelligence (go-to-definition, find references, diagnostics, etc.). | | Experimental. Requires language servers to be installed. | -| `--extensions` | `-e` | Specifies a list of extensions to use for the session. | Extension names | If not provided, all available extensions are used. Use the special term `qwen -e none` to disable all extensions. Example: `qwen -e my-extension -e my-other-extension` | -| `--list-extensions` | `-l` | Lists all available extensions and exits. | | | -| `--proxy` | | Sets the proxy for the CLI. | Proxy URL | Example: `--proxy http://localhost:7890`. | -| `--include-directories` | | Includes additional directories in the workspace for multi-directory support. | Directory paths | Can be specified multiple times or as comma-separated values. 5 directories can be added at maximum. Example: `--include-directories /path/to/project1,/path/to/project2` or `--include-directories /path/to/project1 --include-directories /path/to/project2` | -| `--screen-reader` | | Enables screen reader mode, which adjusts the TUI for better compatibility with screen readers. | | | -| `--version` | | Displays the version of the CLI. | | | -| `--openai-logging` | | Enables logging of OpenAI API calls for debugging and analysis. | | This flag overrides the `enableOpenAILogging` setting in `settings.json`. | -| `--openai-logging-dir` | | Sets a custom directory path for OpenAI API logs. | Directory path | This flag overrides the `openAILoggingDir` setting in `settings.json`. Supports absolute paths, relative paths, and `~` expansion. Example: `qwen --openai-logging-dir "~/qwen-logs" --openai-logging` | +| Argument | Alias | Description | Possible Values | Notes | +| ---------------------------- | ----- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | +| `--model` | `-m` | Specifies the Qwen model to use for this session. | Model name | Example: `npm start -- --model qwen3-coder-plus` | +| `--prompt` | `-p` | Used to pass a prompt directly to the command. This invokes Qwen Code in a non-interactive mode. | Your prompt text | For scripting examples, use the `--output-format json` flag to get structured output. | +| `--prompt-interactive` | `-i` | Starts an interactive session with the provided prompt as the initial input. | Your prompt text | The prompt is processed within the interactive session, not before it. Cannot be used when piping input from stdin. Example: `qwen -i "explain this code"` | +| `--system-prompt` | | Overrides the built-in main session system prompt for this run. | Your prompt text | Loaded context files such as `QWEN.md` are still appended after this override. Can be combined with `--append-system-prompt`. | +| `--append-system-prompt` | | Appends extra instructions to the main session system prompt for this run. | Your prompt text | Applied after the built-in prompt and loaded context files. Can be combined with `--system-prompt`. See [Headless Mode](../features/headless) for examples. | +| `--output-format` | `-o` | Specifies the format of the CLI output for non-interactive mode. | `text`, `json`, `stream-json` | `text`: (Default) The standard human-readable output. `json`: A machine-readable JSON output emitted at the end of execution. `stream-json`: Streaming JSON messages emitted as they occur during execution. For structured output and scripting, use the `--output-format json` or `--output-format stream-json` flag. See [Headless Mode](../features/headless) for detailed information. | +| `--input-format` | | Specifies the format consumed from standard input. | `text`, `stream-json` | `text`: (Default) Standard text input from stdin or command-line arguments. `stream-json`: JSON message protocol via stdin for bidirectional communication. Requirement: `--input-format stream-json` requires `--output-format stream-json` to be set. When using `stream-json`, stdin is reserved for protocol messages. See [Headless Mode](../features/headless) for detailed information. | +| `--include-partial-messages` | | Include partial assistant messages when using `stream-json` output format. When enabled, emits stream events (message_start, content_block_delta, etc.) as they occur during streaming. | | Default: `false`. Requirement: Requires `--output-format stream-json` to be set. See [Headless Mode](../features/headless) for detailed information about stream events. | +| `--sandbox` | `-s` | Enables sandbox mode for this session. | | | +| `--sandbox-image` | | Sets the sandbox image URI. | | | +| `--debug` | `-d` | Enables debug mode for this session, providing more verbose output. | | | +| `--all-files` | `-a` | If set, recursively includes all files within the current directory as context for the prompt. | | | +| `--help` | `-h` | Displays help information about command-line arguments. | | | +| `--show-memory-usage` | | Displays the current memory usage. | | | +| `--yolo` | | Enables YOLO mode, which automatically approves all tool calls. | | | +| `--approval-mode` | | Sets the approval mode for tool calls. | `plan`, `default`, `auto-edit`, `yolo` | Supported modes: `plan`: Analyze only—do not modify files or execute commands. `default`: Require approval for file edits or shell commands (default behavior). `auto-edit`: Automatically approve edit tools (`edit`, `write_file`, `notebook_edit`) while prompting for others. `yolo`: Automatically approve all tool calls (equivalent to `--yolo`). Cannot be used together with `--yolo`. Use `--approval-mode=yolo` instead of `--yolo` for the new unified approach. Example: `qwen --approval-mode auto-edit`
See more about [Approval Mode](../features/approval-mode). | +| `--allowed-tools` | | A comma-separated list of tool names that will bypass the confirmation dialog. | Tool names | Example: `qwen --allowed-tools "Shell(git status)"` | +| `--disabled-slash-commands` | | Slash command names to hide/disable (comma-separated or repeated). Unioned with the `slashCommands.disabled` setting and the `QWEN_DISABLED_SLASH_COMMANDS` environment variable. Matched case-insensitively against the final command name. | Command names | Example: `qwen --disabled-slash-commands "auth,mcp,extensions"` | +| `--telemetry` | | Enables [telemetry](/developers/development/telemetry). | | | +| `--telemetry-target` | | Sets the telemetry target. | | See [telemetry](/developers/development/telemetry) for more information. | +| `--telemetry-otlp-endpoint` | | Sets the OTLP endpoint for telemetry. | | See [telemetry](../../developers/development/telemetry) for more information. | +| `--telemetry-otlp-protocol` | | Sets the OTLP protocol for telemetry (`grpc` or `http`). | | Defaults to `grpc`. See [telemetry](../../developers/development/telemetry) for more information. | +| `--telemetry-log-prompts` | | Enables logging of prompts for telemetry. | | See [telemetry](../../developers/development/telemetry) for more information. | +| `--checkpointing` | | Enables [checkpointing](../features/checkpointing). | | | +| `--acp` | | Enables ACP mode (Agent Client Protocol). Useful for IDE/editor integrations like [Zed](../integration-zed). | | Stable. Replaces the deprecated `--experimental-acp` flag. | +| `--experimental-lsp` | | Enables experimental [LSP (Language Server Protocol)](../features/lsp) feature for code intelligence (go-to-definition, find references, diagnostics, etc.). | | Experimental. Requires language servers to be installed. | +| `--extensions` | `-e` | Specifies a list of extensions to use for the session. | Extension names | If not provided, all available extensions are used. Use the special term `qwen -e none` to disable all extensions. Example: `qwen -e my-extension -e my-other-extension` | +| `--list-extensions` | `-l` | Lists all available extensions and exits. | | | +| `--proxy` | | Sets the proxy for the CLI. | Proxy URL | Example: `--proxy http://localhost:7890`. | +| `--include-directories` | | Includes additional directories in the workspace for multi-directory support. | Directory paths | Can be specified multiple times or as comma-separated values. 5 directories can be added at maximum. Example: `--include-directories /path/to/project1,/path/to/project2` or `--include-directories /path/to/project1 --include-directories /path/to/project2` | +| `--screen-reader` | | Enables screen reader mode, which adjusts the TUI for better compatibility with screen readers. | | | +| `--version` | | Displays the version of the CLI. | | | +| `--openai-logging` | | Enables logging of OpenAI API calls for debugging and analysis. | | This flag overrides the `enableOpenAILogging` setting in `settings.json`. | +| `--openai-logging-dir` | | Sets a custom directory path for OpenAI API logs. | Directory path | This flag overrides the `openAILoggingDir` setting in `settings.json`. Supports absolute paths, relative paths, and `~` expansion. Example: `qwen --openai-logging-dir "~/qwen-logs" --openai-logging` | ## Context Files (Hierarchical Instructional Context) diff --git a/docs/users/features/approval-mode.md b/docs/users/features/approval-mode.md index 94bcfd97f..3f9849de5 100644 --- a/docs/users/features/approval-mode.md +++ b/docs/users/features/approval-mode.md @@ -163,6 +163,8 @@ You can review each proposed change and approve or reject it individually. Auto-Edit Mode instructs Qwen Code to automatically approve file edits while requiring manual approval for shell commands, ideal for accelerating development workflows while maintaining system safety. +Auto-approved edit tools include `edit`, `write_file`, and `notebook_edit`. + ### When to use Auto-Accept Edits Mode - **Daily development**: Ideal for most coding tasks diff --git a/integration-tests/cli/notebook-edit.test.ts b/integration-tests/cli/notebook-edit.test.ts new file mode 100644 index 000000000..746380c5c --- /dev/null +++ b/integration-tests/cli/notebook-edit.test.ts @@ -0,0 +1,239 @@ +/** + * @license + * Copyright 2026 Qwen Team + * SPDX-License-Identifier: Apache-2.0 + */ + +import path from 'node:path'; +import { afterEach, describe, expect, it } from 'vitest'; +import { + TestRig, + createToolCallErrorMessage, + printDebugInfo, + validateModelOutput, +} from '../test-helper.js'; + +type NotebookCell = { + id?: string; + cell_type: 'code' | 'markdown' | 'raw'; + metadata: Record; + source: string | string[]; + execution_count?: number | null; + outputs?: unknown[]; +}; + +type NotebookContent = { + cells: NotebookCell[]; + metadata: Record; + nbformat: number; + nbformat_minor: number; +}; + +const sourceText = (source: string | string[]) => + Array.isArray(source) ? source.join('') : source; + +const promptPath = (filePath: string) => filePath.split(path.sep).join('/'); + +const readNotebook = (rig: TestRig, fileName: string): NotebookContent => + JSON.parse(rig.readFile(fileName)) as NotebookContent; + +const baseNotebook = (cells: NotebookCell[]): NotebookContent => ({ + cells, + metadata: { + kernelspec: { + display_name: 'Python 3', + language: 'python', + name: 'python3', + }, + language_info: { + name: 'python', + }, + }, + nbformat: 4, + nbformat_minor: 5, +}); + +const expectReadThenNotebookEdit = (rig: TestRig, result: string) => { + const logs = rig.readToolLogs(); + const foundTools = logs.map((t) => t.toolRequest.name); + const readIndex = foundTools.findIndex((name) => name === 'read_file'); + const notebookEditIndex = foundTools.findIndex( + (name) => name === 'notebook_edit', + ); + + if (readIndex === -1 || notebookEditIndex === -1) { + printDebugInfo(rig, result, { foundTools }); + } + + expect( + readIndex, + createToolCallErrorMessage('read_file', foundTools, result), + ).toBeGreaterThanOrEqual(0); + expect( + notebookEditIndex, + createToolCallErrorMessage('notebook_edit', foundTools, result), + ).toBeGreaterThan(readIndex); +}; + +const expectNoSuccessfulRawNotebookWrites = ( + rig: TestRig, + notebookFileName: string, +) => { + const rawNotebookWrites = rig + .readToolLogs() + .filter( + (log) => + ['edit', 'write_file'].includes(log.toolRequest.name) && + log.toolRequest.success && + log.toolRequest.args.includes(notebookFileName), + ); + + expect(rawNotebookWrites).toEqual([]); +}; + +describe('notebook_edit integration', () => { + let rig: TestRig; + + afterEach(async () => { + await rig?.cleanup(); + }); + + it('replaces a code cell after reading the notebook and clears stale outputs', async () => { + rig = new TestRig(); + await rig.setup('notebook edit replace code cell clears outputs'); + + const fileName = 'analysis.ipynb'; + const notebookPath = rig.createFile( + fileName, + JSON.stringify( + baseNotebook([ + { + id: 'intro', + cell_type: 'markdown', + metadata: {}, + source: ['# Analysis\n'], + }, + { + id: 'load-data', + cell_type: 'code', + metadata: {}, + source: ['old_value = 1\n', 'print(old_value)\n'], + execution_count: 7, + outputs: [ + { + output_type: 'stream', + name: 'stdout', + text: ['1\n'], + }, + ], + }, + ]), + null, + 1, + ), + ); + + const prompt = `Read the notebook at ${promptPath(notebookPath)} with read_file first. +Then use notebook_edit, not edit or write_file, to replace the code cell whose id is load-data. +Set the new source exactly to: + +result = 41 + 1 +print(result) + +Do not change any other cell.`; + + const result = await rig.run(prompt); + + expectReadThenNotebookEdit(rig, result); + expectNoSuccessfulRawNotebookWrites(rig, fileName); + validateModelOutput(result, null, 'Notebook replace'); + + const notebook = readNotebook(rig, fileName); + const target = notebook.cells.find((cell) => cell.id === 'load-data'); + + expect(target).toBeDefined(); + expect(target?.cell_type).toBe('code'); + expect(sourceText(target!.source).trimEnd()).toBe( + 'result = 41 + 1\nprint(result)', + ); + expect(target?.execution_count).toBeNull(); + expect(target?.outputs).toEqual([]); + expect(sourceText(notebook.cells[0]!.source)).toBe('# Analysis\n'); + }); + + it('inserts a markdown cell and deletes a target cell using notebook_edit', async () => { + rig = new TestRig(); + await rig.setup('notebook edit insert and delete cells'); + + const fileName = 'workflow.ipynb'; + const notebookPath = rig.createFile( + fileName, + JSON.stringify( + baseNotebook([ + { + id: 'intro', + cell_type: 'markdown', + metadata: {}, + source: ['# Workflow\n'], + }, + { + id: 'remove-me', + cell_type: 'markdown', + metadata: {}, + source: ['This temporary cell should be deleted.\n'], + }, + { + id: 'calculate', + cell_type: 'code', + metadata: {}, + source: ['value = 10\n'], + execution_count: null, + outputs: [], + }, + ]), + null, + 1, + ), + ); + + const insertedMarkdown = + '## Inserted Note\nThis cell was inserted by NotebookEdit.'; + const prompt = `Read the notebook at ${promptPath(notebookPath)} with read_file first. +Then use notebook_edit, not edit or write_file, for both changes: +1. Insert a markdown cell after the cell whose id is intro. Its source must be exactly: + +${insertedMarkdown} + +2. Delete the cell whose id is remove-me. +Do not change the calculate code cell.`; + + const result = await rig.run(prompt); + + expectReadThenNotebookEdit(rig, result); + expectNoSuccessfulRawNotebookWrites(rig, fileName); + validateModelOutput(result, null, 'Notebook insert/delete'); + + const successfulNotebookEdits = rig + .readToolLogs() + .filter( + (log) => + log.toolRequest.name === 'notebook_edit' && log.toolRequest.success, + ); + expect(successfulNotebookEdits.length).toBeGreaterThanOrEqual(2); + + const notebook = readNotebook(rig, fileName); + const cellIds = notebook.cells.map((cell) => cell.id); + const insertedCell = notebook.cells.find((cell) => + sourceText(cell.source).includes('Inserted Note'), + ); + + expect(cellIds).toHaveLength(3); + expect(cellIds).not.toContain('remove-me'); + expect(notebook.cells[0]?.id).toBe('intro'); + expect(insertedCell).toBeDefined(); + expect(insertedCell?.cell_type).toBe('markdown'); + expect(sourceText(insertedCell!.source).trimEnd()).toBe(insertedMarkdown); + expect(notebook.cells.at(-1)?.id).toBe('calculate'); + expect(sourceText(notebook.cells.at(-1)!.source)).toBe('value = 10\n'); + }); +}); diff --git a/packages/cli/src/config/config.test.ts b/packages/cli/src/config/config.test.ts index 00246d146..972bec8c8 100644 --- a/packages/cli/src/config/config.test.ts +++ b/packages/cli/src/config/config.test.ts @@ -2388,6 +2388,7 @@ describe('loadCliConfig with includeDirectories', () => { expect(config.getCoreTools()).toEqual([ ToolNames.READ_FILE, ToolNames.EDIT, + ToolNames.NOTEBOOK_EDIT, ToolNames.SHELL, ]); expect(config.getDisableAllHooks()).toBe(true); @@ -2406,6 +2407,7 @@ describe('loadCliConfig with includeDirectories', () => { expect(config.getCoreTools()).toEqual([ ToolNames.READ_FILE, ToolNames.EDIT, + ToolNames.NOTEBOOK_EDIT, ToolNames.SHELL, ]); }); diff --git a/packages/core/src/config/config.test.ts b/packages/core/src/config/config.test.ts index 8a668cfe5..a9badd587 100644 --- a/packages/core/src/config/config.test.ts +++ b/packages/core/src/config/config.test.ts @@ -624,7 +624,12 @@ describe('Server Config (config.ts)', () => { (ToolRegistry.prototype.registerFactory as Mock).mock.calls.map( (call) => call[0], ), - ).toEqual([ToolNames.READ_FILE, ToolNames.EDIT, ToolNames.SHELL]); + ).toEqual([ + ToolNames.READ_FILE, + ToolNames.EDIT, + ToolNames.NOTEBOOK_EDIT, + ToolNames.SHELL, + ]); }); it('skips inline MCP discovery by default (progressive availability)', async () => { @@ -1736,20 +1741,26 @@ describe('Server Config (config.ts)', () => { expect(config.getCoreTools()).toEqual([ ToolNames.READ_FILE, ToolNames.EDIT, + ToolNames.NOTEBOOK_EDIT, ToolNames.SHELL, ]); expect( (registerToolMock as Mock).mock.calls.map((call) => call[0]), - ).toEqual([ToolNames.READ_FILE, ToolNames.EDIT, ToolNames.SHELL]); + ).toEqual([ + ToolNames.READ_FILE, + ToolNames.EDIT, + ToolNames.NOTEBOOK_EDIT, + ToolNames.SHELL, + ]); }); it('registers structured_output in bare mode when jsonSchema is set', async () => { - // Bare mode strips the toolset to READ_FILE/EDIT/SHELL, but the + // Bare mode strips the toolset to READ_FILE/EDIT/NOTEBOOK_EDIT/SHELL, but the // synthetic structured_output tool is the terminal contract for // --json-schema runs. Without it the model loops until // maxSessionTurns and exits via the "plain text" failure path — // expensive in tokens for what's almost always a CI use case. The - // synthetic tool must be registered alongside the bare three. + // synthetic tool must be registered alongside the bare toolset. const config = new Config({ ...baseParams, bareMode: true, @@ -1768,6 +1779,7 @@ describe('Server Config (config.ts)', () => { ).toEqual([ ToolNames.READ_FILE, ToolNames.EDIT, + ToolNames.NOTEBOOK_EDIT, ToolNames.SHELL, ToolNames.STRUCTURED_OUTPUT, ]); @@ -1796,8 +1808,8 @@ describe('Server Config (config.ts)', () => { ToolRegistry: { prototype: { registerFactory: Mock } }; } ).ToolRegistry.prototype.registerFactory; - // Initial bare init registers READ_FILE / EDIT / SHELL / - // STRUCTURED_OUTPUT (asserted by the test above). Reset so we can + // Initial bare init registers READ_FILE / EDIT / NOTEBOOK_EDIT / + // SHELL / STRUCTURED_OUTPUT (asserted by the test above). Reset so we can // observe ONLY the forSubAgent rebuild's calls. (registerToolMock as Mock).mockClear(); @@ -1811,10 +1823,11 @@ describe('Server Config (config.ts)', () => { (call) => call[0], ); expect(registeredNames).not.toContain(ToolNames.STRUCTURED_OUTPUT); - // The bare three still register so the subagent has its toolset. + // The bare tools still register so the subagent has its toolset. expect(registeredNames).toEqual([ ToolNames.READ_FILE, ToolNames.EDIT, + ToolNames.NOTEBOOK_EDIT, ToolNames.SHELL, ]); }); diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index 837927660..03ec0cd0e 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -736,6 +736,7 @@ export interface ConfigInitializeOptions { const DEFAULT_BARE_CORE_TOOLS = [ ToolNames.READ_FILE, ToolNames.EDIT, + ToolNames.NOTEBOOK_EDIT, ToolNames.SHELL, ]; @@ -3593,6 +3594,10 @@ export class Config { const { EditTool } = await import('../tools/edit.js'); return new EditTool(this); }); + await registerLazy(ToolNames.NOTEBOOK_EDIT, async () => { + const { NotebookEditTool } = await import('../tools/notebook-edit.js'); + return new NotebookEditTool(this); + }); await registerLazy(ToolNames.SHELL, async () => { const { ShellTool } = await import('../tools/shell.js'); return new ShellTool(this); @@ -3677,6 +3682,10 @@ export class Config { const { EditTool } = await import('../tools/edit.js'); return new EditTool(this); }); + await registerLazy(ToolNames.NOTEBOOK_EDIT, async () => { + const { NotebookEditTool } = await import('../tools/notebook-edit.js'); + return new NotebookEditTool(this); + }); await registerLazy(ToolNames.WRITE_FILE, async () => { const { WriteFileTool } = await import('../tools/write-file.js'); return new WriteFileTool(this); diff --git a/packages/core/src/core/coreToolScheduler.test.ts b/packages/core/src/core/coreToolScheduler.test.ts index ec9be0ce3..c1bacaf8e 100644 --- a/packages/core/src/core/coreToolScheduler.test.ts +++ b/packages/core/src/core/coreToolScheduler.test.ts @@ -5552,6 +5552,14 @@ describe('extractToolFilePaths', () => { ]); }); + it('extracts notebook_path for notebook_edit', () => { + expect( + extractToolFilePaths('notebook_edit', { + notebook_path: '/proj/analysis.ipynb', + }), + ).toEqual(['/proj/analysis.ipynb']); + }); + it('extracts filePath for lsp (camelCase convention)', () => { expect(extractToolFilePaths('lsp', { filePath: '/proj/b.ts' })).toEqual([ '/proj/b.ts', diff --git a/packages/core/src/core/coreToolScheduler.ts b/packages/core/src/core/coreToolScheduler.ts index 804d6f4cd..859128497 100644 --- a/packages/core/src/core/coreToolScheduler.ts +++ b/packages/core/src/core/coreToolScheduler.ts @@ -314,6 +314,7 @@ const FS_PATH_TOOL_NAMES: ReadonlySet = new Set([ ToolNames.GLOB, ToolNames.LS, ToolNames.LSP, + ToolNames.NOTEBOOK_EDIT, ]); function canonicalToolName(toolName: string): string { @@ -383,6 +384,7 @@ function pushLspPathCandidate(out: string[], v: unknown): void { * Per-tool dispatcher because the field name and shape differ: * * - read_file / edit / write_file → `file_path` + * - notebook_edit → `notebook_path` * - list_directory → `path` (search root) * - glob → `path` (search root, optional) + `pattern` (path-shaped * selector); `/` is the effective glob walked @@ -502,6 +504,13 @@ export function extractToolFilePaths( case ToolNames.READ_FILE: case ToolNames.EDIT: case ToolNames.WRITE_FILE: + push(obj['file_path']); + return out; + + case ToolNames.NOTEBOOK_EDIT: + push(obj['notebook_path']); + return out; + default: push(obj['file_path']); return out; diff --git a/packages/core/src/core/permission-helpers.test.ts b/packages/core/src/core/permission-helpers.test.ts index 6197c1b17..e038eb65a 100644 --- a/packages/core/src/core/permission-helpers.test.ts +++ b/packages/core/src/core/permission-helpers.test.ts @@ -74,4 +74,19 @@ describe('buildPermissionCheckContext', () => { command: `/bin/bash -c 'tail -f ./app.log' && rm -rf /tmp/owned`, }); }); + + it('uses notebook_path as the file path for notebook_edit', () => { + expect( + buildPermissionCheckContext( + 'notebook_edit', + { + notebook_path: '/project/analysis.ipynb', + }, + '/project', + ), + ).toMatchObject({ + toolName: 'notebook_edit', + filePath: '/project/analysis.ipynb', + }); + }); }); diff --git a/packages/core/src/core/permission-helpers.ts b/packages/core/src/core/permission-helpers.ts index c5246d623..1213bbfe9 100644 --- a/packages/core/src/core/permission-helpers.ts +++ b/packages/core/src/core/permission-helpers.ts @@ -49,11 +49,18 @@ export function buildPermissionCheckContext( : path.resolve(targetDir, toolParams['directory']) : undefined; - // Extract file path — tools use 'file_path' or 'path' (LS / grep / glob). + // Extract file path — tools use 'file_path', 'notebook_path', or + // 'path' (LS / grep / glob). let filePath = typeof toolParams['file_path'] === 'string' ? toolParams['file_path'] : undefined; + if ( + filePath === undefined && + typeof toolParams['notebook_path'] === 'string' + ) { + filePath = toolParams['notebook_path']; + } if (filePath === undefined && typeof toolParams['path'] === 'string') { // LS uses absolute paths; grep/glob may be relative to targetDir. filePath = path.isAbsolute(toolParams['path']) diff --git a/packages/core/src/extension/claude-converter.test.ts b/packages/core/src/extension/claude-converter.test.ts index 1c07118e6..a55a8bbd8 100644 --- a/packages/core/src/extension/claude-converter.test.ts +++ b/packages/core/src/extension/claude-converter.test.ts @@ -10,6 +10,7 @@ import * as path from 'node:path'; import * as os from 'node:os'; import { convertClaudeToQwenConfig, + convertClaudeAgentConfig, mergeClaudeConfigs, isClaudePluginConfig, convertClaudePluginPackage, @@ -79,6 +80,18 @@ describe('convertClaudeToQwenConfig', () => { }); }); +describe('convertClaudeAgentConfig', () => { + it('should map Claude NotebookEdit to Qwen NotebookEdit', () => { + const result = convertClaudeAgentConfig({ + name: 'notebook-agent', + description: 'Works on notebooks', + tools: ['Read', 'NotebookEdit', 'Edit'], + }); + + expect(result['tools']).toEqual(['ReadFile', 'NotebookEdit', 'Edit']); + }); +}); + describe('mergeClaudeConfigs', () => { it('should merge marketplace and plugin configs', () => { const marketplacePlugin: ClaudeMarketplacePluginConfig = { diff --git a/packages/core/src/extension/claude-converter.ts b/packages/core/src/extension/claude-converter.ts index 0f8600832..2df8b9df0 100644 --- a/packages/core/src/extension/claude-converter.ts +++ b/packages/core/src/extension/claude-converter.ts @@ -102,7 +102,7 @@ const CLAUDE_TOOLS_MAPPING: Record = { Glob: 'Glob', Grep: 'Grep', KillShell: 'None', - NotebookEdit: 'None', + NotebookEdit: 'NotebookEdit', Read: 'ReadFile', Skill: 'Skill', Task: 'Task', diff --git a/packages/core/src/permissions/permission-manager.test.ts b/packages/core/src/permissions/permission-manager.test.ts index 2ac84031e..32d529d5a 100644 --- a/packages/core/src/permissions/permission-manager.test.ts +++ b/packages/core/src/permissions/permission-manager.test.ts @@ -40,6 +40,8 @@ describe('resolveToolName', () => { expect(resolveToolName('ReadFile')).toBe('read_file'); expect(resolveToolName('ReadFileTool')).toBe('read_file'); expect(resolveToolName('EditTool')).toBe('edit'); + expect(resolveToolName('NotebookEdit')).toBe('notebook_edit'); + expect(resolveToolName('NotebookEditTool')).toBe('notebook_edit'); expect(resolveToolName('WriteFileTool')).toBe('write_file'); }); @@ -77,6 +79,7 @@ describe('getSpecifierKind', () => { it('returns "path" for file read/edit tools', async () => { expect(getSpecifierKind('read_file')).toBe('path'); expect(getSpecifierKind('edit')).toBe('path'); + expect(getSpecifierKind('notebook_edit')).toBe('path'); expect(getSpecifierKind('write_file')).toBe('path'); expect(getSpecifierKind('grep_search')).toBe('path'); expect(getSpecifierKind('glob')).toBe('path'); @@ -108,8 +111,9 @@ describe('toolMatchesRuleToolName', () => { expect(toolMatchesRuleToolName('read_file', 'list_directory')).toBe(true); }); - it('"Edit" (edit) covers write_file', async () => { + it('"Edit" (edit) covers write_file and notebook_edit', async () => { expect(toolMatchesRuleToolName('edit', 'write_file')).toBe(true); + expect(toolMatchesRuleToolName('edit', 'notebook_edit')).toBe(true); }); it('"Bash" (run_shell_command) covers monitor', async () => { @@ -704,10 +708,11 @@ describe('matchesRule', () => { }); // Meta-category matching: Edit - it('Edit rule matches edit and write_file', async () => { + it('Edit rule matches edit, write_file, and notebook_edit', async () => { const rule = parseRule('Edit'); expect(matchesRule(rule, 'edit')).toBe(true); expect(matchesRule(rule, 'write_file')).toBe(true); + expect(matchesRule(rule, 'notebook_edit')).toBe(true); expect(matchesRule(rule, 'read_file')).toBe(false); // not an edit tool }); @@ -765,6 +770,31 @@ describe('matchesRule', () => { ).toBe(false); }); + it('Edit path specifier matches notebook_edit too', async () => { + const rule = parseRule('Edit(/src/**/*.ipynb)'); + const pathCtx = { projectRoot: '/project', cwd: '/project' }; + expect( + matchesRule( + rule, + 'notebook_edit', + undefined, + '/project/src/analysis.ipynb', + undefined, + pathCtx, + ), + ).toBe(true); + expect( + matchesRule( + rule, + 'notebook_edit', + undefined, + '/project/docs/analysis.ipynb', + undefined, + pathCtx, + ), + ).toBe(false); + }); + // WebFetch domain matching it('WebFetch domain specifier', async () => { const rule = parseRule('WebFetch(domain:example.com)'); @@ -1504,6 +1534,12 @@ describe('PermissionManager', () => { expect(await pm.isToolEnabled('run_shell_command')).toBe(false); }); + it('Edit deny rule disables notebook_edit', async () => { + pm = new PermissionManager(makeConfig({ permissionsDeny: ['Edit'] })); + pm.initialize(); + expect(await pm.isToolEnabled('notebook_edit')).toBe(false); + }); + it('coreTools allowlist: listed tool is enabled', async () => { pm = new PermissionManager( makeConfig({ coreTools: ['read_file', 'Bash'] }), @@ -1519,6 +1555,14 @@ describe('PermissionManager', () => { expect(await pm.isToolEnabled('read_file')).toBe(true); expect(await pm.isToolEnabled('run_shell_command')).toBe(false); expect(await pm.isToolEnabled('edit')).toBe(false); + expect(await pm.isToolEnabled('notebook_edit')).toBe(false); + }); + + it('coreTools allowlist: NotebookEdit alias enables notebook_edit', async () => { + pm = new PermissionManager(makeConfig({ coreTools: ['NotebookEdit'] })); + pm.initialize(); + expect(await pm.isToolEnabled('notebook_edit')).toBe(true); + expect(await pm.isToolEnabled('edit')).toBe(false); }); it('coreTools with specifier: tool-level check strips specifier', async () => { @@ -1773,6 +1817,7 @@ describe('getRuleDisplayName', () => { it('maps edit tools to "Edit" meta-category', async () => { expect(getRuleDisplayName('edit')).toBe('Edit'); expect(getRuleDisplayName('write_file')).toBe('Edit'); + expect(getRuleDisplayName('notebook_edit')).toBe('Edit'); }); it('maps shell to "Bash"', async () => { @@ -1848,6 +1893,14 @@ describe('buildPermissionRules', () => { expect(rules).toEqual(['Edit(//tmp/**)']); }); + it('generates Edit rule scoped to parent directory for notebook_edit', async () => { + const rules = buildPermissionRules({ + toolName: 'notebook_edit', + filePath: '/tmp/analysis.ipynb', + }); + expect(rules).toEqual(['Edit(//tmp/**)']); + }); + it('falls back to bare display name when no filePath', async () => { const rules = buildPermissionRules({ toolName: 'read_file' }); expect(rules).toEqual(['Read']); diff --git a/packages/core/src/permissions/permission-manager.ts b/packages/core/src/permissions/permission-manager.ts index 24f7b09b8..20bd0f4d2 100644 --- a/packages/core/src/permissions/permission-manager.ts +++ b/packages/core/src/permissions/permission-manager.ts @@ -468,6 +468,7 @@ export class PermissionManager { 'read_file', 'write_file', 'edit', + 'notebook_edit', 'glob', 'grep_search', 'run_shell_command', diff --git a/packages/core/src/permissions/rule-parser.ts b/packages/core/src/permissions/rule-parser.ts index 2e6ea456a..5ebd6de52 100644 --- a/packages/core/src/permissions/rule-parser.ts +++ b/packages/core/src/permissions/rule-parser.ts @@ -52,6 +52,11 @@ export const TOOL_NAME_ALIASES: Readonly> = { Edit: 'edit', EditTool: 'edit', + // Notebook Edit tool — also matched by "Edit" meta-category rules + notebook_edit: 'notebook_edit', + NotebookEdit: 'notebook_edit', + NotebookEditTool: 'notebook_edit', + // Write File tool — also matched by "Edit" meta-category rules write_file: 'write_file', WriteFile: 'write_file', @@ -153,7 +158,7 @@ const READ_TOOLS = new Set([ * * Per Claude Code docs: "Edit rules apply to all built-in tools that edit files." */ -const EDIT_TOOLS = new Set(['edit', 'write_file']); +const EDIT_TOOLS = new Set(['edit', 'write_file', 'notebook_edit']); /** * WebFetch tools. @@ -319,6 +324,7 @@ const CANONICAL_TO_RULE_DISPLAY: Readonly> = { // Edit meta-category edit: 'Edit', write_file: 'Edit', + notebook_edit: 'Edit', // Shell run_shell_command: 'Bash', // Monitor @@ -355,7 +361,12 @@ export function getRuleDisplayName(canonicalToolName: string): string { * Directory-targeted tools (list_directory, grep_search, glob) already receive * a directory path, so they use it as-is. */ -const FILE_TARGETED_TOOLS = new Set(['read_file', 'edit', 'write_file']); +const FILE_TARGETED_TOOLS = new Set([ + 'read_file', + 'edit', + 'write_file', + 'notebook_edit', +]); /** * Build minimum-scope permission rule strings from a permission check context. diff --git a/packages/core/src/services/fileReadCache.test.ts b/packages/core/src/services/fileReadCache.test.ts index 182693351..894b0cb85 100644 --- a/packages/core/src/services/fileReadCache.test.ts +++ b/packages/core/src/services/fileReadCache.test.ts @@ -332,6 +332,17 @@ describe('FileReadCache', () => { expect(afterWrite.entry.lastReadCacheable).toBe(true); } }); + + it('can record structured writes as non-cacheable', () => { + const cache = new FileReadCache(); + const entry = cache.recordWrite('/x/notebook.ipynb', makeStats(), { + cacheable: false, + }); + + expect(entry.lastReadWasFull).toBe(true); + expect(entry.lastReadCacheable).toBe(false); + expect(entry.lastReadAt).toBe(entry.lastWriteAt); + }); }); describe('read-then-write-then-read ordering', () => { diff --git a/packages/core/src/services/fileReadCache.ts b/packages/core/src/services/fileReadCache.ts index cae6f7ff3..0dbd1d7b4 100644 --- a/packages/core/src/services/fileReadCache.ts +++ b/packages/core/src/services/fileReadCache.ts @@ -149,7 +149,8 @@ export class FileReadCache { * output was not truncated. Pass `false` for ranged reads OR * for full-request reads whose content was truncated by the * truncate-tool-output limit; both leave the model without - * sight of every current byte. + * sight of every current byte. This gates the `file_unchanged` + * fast-path and notebook-specific prior-read checks. * - `cacheable` — the produced content is plain text (vs. binary / * image / audio / video / PDF / notebook). This flag is purely * about content type, not about whether the read was complete: @@ -230,22 +231,24 @@ export class FileReadCache { * see its own write as a "stale" external change. * * Read metadata is **always** refreshed alongside the write, not - * just for brand-new entries: the model authored the entire current - * content, so for prior-read enforcement purposes it has now "seen" - * all bytes — regardless of whether the prior recordRead happened - * to be partial (`lastReadWasFull=false`) or non-cacheable - * (`lastReadCacheable=false`). Without this, a sequence such as - * `ReadFile(limit=10)` → `WriteFile` (full content) → `Edit` would - * be rejected on the Edit because `lastReadWasFull=false` from the - * earlier partial read would persist through the write. + * just for brand-new entries: the model authored the current content + * produced by the mutating tool, so for prior-read enforcement purposes + * it has now "seen" the bytes that tool wrote. Plain text writers use + * the default `cacheable: true`; structured writers such as notebook cell + * editors can set `cacheable: false` so regular Edit / WriteFile still + * reject the file as a non-text payload. */ - recordWrite(absPath: string, stats: Stats): FileReadEntry { + recordWrite( + absPath: string, + stats: Stats, + opts: { cacheable?: boolean } = {}, + ): FileReadEntry { const entry = this.upsert(absPath, stats); const now = Date.now(); entry.lastWriteAt = now; entry.lastReadAt = now; entry.lastReadWasFull = true; - entry.lastReadCacheable = true; + entry.lastReadCacheable = opts.cacheable ?? true; // The model authored the current bytes and that result is in // history, so the fast-path may serve a placeholder again. entry.readResidentInHistory = true; diff --git a/packages/core/src/tools/edit.test.ts b/packages/core/src/tools/edit.test.ts index c6a09218e..2ad797dcb 100644 --- a/packages/core/src/tools/edit.test.ts +++ b/packages/core/src/tools/edit.test.ts @@ -1216,6 +1216,7 @@ describe('EditTool', () => { expect(result.error?.message).toMatch( /binary \/ image \/ audio \/ video \/ PDF \/ notebook payload/, ); + expect(result.error?.message).toContain('notebook_edit'); expect(result.error?.message).not.toMatch(/Use the read_file tool first/); // EditTool's verb is "edit", not "overwrite" — using the // wrong one here would be confusing for in-place edits. diff --git a/packages/core/src/tools/notebook-edit.test.ts b/packages/core/src/tools/notebook-edit.test.ts new file mode 100644 index 000000000..4ab985c17 --- /dev/null +++ b/packages/core/src/tools/notebook-edit.test.ts @@ -0,0 +1,800 @@ +/** + * @license + * Copyright 2026 Qwen Team + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest'; +import fs from 'node:fs'; +import os from 'node:os'; +import path from 'node:path'; +import { execFileSync } from 'node:child_process'; +import type { Config } from '../config/config.js'; +import { ApprovalMode } from '../config/config.js'; +import { FileDiscoveryService } from '../services/fileDiscoveryService.js'; +import { FileReadCache } from '../services/fileReadCache.js'; +import { StandardFileSystemService } from '../services/fileSystemService.js'; +import { CommitAttributionService } from '../services/commitAttribution.js'; +import { createMockWorkspaceContext } from '../test-utils/mockWorkspaceContext.js'; +import { ToolErrorType } from './tool-error.js'; +import type { ToolInvocation, ToolResult } from './tools.js'; +import { applyNotebookEdit, NotebookEditTool } from './notebook-edit.js'; + +vi.mock('../telemetry/loggers.js', () => ({ + logFileOperation: vi.fn(), +})); + +describe('NotebookEditTool', () => { + let tempDir: string; + let fileReadCache: FileReadCache; + let config: Config; + let tool: NotebookEditTool; + let mockFileHistoryService: { trackEdit: ReturnType }; + const abortSignal = new AbortController().signal; + + beforeEach(() => { + CommitAttributionService.resetInstance(); + tempDir = fs.mkdtempSync(path.join(os.tmpdir(), 'notebook-edit-test-')); + fileReadCache = new FileReadCache(); + mockFileHistoryService = { trackEdit: vi.fn() }; + config = { + getTargetDir: () => tempDir, + getProjectRoot: () => tempDir, + getApprovalMode: vi.fn().mockReturnValue(ApprovalMode.DEFAULT), + setApprovalMode: vi.fn(), + getWorkspaceContext: () => createMockWorkspaceContext(tempDir), + getFileService: () => new FileDiscoveryService(tempDir), + getFileSystemService: () => new StandardFileSystemService(), + getDefaultFileEncoding: () => 'utf-8', + getFileReadCache: () => fileReadCache, + getFileHistoryService: () => mockFileHistoryService, + getFileReadCacheDisabled: () => false, + getGeminiClient: vi.fn(), + getBaseLlmClient: vi.fn(), + getIdeMode: () => false, + getApiKey: () => 'test-api-key', + getModel: () => 'test-model', + getSandbox: () => false, + getDebugMode: () => false, + getQuestion: () => undefined, + getFullContext: () => false, + getToolDiscoveryCommand: () => undefined, + getToolCallCommand: () => undefined, + getMcpServerCommand: () => undefined, + getMcpServers: () => undefined, + getUserAgent: () => 'test-agent', + getUserMemory: () => '', + setUserMemory: vi.fn(), + getGeminiMdFileCount: () => 0, + setGeminiMdFileCount: vi.fn(), + getToolRegistry: () => ({}) as never, + } as unknown as Config; + tool = new NotebookEditTool(config); + }); + + afterEach(() => { + CommitAttributionService.resetInstance(); + fs.rmSync(tempDir, { recursive: true, force: true }); + }); + + function writeNotebook(name: string, notebook: Record) { + const filePath = path.join(tempDir, name); + fs.writeFileSync(filePath, JSON.stringify(notebook, null, 1), 'utf-8'); + return filePath; + } + + function seedNotebookRead(filePath: string) { + fileReadCache.recordRead(filePath, fs.statSync(filePath), { + full: true, + cacheable: false, + }); + } + + function buildInvocation(params: Parameters[0]) { + return tool.build(params) as ToolInvocation< + Parameters[0], + ToolResult + >; + } + + it('replaces a code cell by real ID and clears stale outputs', async () => { + const filePath = writeNotebook('analysis.ipynb', { + nbformat: 4, + nbformat_minor: 5, + cells: [ + { + cell_type: 'code', + id: 'load-data', + source: ['x = 1\n'], + execution_count: 7, + outputs: [{ output_type: 'stream', text: ['old\n'] }], + metadata: {}, + }, + ], + metadata: { language_info: { name: 'python' } }, + }); + seedNotebookRead(filePath); + + const result = await buildInvocation({ + notebook_path: filePath, + cell_id: 'load-data', + new_source: 'x = 2\nprint(x)', + }).execute(abortSignal); + + expect(result.error).toBeUndefined(); + const updated = JSON.parse(fs.readFileSync(filePath, 'utf-8')); + expect(updated.cells[0].source).toEqual(['x = 2\n', 'print(x)']); + expect(updated.cells[0].execution_count).toBeNull(); + expect(updated.cells[0].outputs).toEqual([]); + expect(result.llmContent).toContain('replace cell load-data'); + + const cacheState = fileReadCache.check(fs.statSync(filePath)); + expect(cacheState.state).toBe('fresh'); + if (cacheState.state === 'fresh') { + expect(cacheState.entry.lastReadWasFull).toBe(true); + expect(cacheState.entry.lastReadCacheable).toBe(false); + } + }); + + it('replaces a code cell in a UTF-8 BOM notebook and preserves the BOM', async () => { + const filePath = path.join(tempDir, 'bom-replace.ipynb'); + fs.writeFileSync( + filePath, + `\ufeff${JSON.stringify( + { + nbformat: 4, + nbformat_minor: 5, + cells: [ + { + cell_type: 'code', + id: 'load-data', + source: ['x = 1\n'], + execution_count: 7, + outputs: [{ output_type: 'stream', text: ['old\n'] }], + metadata: {}, + }, + ], + metadata: { language_info: { name: 'python' } }, + }, + null, + 1, + )}`, + 'utf-8', + ); + seedNotebookRead(filePath); + + const result = await buildInvocation({ + notebook_path: filePath, + cell_id: 'load-data', + new_source: 'x = 2\nprint(x)', + }).execute(abortSignal); + + expect(result.error).toBeUndefined(); + const updatedBuffer = fs.readFileSync(filePath); + expect([...updatedBuffer.subarray(0, 3)]).toEqual([0xef, 0xbb, 0xbf]); + const updated = JSON.parse(updatedBuffer.toString('utf-8').slice(1)); + expect(updated.cells[0].source).toEqual(['x = 2\n', 'print(x)']); + expect(updated.cells[0].execution_count).toBeNull(); + expect(updated.cells[0].outputs).toEqual([]); + }); + + it('replaces by cell-N fallback and converts code to markdown cleanly', async () => { + const filePath = writeNotebook('convert.ipynb', { + nbformat: 4, + nbformat_minor: 5, + cells: [ + { + cell_type: 'code', + source: 'print("old")', + execution_count: 1, + outputs: [{ output_type: 'stream', text: 'old\n' }], + metadata: {}, + }, + ], + metadata: {}, + }); + seedNotebookRead(filePath); + + const result = await buildInvocation({ + notebook_path: filePath, + cell_id: 'cell-0', + cell_type: 'markdown', + new_source: '# Notes', + }).execute(abortSignal); + + expect(result.error).toBeUndefined(); + const updated = JSON.parse(fs.readFileSync(filePath, 'utf-8')); + expect(updated.cells[0].cell_type).toBe('markdown'); + expect(updated.cells[0].source).toBe('# Notes'); + expect(updated.cells[0]).not.toHaveProperty('outputs'); + expect(updated.cells[0]).not.toHaveProperty('execution_count'); + }); + + it('converts markdown to code with code-only fields', async () => { + const filePath = writeNotebook('convert-to-code.ipynb', { + nbformat: 4, + nbformat_minor: 5, + cells: [ + { + cell_type: 'markdown', + id: 'intro', + source: ['# Intro'], + metadata: {}, + }, + ], + metadata: {}, + }); + seedNotebookRead(filePath); + + const result = await buildInvocation({ + notebook_path: filePath, + cell_id: 'intro', + cell_type: 'code', + new_source: 'print("hi")', + }).execute(abortSignal); + + expect(result.error).toBeUndefined(); + const updated = JSON.parse(fs.readFileSync(filePath, 'utf-8')); + expect(updated.cells[0].cell_type).toBe('code'); + expect(updated.cells[0].source).toEqual(['print("hi")']); + expect(updated.cells[0].execution_count).toBeNull(); + expect(updated.cells[0].outputs).toEqual([]); + }); + + it('inserts after a target cell and generates an nbformat 4.5 cell ID', async () => { + const raw = JSON.stringify({ + nbformat: 4, + nbformat_minor: 5, + cells: [ + { cell_type: 'markdown', id: 'cell-1', source: ['# A'], metadata: {} }, + { cell_type: 'code', id: 'cell-2', source: ['a = 1'], metadata: {} }, + ], + metadata: {}, + }); + + const result = applyNotebookEdit(raw, { + notebook_path: '/tmp/insert.ipynb', + edit_mode: 'insert', + cell_id: 'cell-1', + cell_type: 'markdown', + new_source: '## Inserted', + }); + + const updated = JSON.parse(result.updatedContent); + expect(updated.cells).toHaveLength(3); + expect(updated.cells[1].cell_type).toBe('markdown'); + expect(updated.cells[1].source).toEqual(['## Inserted']); + expect(updated.cells[1].id).toBe('qwen-cell-1'); + expect(result.editedCellId).toBe('qwen-cell-1'); + }); + + it('preserves adjacent source style for inserted cells in mixed-format notebooks', async () => { + const raw = JSON.stringify({ + nbformat: 4, + nbformat_minor: 5, + cells: [ + { cell_type: 'markdown', id: 'intro', source: '# Intro', metadata: {} }, + { + cell_type: 'code', + id: 'code', + source: ['value = 1\n'], + metadata: {}, + }, + ], + metadata: {}, + }); + + const result = applyNotebookEdit(raw, { + notebook_path: '/tmp/insert.ipynb', + edit_mode: 'insert', + cell_id: 'intro', + cell_type: 'markdown', + new_source: '## Inserted', + }); + + const updated = JSON.parse(result.updatedContent); + expect(updated.cells[1].source).toBe('## Inserted'); + }); + + it('preserves notebook JSON indentation and trailing newline style on edit', () => { + const raw = JSON.stringify( + { + nbformat: 4, + nbformat_minor: 5, + cells: [ + { + cell_type: 'markdown', + id: 'intro', + source: '# Intro', + metadata: {}, + }, + ], + metadata: {}, + }, + null, + 2, + ); + + const result = applyNotebookEdit(raw, { + notebook_path: '/tmp/format.ipynb', + cell_id: 'intro', + new_source: '# Updated', + }); + + expect(result.updatedContent).toContain('\n "cells"'); + expect(result.updatedContent.endsWith('\n')).toBe(false); + }); + + it('rejects ambiguous fallback-like cell IDs', async () => { + const filePath = writeNotebook('ambiguous.ipynb', { + nbformat: 4, + nbformat_minor: 5, + cells: [ + { + cell_type: 'markdown', + id: 'cell-1', + source: ['real id'], + metadata: {}, + }, + { + cell_type: 'markdown', + source: ['fallback id'], + metadata: {}, + }, + ], + metadata: {}, + }); + seedNotebookRead(filePath); + + const result = await buildInvocation({ + notebook_path: filePath, + cell_id: 'cell-1', + new_source: 'updated', + }).execute(abortSignal); + + expect(result.error?.type).toBe(ToolErrorType.INVALID_TOOL_PARAMS); + expect(result.llmContent).toContain('ambiguous'); + }); + + it('inserts at the beginning when no cell_id is provided', async () => { + const filePath = writeNotebook('insert-start.ipynb', { + nbformat: 4, + nbformat_minor: 4, + cells: [{ cell_type: 'code', source: ['x = 1'], metadata: {} }], + metadata: {}, + }); + seedNotebookRead(filePath); + + const result = await buildInvocation({ + notebook_path: filePath, + edit_mode: 'insert', + cell_type: 'code', + new_source: 'print("first")', + }).execute(abortSignal); + + expect(result.error).toBeUndefined(); + const updated = JSON.parse(fs.readFileSync(filePath, 'utf-8')); + expect(updated.cells[0].source).toEqual(['print("first")']); + expect(updated.cells[0]).not.toHaveProperty('id'); + }); + + it('deletes a cell without requiring new_source', async () => { + const filePath = writeNotebook('delete.ipynb', { + nbformat: 4, + nbformat_minor: 5, + cells: [ + { cell_type: 'markdown', id: 'keep', source: ['keep'], metadata: {} }, + { cell_type: 'markdown', id: 'drop', source: ['drop'], metadata: {} }, + ], + metadata: {}, + }); + seedNotebookRead(filePath); + + const result = await buildInvocation({ + notebook_path: filePath, + edit_mode: 'delete', + cell_id: 'drop', + }).execute(abortSignal); + + expect(result.error).toBeUndefined(); + const updated = JSON.parse(fs.readFileSync(filePath, 'utf-8')); + expect(updated.cells.map((cell: { id: string }) => cell.id)).toEqual([ + 'keep', + ]); + }); + + it('requires a fresh read after structural edits when fallback IDs can shift', async () => { + const filePath = writeNotebook('fallback-shift.ipynb', { + nbformat: 4, + nbformat_minor: 5, + cells: [ + { cell_type: 'markdown', source: ['A'], metadata: {} }, + { cell_type: 'markdown', source: ['B'], metadata: {} }, + ], + metadata: {}, + }); + seedNotebookRead(filePath); + + const result = await buildInvocation({ + notebook_path: filePath, + edit_mode: 'insert', + cell_type: 'markdown', + new_source: 'inserted', + }).execute(abortSignal); + + expect(result.error).toBeUndefined(); + expect(fileReadCache.check(fs.statSync(filePath)).state).toBe('unknown'); + }); + + it('preserves fresh read state after structural edits when all IDs are stable', async () => { + const filePath = writeNotebook('stable-ids.ipynb', { + nbformat: 4, + nbformat_minor: 5, + cells: [ + { cell_type: 'markdown', id: 'a', source: ['A'], metadata: {} }, + { cell_type: 'markdown', id: 'b', source: ['B'], metadata: {} }, + ], + metadata: {}, + }); + seedNotebookRead(filePath); + + const result = await buildInvocation({ + notebook_path: filePath, + edit_mode: 'insert', + cell_id: 'a', + cell_type: 'markdown', + new_source: 'inserted', + }).execute(abortSignal); + + expect(result.error).toBeUndefined(); + const cacheState = fileReadCache.check(fs.statSync(filePath)); + expect(cacheState.state).toBe('fresh'); + if (cacheState.state === 'fresh') { + expect(cacheState.entry.lastReadWasFull).toBe(true); + } + }); + + it('requires a fresh full notebook read before editing', async () => { + const filePath = writeNotebook('unread.ipynb', { + cells: [{ cell_type: 'code', id: 'a', source: ['x = 1'], metadata: {} }], + metadata: {}, + }); + + const result = await buildInvocation({ + notebook_path: filePath, + cell_id: 'a', + new_source: 'x = 2', + }).execute(abortSignal); + + expect(result.error?.type).toBe(ToolErrorType.EDIT_REQUIRES_PRIOR_READ); + expect(result.llmContent).toContain('has not been fully read'); + }); + + it('rejects edits after a truncated notebook read', async () => { + const filePath = writeNotebook('truncated-read.ipynb', { + cells: [ + { cell_type: 'code', id: 'visible', source: ['x = 1'], metadata: {} }, + { cell_type: 'code', id: 'tail', source: ['x = 2'], metadata: {} }, + ], + metadata: {}, + }); + fileReadCache.recordRead(filePath, fs.statSync(filePath), { + full: false, + cacheable: false, + }); + + const result = await buildInvocation({ + notebook_path: filePath, + cell_id: 'tail', + new_source: 'x = 3', + }).execute(abortSignal); + + expect(result.error?.type).toBe(ToolErrorType.EDIT_REQUIRES_PRIOR_READ); + expect(result.llmContent).toContain('too large for cell-level editing'); + expect(result.llmContent).not.toContain('without offset or limit'); + }); + + it.skipIf(process.platform === 'win32')( + 'rejects non-regular notebook paths with a dedicated error type', + async () => { + const fifoPath = path.join(tempDir, 'notebook-fifo.ipynb'); + execFileSync('mkfifo', [fifoPath]); + + const result = await buildInvocation({ + notebook_path: fifoPath, + cell_id: 'a', + new_source: 'x = 2', + }).execute(abortSignal); + + expect(result.error?.type).toBe(ToolErrorType.TARGET_NOT_REGULAR_FILE); + expect(result.llmContent).toContain('not a regular file'); + }, + ); + + it('rejects stale notebook edits after an external change', async () => { + const filePath = writeNotebook('stale.ipynb', { + cells: [{ cell_type: 'code', id: 'a', source: ['x = 1'], metadata: {} }], + metadata: {}, + }); + seedNotebookRead(filePath); + fs.writeFileSync( + filePath, + JSON.stringify({ + cells: [ + { cell_type: 'code', id: 'a', source: ['x = 100'], metadata: {} }, + ], + metadata: {}, + }), + 'utf-8', + ); + + const result = await buildInvocation({ + notebook_path: filePath, + cell_id: 'a', + new_source: 'x = 2', + }).execute(abortSignal); + + expect(result.error?.type).toBe(ToolErrorType.FILE_CHANGED_SINCE_READ); + }); + + it('returns structured errors for missing cells and invalid JSON', async () => { + const missingCellPath = writeNotebook('missing-cell.ipynb', { + cells: [{ cell_type: 'code', id: 'a', source: ['x = 1'], metadata: {} }], + metadata: {}, + }); + seedNotebookRead(missingCellPath); + + const missingCellResult = await buildInvocation({ + notebook_path: missingCellPath, + cell_id: 'missing', + new_source: 'x = 2', + }).execute(abortSignal); + + expect(missingCellResult.error?.type).toBe( + ToolErrorType.NOTEBOOK_CELL_NOT_FOUND, + ); + + const invalidPath = path.join(tempDir, 'bad.ipynb'); + fs.writeFileSync(invalidPath, 'not json', 'utf-8'); + seedNotebookRead(invalidPath); + + const invalidResult = await buildInvocation({ + notebook_path: invalidPath, + edit_mode: 'insert', + new_source: 'x = 1', + }).execute(abortSignal); + + expect(invalidResult.error?.type).toBe(ToolErrorType.NOTEBOOK_INVALID_JSON); + }); + + it('keeps invalid original notebook errors structured for user-modified content', async () => { + const invalidPath = writeNotebook('bad-original.ipynb', { + nbformat: 4, + nbformat_minor: 5, + cells: [{ cell_type: 'code', id: 'a', source: ['x = 1'], metadata: {} }], + metadata: {}, + }); + seedNotebookRead(invalidPath); + const originalParams = { + notebook_path: invalidPath, + cell_id: 'a', + new_source: 'x = 2', + }; + const modifyContext = tool.getModifyContext(abortSignal); + const currentContent = + await modifyContext.getCurrentContent(originalParams); + const proposedContent = + await modifyContext.getProposedContent(originalParams); + const updatedParams = modifyContext.createUpdatedParams( + currentContent, + proposedContent, + originalParams, + ); + fs.writeFileSync(invalidPath, 'not json', 'utf-8'); + seedNotebookRead(invalidPath); + + const result = await buildInvocation( + structuredClone(updatedParams), + ).execute(abortSignal); + + expect(result.error?.type).toBe(ToolErrorType.NOTEBOOK_INVALID_JSON); + }); + + it('rejects direct attempts to set internal modified notebook content params', () => { + const filePath = writeNotebook('injected-modified-content.ipynb', { + nbformat: 4, + nbformat_minor: 5, + cells: [{ cell_type: 'code', id: 'a', source: ['x = 1'], metadata: {} }], + metadata: {}, + }); + + expect(() => + tool.build({ + notebook_path: filePath, + cell_id: 'a', + new_source: 'x = 2', + modified_notebook_content: JSON.stringify({ + cells: [], + metadata: {}, + }), + } as Parameters[0]), + ).toThrow(/additional properties|modified_notebook_content/i); + }); + + it('rejects qwenignored notebooks during validation', () => { + fs.writeFileSync(path.join(tempDir, '.qwenignore'), '*.ipynb\n', 'utf-8'); + const filePath = writeNotebook('ignored.ipynb', { + cells: [], + metadata: {}, + }); + + expect(() => + tool.build({ + notebook_path: filePath, + edit_mode: 'insert', + new_source: 'x = 1', + }), + ).toThrow(/ignored by \.qwenignore/); + }); + + it('returns a notebook diff for confirmation', async () => { + const filePath = writeNotebook('confirm.ipynb', { + cells: [{ cell_type: 'code', id: 'a', source: ['x = 1'], metadata: {} }], + metadata: {}, + }); + seedNotebookRead(filePath); + + const details = await buildInvocation({ + notebook_path: filePath, + cell_id: 'a', + new_source: 'x = 2', + }).getConfirmationDetails(abortSignal); + + const editDetails = details as Extract; + expect(editDetails.fileDiff).toContain('- "x = 1"'); + expect(editDetails.fileDiff).toContain('+ "x = 2"'); + expect((editDetails as { originalContent: string }).originalContent).toBe( + fs.readFileSync(filePath, 'utf-8'), + ); + }); + + it('applies IDE or inline modified full-notebook content instead of the original cell proposal', async () => { + const filePath = writeNotebook('modified-content.ipynb', { + nbformat: 4, + nbformat_minor: 5, + cells: [{ cell_type: 'code', id: 'a', source: ['x = 1'], metadata: {} }], + metadata: {}, + }); + seedNotebookRead(filePath); + + const originalParams = { + notebook_path: filePath, + cell_id: 'a', + new_source: 'x = 2', + }; + const modifyContext = tool.getModifyContext(abortSignal); + const currentContent = + await modifyContext.getCurrentContent(originalParams); + const proposedContent = + await modifyContext.getProposedContent(originalParams); + const modifiedContent = proposedContent.replace('x = 2', 'x = 99'); + const updatedParams = modifyContext.createUpdatedParams( + currentContent, + modifiedContent, + originalParams, + ); + + const result = await buildInvocation( + structuredClone(updatedParams), + ).execute(abortSignal); + + expect(result.error).toBeUndefined(); + const updated = JSON.parse(fs.readFileSync(filePath, 'utf-8')); + expect(updated.cells[0].source).toEqual(['x = 99']); + expect(result.llmContent).toContain('modified by the user'); + expect( + CommitAttributionService.getInstance().getFileAttribution(filePath), + ).toBeUndefined(); + }); + + it('uses one current-content snapshot for notebook modify previews', async () => { + const filePath = writeNotebook('modify-snapshot.ipynb', { + nbformat: 4, + nbformat_minor: 5, + cells: [{ cell_type: 'code', id: 'a', source: ['x = 1'], metadata: {} }], + metadata: {}, + }); + + const params = { + notebook_path: filePath, + cell_id: 'a', + new_source: 'x = 2', + }; + const modifyContext = tool.getModifyContext(abortSignal); + const currentContent = await modifyContext.getCurrentContent(params); + fs.writeFileSync( + filePath, + JSON.stringify( + { + nbformat: 4, + nbformat_minor: 5, + cells: [ + { cell_type: 'code', id: 'a', source: ['x = 999'], metadata: {} }, + ], + metadata: {}, + }, + null, + 1, + ), + 'utf-8', + ); + + const proposedContent = await modifyContext.getProposedContent(params); + + expect(currentContent).toContain('x = 1'); + expect(proposedContent).toContain('x = 2'); + expect(proposedContent).not.toContain('x = 999'); + }); + + it('records AI-originated notebook writes for commit attribution', async () => { + const filePath = writeNotebook('attribution.ipynb', { + nbformat: 4, + nbformat_minor: 5, + cells: [{ cell_type: 'code', id: 'a', source: ['x = 1'], metadata: {} }], + metadata: {}, + }); + seedNotebookRead(filePath); + + const result = await buildInvocation({ + notebook_path: filePath, + cell_id: 'a', + new_source: 'x = 2', + }).execute(abortSignal); + + expect(result.error).toBeUndefined(); + const attribution = + CommitAttributionService.getInstance().getFileAttribution(filePath); + expect(attribution).toBeDefined(); + expect(attribution!.aiContribution).toBeGreaterThan(0); + expect(mockFileHistoryService.trackEdit).toHaveBeenCalledWith(filePath); + }); + + it('tracks file history before the final freshness check', async () => { + const filePath = writeNotebook('history-before-check.ipynb', { + nbformat: 4, + nbformat_minor: 5, + cells: [{ cell_type: 'code', id: 'a', source: ['x = 1'], metadata: {} }], + metadata: {}, + }); + seedNotebookRead(filePath); + mockFileHistoryService.trackEdit.mockImplementation(async () => { + fs.writeFileSync( + filePath, + JSON.stringify( + { + nbformat: 4, + nbformat_minor: 5, + cells: [ + { cell_type: 'code', id: 'a', source: ['x = 100'], metadata: {} }, + ], + metadata: {}, + }, + null, + 1, + ), + 'utf-8', + ); + }); + + const result = await buildInvocation({ + notebook_path: filePath, + cell_id: 'a', + new_source: 'x = 2', + }).execute(abortSignal); + + expect(mockFileHistoryService.trackEdit).toHaveBeenCalledWith(filePath); + expect(result.error?.type).toBe(ToolErrorType.FILE_CHANGED_SINCE_READ); + const updated = JSON.parse(fs.readFileSync(filePath, 'utf-8')); + expect(updated.cells[0].source).toEqual(['x = 100']); + }); +}); diff --git a/packages/core/src/tools/notebook-edit.ts b/packages/core/src/tools/notebook-edit.ts new file mode 100644 index 000000000..c47f472f7 --- /dev/null +++ b/packages/core/src/tools/notebook-edit.ts @@ -0,0 +1,931 @@ +/** + * @license + * Copyright 2026 Qwen Team + * SPDX-License-Identifier: Apache-2.0 + */ + +import * as fs from 'node:fs'; +import * as path from 'node:path'; +import * as Diff from 'diff'; +import type { Config } from '../config/config.js'; +import { ApprovalMode } from '../config/config.js'; +import { detectLineEnding } from '../services/fileSystemService.js'; +import type { LineEnding } from '../services/fileSystemService.js'; +import type { + ToolCallConfirmationDetails, + ToolEditConfirmationDetails, + ToolInvocation, + ToolLocation, + ToolResult, +} from './tools.js'; +import { + BaseDeclarativeTool, + BaseToolInvocation, + Kind, + ToolConfirmationOutcome, +} from './tools.js'; +import type { PermissionDecision } from '../permissions/types.js'; +import { DEFAULT_DIFF_OPTIONS, getDiffStat } from './diffOptions.js'; +import { FileOperation } from '../telemetry/metrics.js'; +import { FileOperationEvent } from '../telemetry/types.js'; +import { logFileOperation } from '../telemetry/loggers.js'; +import { getSpecificMimeType } from '../utils/fileUtils.js'; +import { makeRelative, shortenPath, unescapePath } from '../utils/paths.js'; +import { + findCellIndex, + getCellDisplayId, + getNotebookLanguage, + hasStableCellIds, + inferInsertedCellSourceArrayStyle, + inferNotebookJsonFormat, + isAmbiguousCellId, + makeCellId, + normalizeEditedCell, + normalizeSource, + parseNotebook, + serializeNotebook, + toNotebookSource, + type EditableNotebookCellType, + type NotebookCell, + type NotebookCellType, +} from '../utils/notebook.js'; +import { ToolDisplayNames, ToolNames } from './tool-names.js'; +import { ToolErrorType } from './tool-error.js'; +import { StructuredToolError } from './priorReadEnforcement.js'; +import type { + ModifiableDeclarativeTool, + ModifyContext, +} from './modifiable-tool.js'; +import { CommitAttributionService } from '../services/commitAttribution.js'; +import { createDebugLogger } from '../utils/debugLogger.js'; + +const debugLogger = createDebugLogger('NOTEBOOK_EDIT'); + +export type NotebookEditMode = 'replace' | 'insert' | 'delete'; + +export interface NotebookEditToolParams { + notebook_path: string; + cell_id?: string; + new_source?: string; + cell_type?: EditableNotebookCellType; + edit_mode?: NotebookEditMode; +} + +interface NotebookEditModifyMetadata { + modifiedByUser?: boolean; + aiProposedContent?: string; + modifiedNotebookContent?: string; +} + +interface NotebookEditResult { + updatedContent: string; + editedCellId: string; + editedCellType?: NotebookCellType; + language: string; + mode: NotebookEditMode; + requiresReadAfterWrite: boolean; +} + +interface PreparedNotebookEdit extends NotebookEditResult { + originalContent: string; + bom: boolean; + encoding: string | undefined; + lineEnding: LineEnding; +} + +type NotebookPriorReadDecision = + | { ok: true } + | { + ok: false; + type: ToolErrorType; + rawMessage: string; + displayMessage: string; + }; + +function rejectNotebookPriorRead( + notebookPath: string, + reason: string, + decision: Exclude, +): NotebookPriorReadDecision { + debugLogger.debug('prior-read-rejected', { + path: notebookPath, + reason, + type: decision.type, + displayMessage: decision.displayMessage, + }); + return decision; +} + +class NotebookEditError extends Error { + constructor( + message: string, + readonly type: ToolErrorType, + ) { + super(message); + this.name = 'NotebookEditError'; + } +} + +function requireNotebookSource( + source: string | undefined, + mode: NotebookEditMode, +): string { + if (mode === 'delete') { + return ''; + } + if (typeof source !== 'string') { + throw new NotebookEditError( + `new_source is required when edit_mode is "${mode}".`, + ToolErrorType.INVALID_TOOL_PARAMS, + ); + } + return source; +} + +function displayCellId(cell: NotebookCell | undefined, index: number): string { + return cell ? getCellDisplayId(cell, index) : `cell-${index}`; +} + +function resolveTargetIndex( + notebook: ReturnType, + cellId: string | undefined, + mode: NotebookEditMode, +): number { + if (!cellId) { + if (mode === 'insert') { + return -1; + } + throw new NotebookEditError( + 'cell_id is required for replace and delete operations.', + ToolErrorType.INVALID_TOOL_PARAMS, + ); + } + + if (isAmbiguousCellId(notebook, cellId)) { + throw new NotebookEditError( + `Cell ID "${cellId}" is ambiguous in the rendered notebook. Re-read the notebook and target a stable real cell ID before editing.`, + ToolErrorType.INVALID_TOOL_PARAMS, + ); + } + + const index = findCellIndex(notebook, cellId); + if (index === -1) { + throw new NotebookEditError( + `Cell with ID "${cellId}" not found in notebook.`, + ToolErrorType.NOTEBOOK_CELL_NOT_FOUND, + ); + } + return index; +} + +function createNotebookCell( + notebook: ReturnType, + cellType: EditableNotebookCellType, + source: string, + preferSourceArray: boolean, +): NotebookCell { + const cell: NotebookCell = { + cell_type: cellType, + metadata: {}, + source: toNotebookSource(source, preferSourceArray), + }; + const id = makeCellId(notebook); + if (id) { + cell.id = id; + } + normalizeEditedCell(cell, cellType); + return cell; +} + +export function applyNotebookEdit( + rawContent: string, + params: NotebookEditToolParams, +): NotebookEditResult { + let notebook: ReturnType; + try { + notebook = parseNotebook(rawContent); + } catch (error) { + throw new NotebookEditError( + error instanceof Error ? error.message : String(error), + ToolErrorType.NOTEBOOK_INVALID_JSON, + ); + } + + const mode = params.edit_mode ?? 'replace'; + const source = requireNotebookSource(params.new_source, mode); + const originalHasStableCellIds = hasStableCellIds(notebook); + const targetIndex = resolveTargetIndex(notebook, params.cell_id, mode); + const language = getNotebookLanguage(notebook); + const jsonFormat = inferNotebookJsonFormat(rawContent); + const buildResult = ( + updatedNotebook: ReturnType, + result: Omit< + NotebookEditResult, + 'updatedContent' | 'requiresReadAfterWrite' + >, + ): NotebookEditResult => { + const structuralEdit = mode === 'insert' || mode === 'delete'; + return { + ...result, + updatedContent: serializeNotebook(updatedNotebook, jsonFormat), + requiresReadAfterWrite: + structuralEdit && + !(originalHasStableCellIds && hasStableCellIds(updatedNotebook)), + }; + }; + + switch (mode) { + case 'insert': { + const cellType = params.cell_type ?? 'code'; + const insertAt = targetIndex === -1 ? 0 : targetIndex + 1; + const newCell = createNotebookCell( + notebook, + cellType, + source, + inferInsertedCellSourceArrayStyle(notebook, insertAt), + ); + notebook.cells.splice(insertAt, 0, newCell); + return buildResult(notebook, { + editedCellId: displayCellId(newCell, insertAt), + editedCellType: cellType, + language, + mode, + }); + } + + case 'delete': { + const [removed] = notebook.cells.splice(targetIndex, 1); + return buildResult(notebook, { + editedCellId: displayCellId(removed, targetIndex), + editedCellType: removed?.cell_type, + language, + mode, + }); + } + + case 'replace': { + const target = notebook.cells[targetIndex]; + if (!target) { + throw new NotebookEditError( + `Cell index ${targetIndex} is out of range.`, + ToolErrorType.NOTEBOOK_CELL_NOT_FOUND, + ); + } + + const finalType = params.cell_type ?? target.cell_type; + target.source = toNotebookSource(source, Array.isArray(target.source)); + normalizeEditedCell(target, finalType); + return buildResult(notebook, { + editedCellId: displayCellId(target, targetIndex), + editedCellType: finalType, + language, + mode, + }); + } + + default: + throw new NotebookEditError( + `Unsupported notebook edit mode: ${mode}`, + ToolErrorType.INVALID_TOOL_PARAMS, + ); + } +} + +function prepareModifiedNotebookContent( + originalContent: string, + modifiedContent: string, + params: NotebookEditToolParams, +): NotebookEditResult { + let notebook: ReturnType; + let originalNotebook: ReturnType; + try { + notebook = parseNotebook(modifiedContent); + originalNotebook = parseNotebook(originalContent); + } catch (error) { + throw new NotebookEditError( + error instanceof Error ? error.message : String(error), + ToolErrorType.NOTEBOOK_INVALID_JSON, + ); + } + + return { + updatedContent: modifiedContent, + editedCellId: params.cell_id ?? 'user-modified-notebook', + editedCellType: undefined, + language: getNotebookLanguage(notebook), + mode: params.edit_mode ?? 'replace', + requiresReadAfterWrite: + !hasStableCellIds(originalNotebook) || !hasStableCellIds(notebook), + }; +} + +async function checkPriorNotebookRead( + config: Config, + notebookPath: string, + options: { expectExisting?: boolean } = {}, +): Promise { + if (config.getFileReadCacheDisabled()) { + return { ok: true }; + } + + let stats: fs.Stats; + try { + stats = await fs.promises.stat(notebookPath); + } catch (error) { + const code = (error as NodeJS.ErrnoException | undefined)?.code; + if (code === 'ENOENT') { + if (options.expectExisting) { + return rejectNotebookPriorRead(notebookPath, 'missing-after-read', { + ok: false, + type: ToolErrorType.FILE_CHANGED_SINCE_READ, + rawMessage: `Notebook ${notebookPath} disappeared after it was read. Re-read it with the ${ToolNames.READ_FILE} tool before editing it.`, + displayMessage: `notebook disappeared after last read; re-run ${ToolNames.READ_FILE} first.`, + }); + } + return { ok: true }; + } + + return rejectNotebookPriorRead(notebookPath, 'stat-failed', { + ok: false, + type: ToolErrorType.PRIOR_READ_VERIFICATION_FAILED, + rawMessage: `Could not stat ${notebookPath} to verify prior notebook read (${code ?? 'unknown error'}). Re-read it with the ${ToolNames.READ_FILE} tool before editing it.`, + displayMessage: `cannot verify prior read of ${notebookPath}; re-run ${ToolNames.READ_FILE} before editing this notebook.`, + }); + } + + if (stats.isDirectory()) { + return rejectNotebookPriorRead(notebookPath, 'target-is-directory', { + ok: false, + type: ToolErrorType.TARGET_IS_DIRECTORY, + rawMessage: `${notebookPath} is a directory. The NotebookEdit tool only operates on .ipynb files.`, + displayMessage: 'path is a directory; cannot edit as a notebook.', + }); + } + + if (!stats.isFile()) { + return rejectNotebookPriorRead(notebookPath, 'target-not-regular-file', { + ok: false, + type: ToolErrorType.TARGET_NOT_REGULAR_FILE, + rawMessage: `${notebookPath} is not a regular file. The NotebookEdit tool only operates on .ipynb files.`, + displayMessage: 'special file; cannot edit as a notebook.', + }); + } + + const status = config.getFileReadCache().check(stats); + if ( + status.state === 'fresh' && + status.entry.lastReadAt !== undefined && + status.entry.lastReadWasFull + ) { + return { ok: true }; + } + + if ( + status.state === 'fresh' && + status.entry.lastReadAt !== undefined && + !status.entry.lastReadWasFull + ) { + return rejectNotebookPriorRead(notebookPath, 'truncated-cache-entry', { + ok: false, + type: ToolErrorType.EDIT_REQUIRES_PRIOR_READ, + rawMessage: `Notebook ${notebookPath} is too large for cell-level editing because its rendered output was truncated when read. Reduce the notebook output size or split the notebook before editing cells.`, + displayMessage: 'notebook too large for cell-level editing.', + }); + } + + if (status.state === 'stale') { + return rejectNotebookPriorRead(notebookPath, 'stale-cache-entry', { + ok: false, + type: ToolErrorType.FILE_CHANGED_SINCE_READ, + rawMessage: `Notebook ${notebookPath} has been modified since you last read it. Re-read it with the ${ToolNames.READ_FILE} tool before editing it.`, + displayMessage: `notebook changed since last read; re-run ${ToolNames.READ_FILE} first.`, + }); + } + + return rejectNotebookPriorRead(notebookPath, `cache-${status.state}`, { + ok: false, + type: ToolErrorType.EDIT_REQUIRES_PRIOR_READ, + rawMessage: `Notebook ${notebookPath} has not been fully read in this session. Use the ${ToolNames.READ_FILE} tool first, without offset or limit, before editing cells.`, + displayMessage: `${ToolNames.READ_FILE} required before editing this notebook.`, + }); +} + +class NotebookEditInvocation extends BaseToolInvocation< + NotebookEditToolParams, + ToolResult +> { + constructor( + private readonly config: Config, + params: NotebookEditToolParams, + private readonly modifyMetadata?: NotebookEditModifyMetadata, + ) { + super(params); + } + + override toolLocations(): ToolLocation[] { + return [{ path: this.params.notebook_path }]; + } + + override getDescription(): string { + const relativePath = makeRelative( + this.params.notebook_path, + this.config.getTargetDir(), + ); + const mode = this.params.edit_mode ?? 'replace'; + const cell = this.params.cell_id ?? 'beginning'; + return `${mode} notebook cell ${cell} in ${shortenPath(relativePath)}`; + } + + override async getDefaultPermission(): Promise { + return 'ask'; + } + + override async getConfirmationDetails( + abortSignal: AbortSignal, + ): Promise { + const prepared = await this.prepareEdit(abortSignal); + const fileName = path.basename(this.params.notebook_path); + const fileDiff = Diff.createPatch( + fileName, + prepared.originalContent, + prepared.updatedContent, + 'Current', + 'Proposed', + DEFAULT_DIFF_OPTIONS, + ); + + const confirmationDetails: ToolEditConfirmationDetails = { + type: 'edit', + title: `Confirm Notebook Edit: ${shortenPath(makeRelative(this.params.notebook_path, this.config.getTargetDir()))}`, + fileName, + filePath: this.params.notebook_path, + fileDiff, + originalContent: prepared.originalContent, + newContent: prepared.updatedContent, + onConfirm: async (outcome: ToolConfirmationOutcome) => { + if (outcome === ToolConfirmationOutcome.ProceedAlways) { + this.config.setApprovalMode(ApprovalMode.AUTO_EDIT); + } + }, + }; + return confirmationDetails; + } + + private async prepareEdit( + abortSignal: AbortSignal, + ): Promise { + const preDecision = await checkPriorNotebookRead( + this.config, + this.params.notebook_path, + ); + if (!preDecision.ok) { + throw new StructuredToolError(preDecision.rawMessage, preDecision.type); + } + + let originalContent: string; + let bom = false; + let encoding: string | undefined; + let lineEnding: LineEnding = 'lf'; + try { + const fileInfo = await this.config.getFileSystemService().readTextFile({ + path: this.params.notebook_path, + }); + originalContent = fileInfo.content; + bom = fileInfo._meta?.bom ?? false; + encoding = fileInfo._meta?.encoding; + lineEnding = + fileInfo._meta?.lineEnding ?? detectLineEnding(fileInfo.content); + } catch (error) { + if (abortSignal.aborted) { + throw error; + } + const code = (error as NodeJS.ErrnoException | undefined)?.code; + if (code === 'ENOENT') { + throw new StructuredToolError( + `Notebook file not found: ${this.params.notebook_path}`, + ToolErrorType.FILE_NOT_FOUND, + ); + } + throw new StructuredToolError( + `Error reading notebook: ${ + error instanceof Error ? error.message : String(error) + }`, + ToolErrorType.READ_CONTENT_FAILURE, + ); + } + + const postDecision = await checkPriorNotebookRead( + this.config, + this.params.notebook_path, + { expectExisting: true }, + ); + if (!postDecision.ok) { + throw new StructuredToolError(postDecision.rawMessage, postDecision.type); + } + + try { + if (this.modifyMetadata?.modifiedNotebookContent !== undefined) { + return { + ...prepareModifiedNotebookContent( + originalContent, + this.modifyMetadata.modifiedNotebookContent, + this.params, + ), + originalContent, + bom, + encoding, + lineEnding, + }; + } + + return { + ...applyNotebookEdit(originalContent, this.params), + originalContent, + bom, + encoding, + lineEnding, + }; + } catch (error) { + if (error instanceof NotebookEditError) { + throw new StructuredToolError(error.message, error.type); + } + throw error; + } + } + + override async execute(signal: AbortSignal): Promise { + let prepared: PreparedNotebookEdit; + try { + prepared = await this.prepareEdit(signal); + } catch (error) { + if (signal.aborted) { + throw error; + } + const errorType = + error instanceof StructuredToolError + ? error.errorType + : ToolErrorType.NOTEBOOK_EDIT_FAILURE; + const message = error instanceof Error ? error.message : String(error); + return { + llmContent: message, + returnDisplay: `Error: ${message}`, + error: { + message, + type: errorType, + }, + }; + } + + try { + try { + await this.config + .getFileHistoryService() + .trackEdit(this.params.notebook_path); + } catch { + // File history is best-effort; never block core tool operations. + } + + const writeDecision = await checkPriorNotebookRead( + this.config, + this.params.notebook_path, + { expectExisting: true }, + ); + if (!writeDecision.ok) { + return { + llmContent: writeDecision.rawMessage, + returnDisplay: `Error: ${writeDecision.displayMessage}`, + error: { + message: writeDecision.rawMessage, + type: writeDecision.type, + }, + }; + } + + await this.config.getFileSystemService().writeTextFile({ + path: this.params.notebook_path, + content: prepared.updatedContent, + _meta: { + bom: prepared.bom, + encoding: prepared.encoding, + lineEnding: prepared.lineEnding, + }, + }); + + if (!this.modifyMetadata?.modifiedByUser) { + CommitAttributionService.getInstance().recordEdit( + this.params.notebook_path, + prepared.originalContent, + prepared.updatedContent, + ); + } + + try { + const postWriteStats = fs.statSync(this.params.notebook_path); + const cache = this.config.getFileReadCache(); + if (prepared.requiresReadAfterWrite) { + cache.invalidate(postWriteStats); + } else { + cache.recordWrite(this.params.notebook_path, postWriteStats, { + cacheable: false, + }); + } + } catch (error) { + // Non-fatal: the write succeeded. A subsequent read will re-stat and + // refresh the cache if this best-effort cache update failed. + debugLogger.warn( + `[NotebookEdit] post-write cache update failed for ${this.params.notebook_path}: ${ + error instanceof Error ? error.message : String(error) + }`, + ); + } + + const fileName = path.basename(this.params.notebook_path); + const fileDiff = Diff.createPatch( + fileName, + prepared.originalContent, + prepared.updatedContent, + 'Current', + 'Proposed', + DEFAULT_DIFF_OPTIONS, + ); + const diffStat = getDiffStat( + fileName, + prepared.originalContent, + this.modifyMetadata?.aiProposedContent ?? prepared.updatedContent, + prepared.updatedContent, + ); + + logFileOperation( + this.config, + new FileOperationEvent( + NotebookEditTool.Name, + FileOperation.UPDATE, + prepared.updatedContent.split('\n').length, + getSpecificMimeType(this.params.notebook_path), + '.ipynb', + prepared.language, + ), + ); + + const displayResult = { + fileDiff, + fileName, + originalContent: prepared.originalContent, + newContent: prepared.updatedContent, + diffStat, + }; + + const llmContent = + this.modifyMetadata?.modifiedNotebookContent !== undefined + ? `Notebook ${this.params.notebook_path} has been updated. Notebook content was modified by the user before approval; the final saved notebook may differ from the original ${prepared.mode} cell ${prepared.editedCellId} proposal.` + : `Notebook ${this.params.notebook_path} has been updated. ${prepared.mode} cell ${prepared.editedCellId}.${ + prepared.mode === 'delete' + ? '' + : `\n\nUpdated source:\n\n---\n\n${normalizeSource(this.params.new_source ?? '')}` + }`; + return { + llmContent, + returnDisplay: displayResult, + resultFilePaths: [this.params.notebook_path], + }; + } catch (error) { + const message = error instanceof Error ? error.message : String(error); + return { + llmContent: `Error writing notebook: ${message}`, + returnDisplay: `Error writing notebook: ${message}`, + error: { + message, + type: ToolErrorType.FILE_WRITE_FAILURE, + }, + }; + } + } +} + +export class NotebookEditTool + extends BaseDeclarativeTool + implements ModifiableDeclarativeTool +{ + static readonly Name = ToolNames.NOTEBOOK_EDIT; + private readonly modifyMetadataByParams = new WeakMap< + NotebookEditToolParams, + NotebookEditModifyMetadata + >(); + private readonly modifyMetadataByKey = new Map< + string, + NotebookEditModifyMetadata[] + >(); + + constructor(private readonly config: Config) { + super( + NotebookEditTool.Name, + ToolDisplayNames.NOTEBOOK_EDIT, + `Edits a Jupyter notebook (.ipynb) safely at the cell level. Use this instead of ${ToolNames.EDIT} or ${ToolNames.WRITE_FILE} for notebook cells. Supports replacing, inserting, and deleting cells. Always read the notebook first with ${ToolNames.READ_FILE}; then use the cell IDs shown in that output.`, + Kind.Edit, + { + properties: { + notebook_path: { + description: + 'Absolute path to the Jupyter notebook file to edit. Must end with .ipynb.', + type: 'string', + }, + cell_id: { + description: + 'Target cell ID from read_file output, or cell-N 0-based fallback. Required for replace and delete. For insert, the new cell is inserted after this cell; if omitted, inserted at the beginning.', + type: 'string', + }, + new_source: { + description: + 'New source content for replace and insert operations. Not required for delete.', + type: 'string', + }, + cell_type: { + description: + 'Cell type for inserted cells or type conversion on replace.', + type: 'string', + enum: ['code', 'markdown'], + }, + edit_mode: { + description: 'Notebook edit operation. Defaults to replace.', + type: 'string', + enum: ['replace', 'insert', 'delete'], + }, + }, + required: ['notebook_path'], + additionalProperties: false, + type: 'object', + }, + ); + } + + protected override validateToolParamValues( + params: NotebookEditToolParams, + ): string | null { + params.notebook_path = unescapePath(params.notebook_path.trim()); + + if (!params.notebook_path) { + return "The 'notebook_path' parameter must be non-empty."; + } + + if (!path.isAbsolute(params.notebook_path)) { + return `Notebook path must be absolute: ${params.notebook_path}`; + } + + if (path.extname(params.notebook_path).toLowerCase() !== '.ipynb') { + return 'File must be a Jupyter notebook (.ipynb). Use the edit tool for other file types.'; + } + + const mode = params.edit_mode ?? 'replace'; + if (!['replace', 'insert', 'delete'].includes(mode)) { + return "edit_mode must be 'replace', 'insert', or 'delete'."; + } + + if (params.cell_type && !['code', 'markdown'].includes(params.cell_type)) { + return "cell_type must be 'code' or 'markdown'."; + } + + if (mode !== 'insert' && !params.cell_id) { + return 'cell_id is required for replace and delete operations.'; + } + + if (mode !== 'delete' && typeof params.new_source !== 'string') { + return `new_source is required when edit_mode is "${mode}".`; + } + + const fileService = this.config.getFileService(); + if (fileService.shouldQwenIgnoreFile(params.notebook_path)) { + return `File path '${params.notebook_path}' is ignored by .qwenignore pattern(s).`; + } + + return null; + } + + protected createInvocation( + params: NotebookEditToolParams, + ): ToolInvocation { + return new NotebookEditInvocation( + this.config, + params, + this.consumeModifyMetadata(params), + ); + } + + private getModifyMetadataKey(params: NotebookEditToolParams): string { + return JSON.stringify({ + notebook_path: unescapePath(params.notebook_path.trim()), + cell_id: params.cell_id, + new_source: params.new_source, + cell_type: params.cell_type, + edit_mode: params.edit_mode, + }); + } + + private rememberModifyMetadata( + params: NotebookEditToolParams, + metadata: NotebookEditModifyMetadata, + ): void { + this.modifyMetadataByParams.set(params, metadata); + const key = this.getModifyMetadataKey(params); + const queue = this.modifyMetadataByKey.get(key) ?? []; + queue.push(metadata); + this.modifyMetadataByKey.set(key, queue); + } + + private consumeModifyMetadata( + params: NotebookEditToolParams, + ): NotebookEditModifyMetadata | undefined { + const metadata = this.modifyMetadataByParams.get(params); + if (metadata) { + this.modifyMetadataByParams.delete(params); + this.removeQueuedModifyMetadata(params, metadata); + return metadata; + } + + const key = this.getModifyMetadataKey(params); + const queue = this.modifyMetadataByKey.get(key); + const queuedMetadata = queue?.shift(); + if (queue && queue.length === 0) { + this.modifyMetadataByKey.delete(key); + } + return queuedMetadata; + } + + private removeQueuedModifyMetadata( + params: NotebookEditToolParams, + metadata: NotebookEditModifyMetadata, + ): void { + const key = this.getModifyMetadataKey(params); + const queue = this.modifyMetadataByKey.get(key); + if (!queue) { + return; + } + + const index = queue.indexOf(metadata); + if (index !== -1) { + queue.splice(index, 1); + } + if (queue.length === 0) { + this.modifyMetadataByKey.delete(key); + } + } + + getModifyContext( + _abortSignal: AbortSignal, + ): ModifyContext { + const currentContentSnapshots = new Map(); + const readCurrentContentSnapshot = async ( + params: NotebookEditToolParams, + ): Promise => { + const cached = currentContentSnapshots.get(params.notebook_path); + if (cached !== undefined) { + return cached; + } + + const { content } = await this.config + .getFileSystemService() + .readTextFile({ path: params.notebook_path }); + currentContentSnapshots.set(params.notebook_path, content); + return content; + }; + + return { + getFilePath: (params: NotebookEditToolParams) => params.notebook_path, + getCurrentContent: async ( + params: NotebookEditToolParams, + ): Promise => readCurrentContentSnapshot(params), + getProposedContent: async ( + params: NotebookEditToolParams, + ): Promise => { + const content = await readCurrentContentSnapshot(params); + return applyNotebookEdit(content, params).updatedContent; + }, + createUpdatedParams: ( + oldContent: string, + modifiedProposedContent: string, + originalParams: NotebookEditToolParams, + ): NotebookEditToolParams => { + let aiProposedContent = + this.modifyMetadataByParams.get(originalParams)?.aiProposedContent; + if (aiProposedContent === undefined) { + try { + aiProposedContent = applyNotebookEdit( + oldContent, + originalParams, + ).updatedContent; + } catch { + aiProposedContent = modifiedProposedContent; + } + } + const updatedParams: NotebookEditToolParams = { + ...originalParams, + }; + this.rememberModifyMetadata(updatedParams, { + aiProposedContent, + modifiedByUser: true, + modifiedNotebookContent: modifiedProposedContent, + }); + return updatedParams; + }, + }; + } +} diff --git a/packages/core/src/tools/priorReadEnforcement.ts b/packages/core/src/tools/priorReadEnforcement.ts index e7c61eb94..0d2d5eba1 100644 --- a/packages/core/src/tools/priorReadEnforcement.ts +++ b/packages/core/src/tools/priorReadEnforcement.ts @@ -289,8 +289,11 @@ export async function checkPriorRead( `notebook payload that the ${ToolNames.READ_FILE} tool returns ` + `as a structured value rather than as plain text. The Edit / ` + `WriteFile tools cannot mutate that payload safely — re-reading ` + - `it would not change this. Use a different mechanism (e.g. shell ` + - `tool with a binary-aware writer) if you need to ${verbBare} it.`; + `it would not change this. If this is a Jupyter notebook (.ipynb), ` + + `use the ${ToolNames.NOTEBOOK_EDIT} tool for cell-level edits after ` + + `reading it. For other non-text files, use a different mechanism ` + + `(e.g. shell tool with an appropriate writer) if you need to ` + + `${verbBare} it.`; return { ok: false, type: ToolErrorType.EDIT_REQUIRES_PRIOR_READ, diff --git a/packages/core/src/tools/read-file.test.ts b/packages/core/src/tools/read-file.test.ts index ab59c3029..b0e19b617 100644 --- a/packages/core/src/tools/read-file.test.ts +++ b/packages/core/src/tools/read-file.test.ts @@ -158,6 +158,29 @@ describe('ReadFileTool', () => { 'Limit must be a positive number', ); }); + + it('should reject offset or limit for notebook files', () => { + const params: ReadFileToolParams = { + file_path: path.join(tempRootDir, 'test.ipynb'), + offset: 0, + limit: 10, + }; + + expect(() => tool.build(params)).toThrow( + 'offset and limit are not supported for Jupyter notebook (.ipynb) files', + ); + }); + + it('should reject pages for notebook files', () => { + const params: ReadFileToolParams = { + file_path: path.join(tempRootDir, 'test.ipynb'), + pages: '1', + }; + + expect(() => tool.build(params)).toThrow( + 'pages is not supported for Jupyter notebook (.ipynb) files', + ); + }); }); describe('getDefaultPermission', () => { @@ -505,6 +528,36 @@ describe('ReadFileTool', () => { expect(result.returnDisplay).toBe('Read notebook: test.ipynb'); }); + it('records truncated notebook reads as not full', async () => { + const nbPath = path.join(tempRootDir, 'large.ipynb'); + const notebook = { + cells: Array.from({ length: 200 }, (_, i) => ({ + cell_type: 'code', + source: ['x = ' + 'a'.repeat(600) + '\n'], + execution_count: i + 1, + outputs: [{ output_type: 'stream', text: ['result '.repeat(100)] }], + metadata: {}, + })), + metadata: { language_info: { name: 'python' } }, + }; + await fsp.writeFile(nbPath, JSON.stringify(notebook), 'utf-8'); + const invocation = tool.build({ + file_path: nbPath, + }) as ToolInvocation; + + const result = await invocation.execute(abortSignal); + expect(typeof result.llmContent).toBe('string'); + expect(result.llmContent).toContain('remaining cells truncated'); + expect(result.llmContent).not.toContain('Showing lines'); + + const status = fileReadCache.check(fs.statSync(nbPath)); + expect(status.state).toBe('fresh'); + if (status.state === 'fresh') { + expect(status.entry.lastReadWasFull).toBe(false); + expect(status.entry.lastReadCacheable).toBe(false); + } + }); + it('should reject invalid pages parameter', () => { const params: ReadFileToolParams = { file_path: '/tmp/test.pdf', @@ -544,7 +597,9 @@ describe('ReadFileTool', () => { file_path: path.join(tempRootDir, 'test.txt'), pages: '', }; - expect(() => tool.build(params)).not.toThrow(); + const invocation = tool.build(params); + + expect(invocation.params.pages).toBeUndefined(); }); it('should support offset and limit for text files', async () => { diff --git a/packages/core/src/tools/read-file.ts b/packages/core/src/tools/read-file.ts index 4e8528c80..a7c9ea714 100644 --- a/packages/core/src/tools/read-file.ts +++ b/packages/core/src/tools/read-file.ts @@ -156,10 +156,9 @@ class ReadFileToolInvocation extends BaseToolInvocation< const cacheEnabled = !this.config.getFileReadCacheDisabled(); const useFastPath = cacheEnabled && !isAutoMem; const cache = this.config.getFileReadCache(); - // A "full" Read consumes the whole file: no offset, no limit, no PDF - // page range. Only full Reads are eligible for the file_unchanged - // fast-path; range-scoped Reads always go through, since the model - // may legitimately ask for a different slice next time. + // A request-level "full" Read asks for the whole file: no offset, + // no limit, no PDF page range. The cache entry is only marked as + // full later if the produced content was not truncated. const isFullRead = this.params.offset === undefined && this.params.limit === undefined && @@ -245,7 +244,9 @@ class ReadFileToolInvocation extends BaseToolInvocation< // level if the produced content was not truncated, otherwise // the model only saw the head and a follow-up `file_unchanged` // placeholder would falsely imply "you've already seen - // everything". + // everything". NotebookEdit also requires this flag so a + // truncated notebook render does not authorize structured writes + // against unseen cells. // // The stat we record is the one taken inside `processSingleFileContent` // and surfaced via `result.stats`. The internal stat happens @@ -274,7 +275,11 @@ class ReadFileToolInvocation extends BaseToolInvocation< } let llmContent: PartUnion; - if (result.isTruncated) { + if ( + result.isTruncated && + result.linesShown && + result.originalLineCount !== undefined + ) { const [start, end] = result.linesShown!; const total = result.originalLineCount!; llmContent = `Showing lines ${start}-${end} of ${total} total lines.\n\n---\n\n${result.llmContent}`; @@ -432,6 +437,23 @@ export class ReadFileTool extends BaseDeclarativeTool< return 'Limit must be a positive number'; } + if (params.pages !== undefined) { + const pages = params.pages.trim(); + params.pages = pages.length > 0 ? pages : undefined; + } + + const ext = path.extname(filePath).toLowerCase(); + if ( + (params.offset !== undefined || params.limit !== undefined) && + ext === '.ipynb' + ) { + return 'offset and limit are not supported for Jupyter notebook (.ipynb) files. Notebooks are always read in full with structured cell output.'; + } + + if (params.pages !== undefined && ext === '.ipynb') { + return 'pages is not supported for Jupyter notebook (.ipynb) files. Notebooks are always read in full with structured cell output.'; + } + if (params.pages) { const parsed = parsePDFPageRange(params.pages); if (!parsed) { diff --git a/packages/core/src/tools/tool-error.ts b/packages/core/src/tools/tool-error.ts index 84a2a884b..077d0aec4 100644 --- a/packages/core/src/tools/tool-error.ts +++ b/packages/core/src/tools/tool-error.ts @@ -46,17 +46,11 @@ export enum ToolErrorType { // view, so the model has not seen the full text content the // mutation could touch. // 3. The file is a structural dead end that no amount of - // re-reading can change: - // - non-text payloads (binary / image / audio / video / - // PDF / notebook) — read_file returns these as - // structured values that Edit / WriteFile cannot mutate - // safely; the rejection message tells the model to use - // a different tool (shell with a binary-aware writer). - // - special files (FIFO / socket / character or block - // device) — read_file rejects these as "not a regular - // file", so an enforcement loop on read_file would - // never terminate; the rejection points the model at - // shell instead. + // re-reading can change: non-text payloads (binary / image / + // audio / video / PDF / notebook). read_file returns these as + // structured values that Edit / WriteFile cannot mutate safely; + // the rejection message tells the model to use a different tool + // (shell with a binary-aware writer). // // Despite the `EDIT_` prefix this code is shared between EditTool // and WriteFileTool: the boundary it guards is "the model is about @@ -88,6 +82,15 @@ export enum ToolErrorType { // cannot verify. Operators monitoring on error codes can route this // separately. PRIOR_READ_VERIFICATION_FAILED = 'prior_read_verification_failed', + // Returned when a path resolves but is not a regular file (FIFO / socket / + // character or block device). Re-reading cannot make these editable, so this + // is distinct from EDIT_REQUIRES_PRIOR_READ to avoid read/edit retry loops. + TARGET_NOT_REGULAR_FILE = 'target_not_regular_file', + + // Notebook-specific Errors + NOTEBOOK_EDIT_FAILURE = 'notebook_edit_failure', + NOTEBOOK_INVALID_JSON = 'notebook_invalid_json', + NOTEBOOK_CELL_NOT_FOUND = 'notebook_cell_not_found', // Glob-specific Errors GLOB_EXECUTION_ERROR = 'glob_execution_error', diff --git a/packages/core/src/tools/tool-names.ts b/packages/core/src/tools/tool-names.ts index 25a11bf63..a8dc79426 100644 --- a/packages/core/src/tools/tool-names.ts +++ b/packages/core/src/tools/tool-names.ts @@ -40,6 +40,7 @@ export const ToolNames = { SEND_MESSAGE: 'send_message', STRUCTURED_OUTPUT: 'structured_output', MONITOR: 'monitor', + NOTEBOOK_EDIT: 'notebook_edit', TOOL_SEARCH: 'tool_search', ENTER_WORKTREE: 'enter_worktree', EXIT_WORKTREE: 'exit_worktree', @@ -73,6 +74,7 @@ export const ToolDisplayNames = { SEND_MESSAGE: 'SendMessage', STRUCTURED_OUTPUT: 'StructuredOutput', MONITOR: 'Monitor', + NOTEBOOK_EDIT: 'NotebookEdit', TOOL_SEARCH: 'ToolSearch', ENTER_WORKTREE: 'EnterWorktree', EXIT_WORKTREE: 'ExitWorktree', diff --git a/packages/core/src/tools/write-file.test.ts b/packages/core/src/tools/write-file.test.ts index d9e4b8fd8..9905f829c 100644 --- a/packages/core/src/tools/write-file.test.ts +++ b/packages/core/src/tools/write-file.test.ts @@ -1116,6 +1116,7 @@ describe('WriteFileTool', () => { .build({ file_path: filePath, content: 'clobber' }) .execute(abortSignal); expect(result.error?.type).toBe(ToolErrorType.EDIT_REQUIRES_PRIOR_READ); + expect(result.error?.message).toContain('notebook_edit'); // Verb in the dead-end guidance must read correctly for // overwrite (the WriteFile path), not "edit". expect(result.error?.message).toMatch(/if you need to overwrite it\./); diff --git a/packages/core/src/utils/fileUtils.ts b/packages/core/src/utils/fileUtils.ts index 60d9b3f9d..48db20c38 100644 --- a/packages/core/src/utils/fileUtils.ts +++ b/packages/core/src/utils/fileUtils.ts @@ -22,7 +22,7 @@ import { isNodeError } from './errors.js'; import type { InputModalities } from '../core/contentGenerator.js'; import { detectEncodingFromBuffer } from './systemEncoding.js'; import { extractPDFText, parsePDFPageRange } from './pdf.js'; -import { readNotebook } from './notebook.js'; +import { readNotebookWithMetadata } from './notebook.js'; const debugLogger = createDebugLogger('FILE_UTILS'); @@ -818,7 +818,7 @@ export interface ProcessedFileReadResult { error?: string; // Optional error message for the LLM if file processing failed errorType?: ToolErrorType; // Structured error type originalLineCount?: number; // For text files, the total number of lines in the original file - isTruncated?: boolean; // For text files, indicates if content was truncated + isTruncated?: boolean; // Indicates if displayed content was truncated linesShown?: [number, number]; // For text files [startLine, endLine] (1-based for display) /** * The Stats taken at the start of the read pipeline, before the @@ -1172,10 +1172,13 @@ export async function processSingleFileContent( } case 'notebook': { try { - const content = await readNotebook(filePath); + const { content, isTruncated } = + await readNotebookWithMetadata(filePath); return { llmContent: content, returnDisplay: `Read notebook: ${relativePathForDisplay}`, + isTruncated, + stats, }; } catch (e: unknown) { const msg = e instanceof Error ? e.message : String(e); diff --git a/packages/core/src/utils/notebook.test.ts b/packages/core/src/utils/notebook.test.ts index 9ff577eac..2f9419d9e 100644 --- a/packages/core/src/utils/notebook.test.ts +++ b/packages/core/src/utils/notebook.test.ts @@ -5,7 +5,21 @@ */ import { describe, it, expect, afterEach } from 'vitest'; -import { readNotebook } from './notebook.js'; +import { + findCellIndex, + getCellDisplayId, + hasStableCellIds, + inferInsertedCellSourceArrayStyle, + inferNotebookJsonFormat, + isAmbiguousCellId, + makeCellId, + parseCellId, + parseNotebook, + readNotebook, + readNotebookWithMetadata, + serializeNotebook, + toNotebookSource, +} from './notebook.js'; import path from 'node:path'; import os from 'node:os'; import fsp from 'node:fs/promises'; @@ -369,6 +383,26 @@ describe('notebook utilities', () => { expect(result.length).toBeLessThan(120000); }); + it('reports when notebook cell rendering is truncated', async () => { + const cells = Array.from({ length: 200 }, (_, i) => ({ + cell_type: 'code' as const, + source: ['x = ' + 'a'.repeat(600) + '\n'], + execution_count: i + 1, + outputs: [ + { output_type: 'stream' as const, text: ['result '.repeat(100)] }, + ], + metadata: {}, + })); + const filePath = await writeNotebook('big-metadata.ipynb', { + cells, + metadata: { language_info: { name: 'python' } }, + }); + + const result = await readNotebookWithMetadata(filePath); + expect(result.isTruncated).toBe(true); + expect(result.content).toContain('remaining cells truncated'); + }); + it('should throw on invalid JSON', async () => { tempDir = await fsp.mkdtemp(path.join(os.tmpdir(), 'notebook-test-')); const filePath = path.join(tempDir, 'bad.ipynb'); @@ -376,4 +410,184 @@ describe('notebook utilities', () => { await expect(readNotebook(filePath)).rejects.toThrow(); }); + + it('should parse notebooks with a leading UTF-8 BOM', async () => { + const notebook = { + cells: [ + { + cell_type: 'code', + source: ['print("bom")'], + execution_count: null, + outputs: [], + metadata: {}, + }, + ], + metadata: { language_info: { name: 'python' } }, + }; + tempDir = await fsp.mkdtemp(path.join(os.tmpdir(), 'notebook-test-')); + const filePath = path.join(tempDir, 'bom.ipynb'); + await fsp.writeFile(filePath, `\ufeff${JSON.stringify(notebook)}`, 'utf-8'); + + expect( + parseNotebook(`\ufeff${JSON.stringify(notebook)}`).cells, + ).toHaveLength(1); + await expect(readNotebookWithMetadata(filePath)).resolves.toMatchObject({ + isTruncated: false, + }); + }); + + it('should parse cell-N IDs as zero-based indexes', () => { + expect(parseCellId('cell-0')).toBe(0); + expect(parseCellId('cell-12')).toBe(12); + expect(parseCellId('abc-12')).toBeUndefined(); + expect(parseCellId('cell-nope')).toBeUndefined(); + }); + + it('should find cells by the same IDs read_file displays', () => { + const notebook = parseNotebook( + JSON.stringify({ + cells: [ + { cell_type: 'code', id: 'real-id', source: '', metadata: {} }, + { cell_type: 'code', source: '', metadata: {} }, + ], + metadata: {}, + }), + ); + + expect(getCellDisplayId(notebook.cells[0]!, 0)).toBe('real-id'); + expect(getCellDisplayId(notebook.cells[1]!, 1)).toBe('cell-1'); + expect(findCellIndex(notebook, 'real-id')).toBe(0); + expect(findCellIndex(notebook, 'cell-1')).toBe(1); + expect(findCellIndex(notebook, 'cell-0')).toBe(-1); + expect(findCellIndex(notebook, 'missing')).toBe(-1); + }); + + it('should reject ambiguous displayed cell IDs', () => { + const notebook = parseNotebook( + JSON.stringify({ + cells: [ + { cell_type: 'code', id: 'cell-1', source: '', metadata: {} }, + { cell_type: 'code', source: '', metadata: {} }, + ], + metadata: {}, + }), + ); + + expect(isAmbiguousCellId(notebook, 'cell-1')).toBe(true); + expect(findCellIndex(notebook, 'cell-1')).toBe(-1); + }); + + it('should preserve newline boundaries when converting source to arrays', () => { + expect(toNotebookSource('a\nb\n', true)).toEqual(['a\n', 'b\n']); + expect(toNotebookSource('a\nb', true)).toEqual(['a\n', 'b']); + expect(toNotebookSource('', true)).toEqual([]); + expect(toNotebookSource('a\nb\n', false)).toBe('a\nb\n'); + }); + + it('should preserve notebook JSON indentation and trailing newline style', () => { + const raw = JSON.stringify( + { + cells: [{ cell_type: 'markdown', source: '# Title', metadata: {} }], + metadata: {}, + }, + null, + 2, + ); + const notebook = parseNotebook(raw); + notebook.cells[0]!.source = '# Updated'; + + const format = inferNotebookJsonFormat(raw); + const serialized = serializeNotebook(notebook, format); + + expect(format).toEqual({ indent: 2, trailingNewline: false }); + expect(serialized).toContain('\n "cells"'); + expect(serialized.endsWith('\n')).toBe(false); + }); + + it('should preserve compact notebook JSON when serializing after edits', () => { + const raw = JSON.stringify({ + cells: [{ cell_type: 'markdown', source: '# Title', metadata: {} }], + metadata: {}, + }); + const notebook = parseNotebook(raw); + notebook.cells[0]!.source = '# Updated'; + + const format = inferNotebookJsonFormat(raw); + const serialized = serializeNotebook(notebook, format); + + expect(format).toEqual({ indent: undefined, trailingNewline: false }); + expect(serialized).toBe(JSON.stringify(notebook)); + }); + + it('should infer inserted source style from adjacent cells', () => { + const notebook = parseNotebook( + JSON.stringify({ + cells: [ + { cell_type: 'markdown', source: '# string source', metadata: {} }, + { + cell_type: 'code', + source: ['print("array source")'], + metadata: {}, + }, + ], + metadata: {}, + }), + ); + + expect(inferInsertedCellSourceArrayStyle(notebook, 1)).toBe(false); + expect(inferInsertedCellSourceArrayStyle(notebook, 0)).toBe(false); + expect(inferInsertedCellSourceArrayStyle(notebook, 2)).toBe(true); + }); + + it('should generate deterministic cell IDs that cannot collide with cell-N fallbacks', () => { + const notebook = parseNotebook( + JSON.stringify({ + nbformat: 4, + nbformat_minor: 5, + cells: [ + { cell_type: 'code', id: 'qwen-cell-1', source: '', metadata: {} }, + { cell_type: 'code', source: '', metadata: {} }, + ], + metadata: {}, + }), + ); + + expect(hasStableCellIds(notebook)).toBe(false); + expect(makeCellId(notebook)).toBe('qwen-cell-2'); + notebook.cells.push({ + cell_type: 'code', + id: 'qwen-cell-2', + source: '', + metadata: {}, + }); + expect(makeCellId(notebook)).toBe('qwen-cell-3'); + }); + + it('should not generate cell IDs for old notebook formats', () => { + const notebook = parseNotebook( + JSON.stringify({ + nbformat: 4, + nbformat_minor: 4, + cells: [], + metadata: {}, + }), + ); + + expect(makeCellId(notebook)).toBeUndefined(); + }); + + it('should reject notebook JSON without a cells array', () => { + expect(() => parseNotebook(JSON.stringify({ metadata: {} }))).toThrow( + 'missing cells array', + ); + }); + + it('should reject non-object notebook cells', () => { + expect(() => + parseNotebook(JSON.stringify({ cells: [null], metadata: {} })), + ).toThrow('cell at index 0 is not an object'); + expect(() => + parseNotebook(JSON.stringify({ cells: ['not a cell'], metadata: {} })), + ).toThrow('cell at index 0 is not an object'); + }); }); diff --git a/packages/core/src/utils/notebook.ts b/packages/core/src/utils/notebook.ts index 838f7d4b4..42da415eb 100644 --- a/packages/core/src/utils/notebook.ts +++ b/packages/core/src/utils/notebook.ts @@ -42,7 +42,7 @@ function sanitizeMimeTypes(keys: string[]): string[] { /** * Jupyter Notebook cell output types. */ -interface NotebookCellOutput { +export interface NotebookCellOutput { output_type: 'stream' | 'execute_result' | 'display_data' | 'error'; text?: string | string[]; data?: Record; @@ -54,9 +54,13 @@ interface NotebookCellOutput { /** * Jupyter Notebook cell. */ -interface NotebookCell { - cell_type: 'code' | 'markdown' | 'raw'; +export type NotebookCellType = 'code' | 'markdown' | 'raw'; +export type EditableNotebookCellType = 'code' | 'markdown'; + +export interface NotebookCell { + cell_type: NotebookCellType; source: string | string[]; + metadata?: Record; outputs?: NotebookCellOutput[]; execution_count?: number | null; id?: string; @@ -65,16 +69,218 @@ interface NotebookCell { /** * Jupyter Notebook top-level structure. */ -interface NotebookContent { +export interface NotebookContent { cells: NotebookCell[]; - metadata: { + metadata?: { language_info?: { name?: string }; kernelspec?: { language?: string; display_name?: string }; + [key: string]: unknown; + }; + nbformat?: number; + nbformat_minor?: number; +} + +export interface NotebookReadResult { + content: string; + isTruncated: boolean; +} + +export interface NotebookJsonFormat { + indent?: number; + trailingNewline: boolean; +} + +export function normalizeSource(source: string | string[]): string { + return Array.isArray(source) ? source.join('') : source; +} + +export function parseNotebook(content: string): NotebookContent { + const jsonContent = + content.charCodeAt(0) === 0xfeff ? content.slice(1) : content; + const parsed = JSON.parse(jsonContent) as unknown; + if (!parsed || typeof parsed !== 'object') { + throw new Error('Invalid notebook: expected a JSON object'); + } + + const notebook = parsed as NotebookContent; + if (!Array.isArray(notebook.cells)) { + throw new Error('Invalid notebook: missing cells array'); + } + + for (let i = 0; i < notebook.cells.length; i++) { + const cell = notebook.cells[i]; + if (!cell || typeof cell !== 'object' || Array.isArray(cell)) { + throw new Error(`Invalid notebook: cell at index ${i} is not an object`); + } + } + + return notebook; +} + +export function inferNotebookJsonFormat(content: string): NotebookJsonFormat { + const lineMatch = content.match(/\n( +)"/); + return { + indent: lineMatch?.[1]?.length, + trailingNewline: content.endsWith('\n'), }; } -function normalizeSource(source: string | string[]): string { - return Array.isArray(source) ? source.join('') : source; +export function serializeNotebook( + notebook: NotebookContent, + format: NotebookJsonFormat = { indent: 1, trailingNewline: true }, +): string { + const serialized = JSON.stringify(notebook, null, format.indent); + return format.trailingNewline ? `${serialized}\n` : serialized; +} + +export function parseCellId(cellId: string): number | undefined { + const match = cellId.match(/^cell-(\d+)$/); + if (!match?.[1]) { + return undefined; + } + + const index = Number.parseInt(match[1], 10); + return Number.isNaN(index) ? undefined : index; +} + +export function getCellDisplayId(cell: NotebookCell, index: number): string { + return typeof cell.id === 'string' && cell.id.length > 0 + ? cell.id + : `cell-${index}`; +} + +export function hasStableCellIds(notebook: NotebookContent): boolean { + return notebook.cells.every( + (cell) => typeof cell.id === 'string' && cell.id.length > 0, + ); +} + +export function findCellIndexesByDisplayId( + notebook: NotebookContent, + cellId: string, +): number[] { + const indexes: number[] = []; + notebook.cells.forEach((cell, index) => { + if (getCellDisplayId(cell, index) === cellId) { + indexes.push(index); + } + }); + return indexes; +} + +export function isAmbiguousCellId( + notebook: NotebookContent, + cellId: string, +): boolean { + return findCellIndexesByDisplayId(notebook, cellId).length > 1; +} + +export function findCellIndex( + notebook: NotebookContent, + cellId: string, +): number { + const indexes = findCellIndexesByDisplayId(notebook, cellId); + return indexes.length === 1 ? indexes[0]! : -1; +} + +export function getNotebookLanguage(notebook: NotebookContent): string { + return ( + notebook.metadata?.language_info?.name ?? + notebook.metadata?.kernelspec?.language ?? + 'python' + ); +} + +export function shouldGenerateCellIds(notebook: NotebookContent): boolean { + return ( + (notebook.nbformat ?? 0) > 4 || + ((notebook.nbformat ?? 0) === 4 && (notebook.nbformat_minor ?? 0) >= 5) + ); +} + +export function makeCellId(notebook: NotebookContent): string | undefined { + if (!shouldGenerateCellIds(notebook)) { + return undefined; + } + + const existingDisplayIds = new Set( + notebook.cells.map((cell, index) => getCellDisplayId(cell, index)), + ); + + let fallbackIndex = 1; + let fallback = `qwen-cell-${fallbackIndex}`; + while (existingDisplayIds.has(fallback)) { + fallbackIndex++; + fallback = `qwen-cell-${fallbackIndex}`; + } + return fallback; +} + +export function inferNotebookSourceArrayStyle( + notebook: NotebookContent, +): boolean { + const sourceCell = notebook.cells.find((cell) => cell.source !== undefined); + return sourceCell ? Array.isArray(sourceCell.source) : true; +} + +export function inferInsertedCellSourceArrayStyle( + notebook: NotebookContent, + insertAt: number, +): boolean { + const previousCell = notebook.cells[insertAt - 1]; + if (previousCell?.source !== undefined) { + return Array.isArray(previousCell.source); + } + + const nextCell = notebook.cells[insertAt]; + if (nextCell?.source !== undefined) { + return Array.isArray(nextCell.source); + } + + return inferNotebookSourceArrayStyle(notebook); +} + +export function toNotebookSource( + source: string, + preferArray: boolean, +): string | string[] { + if (!preferArray) { + return source; + } + + if (source.length === 0) { + return []; + } + + const lines: string[] = []; + let start = 0; + for (let i = 0; i < source.length; i++) { + if (source[i] === '\n') { + lines.push(source.slice(start, i + 1)); + start = i + 1; + } + } + if (start < source.length) { + lines.push(source.slice(start)); + } + return lines; +} + +export function normalizeEditedCell( + cell: NotebookCell, + finalType: NotebookCellType, +): void { + cell.cell_type = finalType; + cell.metadata ??= {}; + + if (finalType === 'code') { + cell.execution_count = null; + cell.outputs = []; + return; + } + + delete cell.execution_count; + delete cell.outputs; } function processOutputText(text: string | string[] | undefined): string { @@ -131,7 +337,7 @@ function processCell( index: number, language: string, ): string { - const cellId = cell.id ?? `cell-${index}`; + const cellId = getCellDisplayId(cell, index); const source = normalizeSource(cell.source); const parts: string[] = []; @@ -181,29 +387,29 @@ function processCell( /** * Read and parse a Jupyter notebook file (.ipynb) into a structured text - * representation. Returns a formatted string with all cells and their outputs. + * representation, plus whether the cell listing was truncated. */ -export async function readNotebook(filePath: string): Promise { +export async function readNotebookWithMetadata( + filePath: string, +): Promise { const raw = await fs.promises.readFile(filePath, 'utf-8'); - const notebook: NotebookContent = JSON.parse(raw); - - const language = - notebook.metadata?.language_info?.name ?? - notebook.metadata?.kernelspec?.language ?? - 'python'; + const notebook = parseNotebook(raw); + const language = getNotebookLanguage(notebook); if (!notebook.cells || notebook.cells.length === 0) { - return '(empty notebook)'; + return { content: '(empty notebook)', isTruncated: false }; } const header = `Jupyter Notebook (${language}, ${notebook.cells.length} cells)`; const cellTexts: string[] = []; let totalLength = header.length; + let isTruncated = false; for (let i = 0; i < notebook.cells.length; i++) { const cellText = processCell(notebook.cells[i]!, i, language); totalLength += cellText.length + 2; // +2 for "\n\n" separator if (totalLength > MAX_NOTEBOOK_OUTPUT_CHARS) { + isTruncated = true; cellTexts.push( `... [${notebook.cells.length - i} remaining cells truncated, total ${notebook.cells.length} cells. Use shell to inspect: cat | jq '.cells[${i}:]']`, ); @@ -212,5 +418,16 @@ export async function readNotebook(filePath: string): Promise { cellTexts.push(cellText); } - return `${header}\n\n${cellTexts.join('\n\n')}`; + return { + content: `${header}\n\n${cellTexts.join('\n\n')}`, + isTruncated, + }; +} + +/** + * Read and parse a Jupyter notebook file (.ipynb) into a structured text + * representation. Returns a formatted string with all cells and their outputs. + */ +export async function readNotebook(filePath: string): Promise { + return (await readNotebookWithMetadata(filePath)).content; } diff --git a/packages/core/src/utils/readManyFiles.test.ts b/packages/core/src/utils/readManyFiles.test.ts index b49fc89e3..57806d79d 100644 --- a/packages/core/src/utils/readManyFiles.test.ts +++ b/packages/core/src/utils/readManyFiles.test.ts @@ -105,6 +105,35 @@ describe('readManyFiles', () => { expect(content).toContain('--- End of content ---'); }); + it('should include truncated notebooks that do not expose text line ranges', async () => { + const relativePath = 'large.ipynb'; + const absolutePath = path.join(tempRootDir, relativePath); + const cells = Array.from({ length: 30 }, (_, index) => ({ + cell_type: 'code', + id: `large-cell-${index}`, + source: [`value_${index} = "${'x'.repeat(4500)}"`], + metadata: {}, + outputs: [], + })); + await fs.writeFile( + absolutePath, + JSON.stringify({ cells, metadata: {} }), + 'utf-8', + ); + const mockConfig = createMockConfig(tempRootDir); + + const result = await readManyFiles(mockConfig, { paths: [relativePath] }); + + const content = contentToString(result.contentParts); + expect(content).toContain('Jupyter Notebook'); + expect(content).toContain('remaining cells truncated'); + expect(content).not.toContain( + 'No files matching the criteria were found', + ); + expect(result.files).toHaveLength(1); + expect(result.files[0]!.filePath).toBe(absolutePath); + }); + it('should return message when no files found', async () => { const mockConfig = createMockConfig(tempRootDir); diff --git a/packages/core/src/utils/readManyFiles.ts b/packages/core/src/utils/readManyFiles.ts index 7ece66e1e..5ddd1b478 100644 --- a/packages/core/src/utils/readManyFiles.ts +++ b/packages/core/src/utils/readManyFiles.ts @@ -201,7 +201,11 @@ async function readFileContent( if (typeof fileReadResult.llmContent === 'string') { let fileContentForLlm = ''; - if (fileReadResult.isTruncated) { + if ( + fileReadResult.isTruncated && + fileReadResult.linesShown && + fileReadResult.originalLineCount !== undefined + ) { const [start, end] = fileReadResult.linesShown!; const total = fileReadResult.originalLineCount!; fileContentForLlm = `Showing lines ${start}-${end} of ${total} total lines.\n---\n${fileReadResult.llmContent}`; diff --git a/packages/sdk-java/qwencode/QWEN.md b/packages/sdk-java/qwencode/QWEN.md index 7eb8e0038..0486fecaf 100644 --- a/packages/sdk-java/qwencode/QWEN.md +++ b/packages/sdk-java/qwencode/QWEN.md @@ -135,7 +135,7 @@ The SDK supports different permission modes for controlling tool execution: - **`default`**: Write tools are denied unless approved via `canUseTool` callback or in `allowedTools`. Read-only tools execute without confirmation. - **`plan`**: Blocks all write tools, instructing AI to present a plan first. -- **`auto-edit`**: Auto-approve edit tools (edit, write_file) while other tools require confirmation. +- **`auto-edit`**: Auto-approve edit tools (`edit`, `write_file`, `notebook_edit`) while other tools require confirmation. - **`yolo`**: All tools execute automatically without confirmation. ### Transport Options diff --git a/packages/sdk-java/qwencode/README.md b/packages/sdk-java/qwencode/README.md index 5da4d2a09..1934c4b5a 100644 --- a/packages/sdk-java/qwencode/README.md +++ b/packages/sdk-java/qwencode/README.md @@ -144,7 +144,7 @@ The SDK supports different permission modes for controlling tool execution: - **`default`**: Write tools are denied unless approved via `canUseTool` callback or in `allowedTools`. Read-only tools execute without confirmation. - **`plan`**: Blocks all write tools, instructing AI to present a plan first. -- **`auto-edit`**: Auto-approve edit tools (edit, write_file) while other tools require confirmation. +- **`auto-edit`**: Auto-approve edit tools (`edit`, `write_file`, `notebook_edit`) while other tools require confirmation. - **`yolo`**: All tools execute automatically without confirmation. ### Session Event Consumers and Assistant Content Consumers diff --git a/packages/sdk-typescript/README.md b/packages/sdk-typescript/README.md index f969ed2cb..a3f9271bf 100644 --- a/packages/sdk-typescript/README.md +++ b/packages/sdk-typescript/README.md @@ -227,7 +227,7 @@ The SDK supports different permission modes for controlling tool execution: - **`default`**: Write tools are denied unless approved via `canUseTool` callback or in `allowedTools`. Read-only tools execute without confirmation. - **`plan`**: Blocks all write tools, instructing AI to present a plan first. -- **`auto-edit`**: Auto-approve edit tools (edit, write_file) while other tools require confirmation. +- **`auto-edit`**: Auto-approve edit tools (`edit`, `write_file`, `notebook_edit`) while other tools require confirmation. - **`yolo`**: All tools execute automatically without confirmation. ### Permission Priority Chain