From 9270814f601d2afd3132d43c9dd490e4bbdc8ddd Mon Sep 17 00:00:00 2001 From: PTFOPlayer Date: Mon, 19 Oct 2026 08:50:31 +0200 Subject: [PATCH] Fix read tool context overflow on notebooks containing images MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Reading a .ipynb dumped the raw JSON into context, and embedded matplotlib figures are ~380 KB of base64 each — one user's notebook produced a 1.4 MB tool result and the provider rejected the request with "Range of input length should be [1, 1006592]". Since each payload sits on a single JSON line, paging with from/to couldn't dodge it either. `read` now renders notebooks with every image/* (and application/pdf) payload replaced by a size-recording placeholder, and any remaining output over 4,000 chars truncated, so a giant text/plain dump can't reproduce the same failure. Cell sources, stream text and execution counts are untouched. Invalid JSON falls back to cutting over-long lines at 2,000 chars. from/to index the rendered text, keeping line numbers stable across reads. The offending file goes from 1,458,837 chars to 6,485. 11 new tests; workspace fmt, clippy -D warnings, cargo doc and 631 tests pass. --- README.md | 2 +- TINYHARNESS.md | 1 + docs/tools-reference.md | 2 + tinyharness-lib/src/tools/read.rs | 423 +++++++++++++++++++++++++++++- 4 files changed, 426 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index c27aa2d..bafb7b4 100644 --- a/README.md +++ b/README.md @@ -170,7 +170,7 @@ Switch modes with `/mode `, use shortcut aliases (`/plan`, `/agent`, `/res | Tool | Category | Description | |------|----------|-------------| | `ls` | ReadOnly | List directory contents | -| `read` | ReadOnly | Read file content (also loads images for multimodal models) | +| `read` | ReadOnly | Read file content (loads images for multimodal models; renders `.ipynb` with base64 outputs stripped) | | `grep` | ReadOnly | Search with regex across files | | `glob` | ReadOnly | Find files by glob pattern | | `web_search` | ReadOnly | Search the web (requires Ollama API key) | diff --git a/TINYHARNESS.md b/TINYHARNESS.md index f2fc5e7..32cc870 100644 --- a/TINYHARNESS.md +++ b/TINYHARNESS.md @@ -96,6 +96,7 @@ Three crates in a Cargo workspace: - **Ctrl+C**: Interrupts current LLM generation; second Ctrl+C exits immediately. - **Configuration**: Stored as JSON in `/settings.json` (`$XDG_CONFIG_HOME/tinyharness/`, else `~/.config/tinyharness/`). Changed via CLI flags (`--url`, `--api-key`), slash commands (`/model`, `/apikey`, `/timeout`, `/retries`, `/contextlimit`, `/autoaccept`, `/settings`, …), or by editing the file. There is no interactive setup wizard; Sockudo credentials are settings-file-only. Persistent prompts are seeded from embedded defaults into `/prompts/`. All locations resolve through `paths::Paths`. - **Image attachments**: Base64 data URIs, used by multimodal models; set via `/image`. +- **Notebook reads**: `read` renders `.ipynb` files with binary outputs (`image/*`, `application/pdf`) replaced by placeholders and text values over 4,000 chars truncated — a single cell's base64 figure is otherwise hundreds of KB and can exceed the provider's input-length limit. `from`/`to` index the rendered text, not the file on disk. Invalid JSON falls back to per-line truncation. - **`async_command!` macro**: Registers commands that need `provider.lock().await`. - **`CommandResult` variants**: `SwitchSession`, `RenameSession`, `Init`, `SkillUse`, `SkillUnload` carry data back to the agent loop. - **`CommandContext`** holds shared mutable state: provider, mode, file context, session ID, skill registry, active skills, pending images, thinking toggle, compaction token usage, workspace context, `paths` (all filesystem locations; `prompts_dir()` derives from it), output writer, exit flag. diff --git a/docs/tools-reference.md b/docs/tools-reference.md index 6558327..574adef 100644 --- a/docs/tools-reference.md +++ b/docs/tools-reference.md @@ -47,6 +47,8 @@ Lists the contents of a single directory. Returns newline-separated file and dir Reads file content with optional line ranges. For image files (png, jpg, webp, gif, bmp), returns a description and the image data is automatically loaded for the model to view. +Jupyter notebooks (`.ipynb`) are rendered as JSON with embedded binary outputs (every `image/*` and `application/pdf` payload) replaced by a `[stripped … base64 chars]` placeholder, and remaining output values longer than 4,000 characters truncated. Cell sources, stream text, `text/plain` results and execution counts are preserved. A single notebook cell can otherwise embed several hundred KB of base64 on one line and push the request past the provider's input limit. If the file is not valid JSON, oversized lines are truncated instead. `from`/`to` index the **rendered** text, not the file on disk, so line numbers stay stable across reads. + | Parameter | Required | Description | |-----------|----------|-------------| | `path` | Yes | The absolute path to the file | diff --git a/tinyharness-lib/src/tools/read.rs b/tinyharness-lib/src/tools/read.rs index 7bb8698..5048cac 100644 --- a/tinyharness-lib/src/tools/read.rs +++ b/tinyharness-lib/src/tools/read.rs @@ -5,11 +5,21 @@ use std::io::{BufRead, BufReader}; use crate::extract_args; use crate::image::SUPPORTED_MIME_TYPES; use crate::tools::tool::{BoxFuture, ToolCategory, build_string_params_schema, make_tool}; +use serde_json::Value; + +/// Maximum characters kept from a single notebook output value (e.g. a huge +/// `text/plain` dataframe dump). Larger values are truncated with a marker. +const MAX_NOTEBOOK_VALUE_CHARS: usize = 4_000; + +/// Hard cap for a single line in the notebook fallback scrubber, used when the +/// `.ipynb` file is not valid JSON (or is minified onto one line) and the +/// structure-aware sanitizer cannot run. +const MAX_SCRUBBED_LINE_CHARS: usize = 2_000; pub fn read_tool_entry() -> crate::tools::tool::Tool { make_tool( "read", - "Read file content. Returns the entire file or a specific line range if from/to are provided. For image files (png, jpg, webp, gif, bmp), returns a description and the image data is automatically loaded for the model to view.", + "Read file content. Returns the entire file or a specific line range if from/to are provided. For image files (png, jpg, webp, gif, bmp), returns a description and the image data is automatically loaded for the model to view. For Jupyter notebooks (.ipynb), renders the JSON with base64 image outputs stripped and oversized text truncated — pass from/to to page through the rendered result.", ToolCategory::ReadOnly, build_string_params_schema( &[("path", "The absolute path to the file to read")], @@ -51,6 +61,14 @@ pub fn read_tool(args: HashMap) -> BoxFuture<'static, String> { let from = args.get("from").and_then(|f| f.parse::().ok()); let to = args.get("to").and_then(|t| t.parse::().ok()); + // Jupyter notebooks: embedded base64 image outputs routinely reach + // hundreds of kilobytes per cell and blow the model's input limit, so + // render a structure-aware, image-free view of the notebook instead of + // dumping the raw JSON. `from`/`to` page through the rendered text. + if is_notebook_file(&path) { + return read_notebook(&path, from, to); + } + match (from, to) { (Some(from), Some(to)) => read_partial(&path, from, to), _ => match fs::read_to_string(&path) { @@ -64,6 +82,167 @@ pub fn read_tool(args: HashMap) -> BoxFuture<'static, String> { }) } +/// Check if a path refers to a Jupyter notebook by extension. +pub fn is_notebook_file(path: &str) -> bool { + path.to_lowercase().ends_with(".ipynb") +} + +/// Read a Jupyter notebook as compact, model-safe JSON. +/// +/// Notebook cells routinely embed rendered figures as base64 payloads +/// (`"image/png": "iVBOR…"` — hundreds of kilobytes *on a single line*), which +/// can push a request past the provider's input-length limit. The JSON is +/// therefore parsed and sanitized before it is returned: +/// +/// * every binary payload (image/*, `application/pdf`) is replaced by a +/// short placeholder recording the mime type and the stripped size — this +/// covers both `outputs[].data` and markdown `attachments`; +/// * remaining text values longer than [`MAX_NOTEBOOK_VALUE_CHARS`] are +/// truncated with a marker, so a giant `text/plain` dump cannot recreate +/// the same problem; +/// * cell sources, stream text, `text/plain` results and execution counts are +/// preserved verbatim. +/// +/// If the file is not valid JSON (or is a notebook-shaped document the +/// sanitizer cannot walk), the raw text is returned with over-long lines +/// truncated instead, so the base64 blow-up is still neutralised. +/// +/// `from`/`to` page through the *rendered* text, not the on-disk file, so the +/// line numbers stay stable across reads. +fn read_notebook(path: &str, from: Option, to: Option) -> String { + let raw = match fs::read_to_string(path) { + Ok(content) => content, + Err(e) => return format!("Error reading file: {e}"), + }; + + let (rendered, mut notes) = match serde_json::from_str::(&raw) { + Ok(mut value) => { + let stripped = strip_binary_payloads(&mut value); + let truncated = truncate_oversized_strings(&mut value); + let text = serde_json::to_string_pretty(&value).unwrap_or_else(|_| raw.clone()); + let mut notes = Vec::new(); + if stripped > 0 { + notes.push(format!("{stripped} image outputs stripped")); + } + if truncated > 0 { + notes.push(format!("{truncated} oversized outputs truncated")); + } + (text, notes) + } + Err(_) => (scrub_huge_lines(&raw), vec!["long lines truncated".into()]), + }; + + let rendered = rendered.trim_end_matches('\n'); + match (from, to) { + (Some(from), Some(to)) => { + let selected: Vec<&str> = rendered.lines().skip(from).take(to).collect(); + if selected.is_empty() { + format!("Error: No lines to read in '{path}' at offset {from}") + } else { + format!( + "({} lines from line {from})\n{}", + selected.len(), + selected.join("\n") + ) + } + } + _ => { + let count = rendered.lines().count(); + if !notes.is_empty() { + notes.insert(0, format!("{count} lines")); + } + let meta = format!("({})", notes.join(", ")); + format!("{meta}\n{rendered}") + } + } +} + +/// Replace every binary (image / PDF) payload in a parsed notebook with a +/// placeholder. Returns the number of payloads stripped. +fn strip_binary_payloads(value: &mut Value) -> usize { + fn is_binary_key(key: &str) -> bool { + key.starts_with("image/") || key == "application/pdf" + } + + /// Base64 payloads are stored either as one string or as a list of + /// strings; both collapse to the same character count. + fn payload_len(value: &Value) -> usize { + match value { + Value::String(s) => s.chars().count(), + Value::Array(items) => items + .iter() + .map(|i| i.as_str().map(str::chars).map_or(0, Iterator::count)) + .sum(), + _ => 0, + } + } + + match value { + Value::Object(map) => { + let mut stripped = 0; + for (key, entry) in map.iter_mut() { + if is_binary_key(key) { + let size = payload_len(entry); + if size > 0 { + *entry = + Value::String(format!("[stripped {key} output — {size} base64 chars]")); + stripped += 1; + continue; + } + } + stripped += strip_binary_payloads(entry); + } + stripped + } + Value::Array(items) => items.iter_mut().map(strip_binary_payloads).sum(), + Value::Null | Value::Bool(_) | Value::Number(_) | Value::String(_) => 0, + } +} + +/// Truncate text values longer than [`MAX_NOTEBOOK_VALUE_CHARS`]. Returns the +/// number of values truncated. +fn truncate_oversized_strings(value: &mut Value) -> usize { + fn cap(s: &str) -> Option { + let total = s.chars().count(); + if total <= MAX_NOTEBOOK_VALUE_CHARS { + return None; + } + let head = s.floor_char_boundary(MAX_NOTEBOOK_VALUE_CHARS); + Some(format!("{}…[truncated, {total} chars total]", &s[..head])) + } + + match value { + Value::Object(map) => map.values_mut().map(truncate_oversized_strings).sum(), + Value::Array(items) => items.iter_mut().map(truncate_oversized_strings).sum(), + Value::String(s) => match cap(s) { + Some(replacement) => { + *s = replacement; + 1 + } + None => 0, + }, + Value::Null | Value::Bool(_) | Value::Number(_) => 0, + } +} + +/// Line-oriented fallback for documents that are not valid JSON: keep the +/// text but cut any line longer than [`MAX_SCRUBBED_LINE_CHARS`], which is +/// where embedded base64 lives in a notebook. +fn scrub_huge_lines(content: &str) -> String { + content + .lines() + .map(|line| { + if line.chars().count() > MAX_SCRUBBED_LINE_CHARS { + let head = line.floor_char_boundary(MAX_SCRUBBED_LINE_CHARS); + format!("{}…[line truncated]", &line[..head]) + } else { + line.to_string() + } + }) + .collect::>() + .join("\n") +} + /// Check if a path refers to an image file by extension. pub fn is_image_file(path: &str) -> bool { let lowercase = path.to_lowercase(); @@ -213,4 +392,246 @@ mod tests { assert_eq!(guess_image_mime("photo.jpeg"), Some("image/jpeg")); assert_eq!(guess_image_mime("main.rs"), None); } + + // ── notebook handling ──────────────────────────────────────────────── + + /// A base64 PNG payload of the requested length, as jupyter stores it. + fn fake_b64(len: usize) -> String { + "A".repeat(len) + } + + fn write_notebook(dir: &std::path::Path, name: &str, body: &str) -> std::path::PathBuf { + let path = dir.join(name); + std::fs::write(&path, body).unwrap(); + path + } + + fn read_nb(path: &std::path::Path) -> String { + read_notebook(&path.to_string_lossy(), None, None) + } + + #[test] + fn is_notebook_file_matches_extension() { + assert!(is_notebook_file("lab/main.ipynb")); + assert!(is_notebook_file("MAIN.IPYNB")); + assert!(!is_notebook_file("main.py")); + assert!(!is_notebook_file("nb.ipynb.bak")); + } + + #[test] + fn notebook_image_outputs_are_stripped_but_source_kept() { + let dir = tempfile::tempdir().unwrap(); + let big = fake_b64(100_000); + let nb = format!( + r#"{{ + "cells": [ + {{"cell_type": "code", "execution_count": 12, "metadata": {{}}, + "outputs": [ + {{"name": "stdout", "output_type": "stream", "text": ["loading done\n"]}}, + {{"output_type": "display_data", "metadata": {{}}, + "data": {{"image/png": "{big}", "text/plain": ["
"]}}}} + ], + "source": ["import cv2\n", "cv2.imshow('x', img)\n"]}} + ], + "metadata": {{}}, "nbformat": 4, "nbformat_minor": 5 +}}"# + ); + let path = write_notebook(dir.path(), "main.ipynb", &nb); + let out = read_nb(&path); + + // The base64 blob is gone, replaced by a size-recording placeholder. + assert!(!out.contains(&fake_b64(1000)), "base64 payload survived"); + assert!(out.contains("[stripped image/png output — 100000 base64 chars]")); + + // Everything the model actually needs survives. + assert!(out.contains("import cv2")); + assert!(out.contains("cv2.imshow")); + assert!(out.contains("loading done")); + assert!(out.contains("
")); + assert!(out.contains("\"execution_count\": 12")); + + // Meta line reports the line count first so the UI card still parses it. + let meta = out.lines().next().unwrap(); + assert!( + meta.starts_with('(') && meta.ends_with(')'), + "meta was: {meta}" + ); + assert!(meta.contains(" lines"), "meta was: {meta}"); + assert!( + meta.contains("1 image outputs stripped"), + "meta was: {meta}" + ); + } + + #[test] + fn notebook_image_payload_stored_as_string_list_is_stripped() { + let dir = tempfile::tempdir().unwrap(); + let nb = format!( + r#"{{"cells": [{{"cell_type": "code", "outputs": [{{"output_type": "display_data", + "data": {{"image/png": ["{}", "{}"]}}}}], "source": ["x=1"]}}], + "nbformat": 4, "nbformat_minor": 5}}"#, + fake_b64(600), + fake_b64(600) + ); + let path = write_notebook(dir.path(), "listed.ipynb", &nb); + let out = read_nb(&path); + assert!(out.contains("[stripped image/png output — 1200 base64 chars]")); + assert!(!out.contains(&fake_b64(600))); + } + + #[test] + fn notebook_markdown_attachments_are_stripped() { + let dir = tempfile::tempdir().unwrap(); + let nb = format!( + r#"{{"cells": [{{"cell_type": "markdown", "attachments": {{"image/jpeg": "{b}"}}, + "source": ["see figure"]}}], "nbformat": 4, "nbformat_minor": 5}}"#, + b = fake_b64(5000) + ); + let path = write_notebook(dir.path(), "attach.ipynb", &nb); + let out = read_nb(&path); + assert!(out.contains("[stripped image/jpeg output")); + assert!(out.contains("see figure")); + assert!(!out.contains(&fake_b64(5000))); + } + + #[test] + fn notebook_oversized_text_output_is_truncated() { + let dir = tempfile::tempdir().unwrap(); + let big = fake_b64(MAX_NOTEBOOK_VALUE_CHARS * 3); + let nb = format!( + r#"{{"cells": [{{"cell_type": "code", "outputs": [{{"output_type": "execute_result", + "data": {{"text/plain": "{big}"}}}}], "source": ["df"]}}], + "nbformat": 4, "nbformat_minor": 5}}"# + ); + let path = write_notebook(dir.path(), "text.ipynb", &nb); + let out = read_nb(&path); + assert!(out.contains(&format!( + "…[truncated, {} chars total]", + MAX_NOTEBOOK_VALUE_CHARS * 3 + ))); + assert!(out.contains("1 oversized outputs truncated")); + assert!(!out.contains(&big)); + assert!(out.len() < MAX_NOTEBOOK_VALUE_CHARS * 2); + } + + #[test] + fn notebook_without_big_outputs_is_left_alone() { + let dir = tempfile::tempdir().unwrap(); + let nb = r#"{"cells": [{"cell_type": "code", "execution_count": 1, + "outputs": [{"output_type": "stream", "name": "stdout", "text": ["hi\n"]}], + "source": ["print('hi')"]}], "metadata": {}, "nbformat": 4, "nbformat_minor": 5}"#; + let path = write_notebook(dir.path(), "clean.ipynb", nb); + let out = read_nb(&path); + assert!(out.contains("print('hi')")); + assert!(out.contains("hi\\n")); + // Nothing was rewritten, so no action notes are appended. + let meta = out.lines().next().unwrap(); + assert!(!meta.contains("stripped"), "meta was: {meta}"); + assert!(!meta.contains("truncated"), "meta was: {meta}"); + } + + #[test] + fn malformed_notebook_falls_back_to_line_scrubbing() { + let dir = tempfile::tempdir().unwrap(); + let body = format!( + "{{{{ not json at all\nsecond line ok\nlong line: {}\nfourth\n", + fake_b64(MAX_SCRUBBED_LINE_CHARS * 4) + ); + let path = write_notebook(dir.path(), "broken.ipynb", &body); + let out = read_nb(&path); + assert!(out.contains("second line ok")); + assert!(out.contains("[line truncated]")); + assert!(!out.contains(&fake_b64(MAX_SCRUBBED_LINE_CHARS * 2))); + assert!(out.lines().next().unwrap().contains("long lines truncated")); + } + + #[test] + fn notebook_from_to_pages_the_rendered_text() { + let dir = tempfile::tempdir().unwrap(); + let nb = format!( + r#"{{"cells": [{{"cell_type": "code", "outputs": [{{"output_type": "display_data", + "data": {{"image/png": "{b}"}}}}], "source": ["alpha\n", "beta"]}}], + "nbformat": 4, "nbformat_minor": 5}}"#, + b = fake_b64(20_000) + ); + let path = write_notebook(dir.path(), "paged.ipynb", &nb); + let full = read_notebook(&path.to_string_lossy(), None, None); + // `from`/`to` index the rendered body, never the meta line — the same + // contract `read_partial` honours for plain files. + let body: Vec<&str> = full.lines().skip(1).collect(); + assert!(body.len() > 4, "fixture should render several lines"); + + let window = read_notebook(&path.to_string_lossy(), Some(0), Some(3)); + assert!(window.starts_with("(3 lines from line 0)")); + assert_eq!(window.lines().skip(1).collect::>(), &body[..3]); + + let middle = read_notebook(&path.to_string_lossy(), Some(1), Some(2)); + assert_eq!(middle.lines().skip(1).collect::>(), &body[1..3]); + + let tail = read_notebook(&path.to_string_lossy(), Some(body.len() - 2), Some(5)); + assert_eq!( + tail.lines().skip(1).collect::>(), + &body[body.len() - 2..] + ); + assert!(tail.contains("\"nbformat_minor\": 5")); + + let empty = read_notebook(&path.to_string_lossy(), Some(body.len() + 10), Some(5)); + assert!(empty.starts_with("Error: No lines to read")); + } + + /// Regression for the bug that motivated this code path: a notebook whose + /// cells embed four ~380 KB figures must not produce a multi-megabyte tool + /// result (that request exceeded the provider's input limit). + #[test] + fn realistic_notebook_with_many_figures_stays_small() { + let dir = tempfile::tempdir().unwrap(); + let cells: Vec = (0..4) + .map(|i| { + format!( + r#"{{"cell_type": "code", "execution_count": {i}, "metadata": {{}}, + "outputs": [{{"output_type": "display_data", "metadata": {{}}, + "data": {{"image/png": "{b}", "text/plain": ["
"]}}}}], + "source": ["cell_{i}\n", "imshow(x)\n"]}}"#, + b = fake_b64(377_440) + ) + }) + .collect(); + let nb = format!( + r#"{{"cells": [{}], "metadata": {{}}, "nbformat": 4, "nbformat_minor": 5}}"#, + cells.join(",") + ); + assert!(nb.len() > 1_400_000, "fixture should mirror the real file"); + + let path = write_notebook(dir.path(), "big.ipynb", &nb); + let out = read_nb(&path); + assert!(out.len() < 4_000, "notebook render was {} chars", out.len()); + assert!(out.contains("cell_0") && out.contains("cell_3")); + assert!(out.contains("4 image outputs stripped")); + } + + #[test] + fn read_tool_dispatches_notebooks_without_image_marker() { + let dir = tempfile::tempdir().unwrap(); + let nb = format!( + r#"{{"cells": [{{"cell_type": "code", "outputs": [{{"data": {{"image/png": "{b}"}}}}], + "source": ["plot()"]}}], "nbformat": 4, "nbformat_minor": 5}}"#, + b = fake_b64(50_000) + ); + let path = write_notebook(dir.path(), "dispatch.ipynb", &nb); + + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + let out = runtime.block_on(read_tool(HashMap::from([( + "path".to_string(), + path.to_string_lossy().to_string(), + )]))); + + // A notebook is not an image: no `[IMAGE]` marker, so the agent layer + // must not try to attach it as a multimodal payload. + assert!(!out.starts_with("[IMAGE]")); + assert!(out.contains("plot()")); + assert!(out.contains("[stripped image/png output")); + } }