diff --git a/CHANGELOG.md b/CHANGELOG.md index 2615431..99ebb68 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,7 @@ All notable changes to this project will be documented in this file. ## Unreleased +- **Auto-diff honors `--max-tokens` like the standard renderer (#20 B14).** The auto-diff path used a bytes/4 cut over the finished document (mid-file, ignoring `--encoding`, omissions unreported). Files are now budgeted one by one with the real tokenizer: a file that would cross the budget is clamped at a `--truncate` boundary (and says so) or omitted, and the notice counts the omitted files. The budget search is shared (`markdown::fit_chunk_to_budget`) with the standard path. - **`--structure` lists imports and exports (#22 B15).** The help promised them but only counts were printed (and the internal lists held placeholders). Each file now gets `**Imports:**` / `**Exports:**` lists (Rust `use`/`pub use`, Python, JS/TS `import`/`export`, Go, Java, C/C++ `#include`), capped at 20 entries. Top-level declarations only. - **Rust `--signatures` group methods under their `impl`/`trait` (#22 B16).** Methods used to be listed flat after `impl A`, leaving several `pub fn new` ambiguous. Each impl/trait is now one entry: `impl A { pub fn new() -> A … }`. Library: members are part of the impl/trait signature's `full_signature` and no longer separate `Signature`s. - **Python `--signatures`: no nested functions, overloads keep their return types (#22 B16).** Functions and classes defined inside a function are no longer listed; `@overload def f(...) -> T: ...` used to lose the decorator and `-> T` (the one-line body defeated the slicer) and now keeps both, with decorator lines no longer carrying stray indentation. diff --git a/src/lib.rs b/src/lib.rs index f73d68d..0b69b99 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -707,28 +707,12 @@ pub fn run_with_args(args: Args, config: Config, prompter: &impl Prompter) -> io &skipped, )?; - // Enforce max_tokens budget (same ~4 bytes/token heuristic as parallel path) + // Files are budgeted one by one inside the renderer. This is only the + // safety net for a header/diff/tree that alone exceeds the budget. if let Some(max_tokens) = final_args.max_tokens { - let max_bytes = max_tokens.saturating_mul(4); - if final_doc.len() > max_bytes { - // Truncate at a valid UTF-8 boundary - let mut truncate_at = max_bytes; - while truncate_at > 0 && !final_doc.is_char_boundary(truncate_at) { - truncate_at -= 1; - } - final_doc.truncate(truncate_at); - - // Close any open code fence so the truncation notice is not - // swallowed by the block. The closer matches the opening - // fence's length — a fixed ``` would not close a longer fence - // chosen because the file itself contains ```. - fences::close_unmatched_backtick_fence(&mut final_doc); - - final_doc.push_str("\n---\n\n"); - final_doc.push_str(&format!( - "_Output truncated: exceeded {} token budget (estimated)._\n", - max_tokens - )); + let encoding = final_args.encoding.parse::().unwrap_or_default(); + if estimate_tokens(encoding, &final_doc) > max_tokens { + final_doc = truncate_to_token_budget(&final_doc, max_tokens, encoding); } } @@ -1183,64 +1167,147 @@ fn generate_markdown_with_diff( if !diff_config.diff_only { output.push_str("## File Contents\n\n"); + // Same per-file, tokenizer-based budget as the standard renderer: a file + // that would cross `--max-tokens` is clamped to what is left (or omitted) + // and the rest are reported, instead of cutting the document mid-file. + let budget = args.max_tokens.map(|max| { + let encoding = args.encoding.parse::().unwrap_or_default(); + (max, encoding, estimate_tokens(encoding, &output)) + }); + let mut tokens_used = budget.map_or(0, |(_, _, used)| used); + // Iterate in relevance order (from sorted_paths) instead of // BTreeMap's alphabetical order — preserves file_relevance_category ordering. - for path in sorted_paths { - if let Some(file_state) = current_state.files.get(path) { - output.push_str(&format!( - "### File: {}\n\n", - fences::inline_code(&path.display().to_string()) - )); - // Same opt-in as the standard renderer. Off by default so an - // mtime-only change does not rewrite the document. The cache - // compares content hashes (file bytes), not mtime. - if args.file_metadata { - output.push_str(&format!("- Size: {} bytes\n", file_state.size)); - let modified = chrono::DateTime::::from(file_state.modified); - let stamp = modified.format("%Y-%m-%d %H:%M:%S UTC"); - output.push_str(&format!("- Modified: {stamp}\n\n")); - } - - // Determine language from file extension (canonical map — - // same fence language as the main rendering path) - let extension = path.extension().and_then(|s| s.to_str()).unwrap_or("text"); - let language = crate::languages::language_for_extension(extension); - - // When --signatures is active, only suppress content for supported code files - let signatures_only = - ts_config.signatures && crate::tree_sitter::is_supported_extension(extension); - - if !signatures_only { - // `output` is a String (`fmt::Write`); the shared writer speaks `io::Write`. - let mut rendered = Vec::new(); - markdown::write_text_content( - &mut rendered, - &file_state.content, - language, - args.line_numbers, - )?; - let rendered = std::str::from_utf8(&rendered) - .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?; - output.push_str(rendered); + let present: Vec<(&PathBuf, &state::FileState)> = sorted_paths + .iter() + .filter_map(|path| current_state.files.get(path).map(|state| (path, state))) + .collect(); + for (idx, (path, file_state)) in present.iter().enumerate() { + let section = diff_file_section(path, file_state, args, ts_config, None)?; + let section_tokens = + budget.map_or(0, |(_, encoding, _)| estimate_tokens(encoding, §ion)); + let fits = budget.is_none_or(|(max, _, _)| tokens_used + section_tokens <= max); + if fits { + tokens_used += section_tokens; + output.push_str(§ion); + } else if let Some((max, encoding, _)) = budget { + // Crossing file: clamp it into the remaining allowance, else omit it. + let remaining_tokens = max.saturating_sub(tokens_used); + let fitted = markdown::fit_chunk_to_budget(remaining_tokens, encoding, |limit| { + diff_file_section(path, file_state, args, ts_config, limit).ok() + }); + let omitted = present.len() - idx - usize::from(fitted.is_some()); + if let Some((chunk, _)) = fitted { + output.push_str(&chunk); } + output.push_str("---\n\n"); + output.push_str(&budget_notice(max, omitted)); + break; + } + } + } - // Tree-sitter enrichment (same as standard path) - let mut enrichment_buf = Vec::new(); - markdown::write_tree_sitter_enrichment( - &mut enrichment_buf, - &file_state.content, - extension, - ts_config, - )?; - if !enrichment_buf.is_empty() { - output.push_str(&String::from_utf8_lossy(&enrichment_buf)); - } + Ok(output) +} - output.push('\n'); - } +/// Cut `doc` so that, with its fence closed and the notice appended, the real +/// token count is within `max_tokens`. Starts from the bytes/4 guess and +/// shrinks until the tokenizer agrees (or nothing is left but the notice). +fn truncate_to_token_budget(doc: &str, max_tokens: usize, encoding: Encoding) -> String { + let mut cut = doc.len().min(max_tokens.saturating_mul(4)); + loop { + while cut > 0 && !doc.is_char_boundary(cut) { + cut -= 1; + } + let mut out = doc[..cut].to_string(); + // Close any open code fence so the notice is not swallowed by the + // block. The closer matches the opening fence's length: a fixed ``` + // would not close a longer fence chosen because the file contains ```. + fences::close_unmatched_backtick_fence(&mut out); + out.push_str("\n---\n\n"); + out.push_str(&format!( + "_Output truncated: exceeded {max_tokens} token budget._\n" + )); + if cut == 0 || estimate_tokens(encoding, &out) <= max_tokens { + return out; } + cut -= (cut / 10).max(1); } +} + +/// The notice that ends a document whose budget ran out. +fn budget_notice(max: usize, omitted: usize) -> String { + format!("_⚠️ Token budget ({max}) reached. {omitted} remaining files omitted._\n\n") +} +/// One `### File:` section of the auto-diff document. `content_limit` clamps +/// the file's content to that many bytes (at a boundary chosen by `--truncate`) +/// and says so, like the standard renderer does for a file that crosses the +/// token budget. +fn diff_file_section( + path: &Path, + file_state: &state::FileState, + args: &Args, + ts_config: &markdown::TreeSitterConfig, + content_limit: Option, +) -> io::Result { + let mut output = String::new(); + let heading = fences::inline_code(&path.display().to_string()); + output.push_str(&format!("### File: {heading}\n\n")); + // Same opt-in as the standard renderer. Off by default so an + // mtime-only change does not rewrite the document. The cache + // compares content hashes (file bytes), not mtime. + if args.file_metadata { + let size = file_state.size; + let modified = chrono::DateTime::::from(file_state.modified); + let stamp = modified.format("%Y-%m-%d %H:%M:%S UTC"); + output.push_str(&format!("- Size: {size} bytes\n- Modified: {stamp}\n\n")); + } + + // Determine language from file extension (canonical map — + // same fence language as the main rendering path) + let extension = path.extension().and_then(|s| s.to_str()).unwrap_or("text"); + let language = crate::languages::language_for_extension(extension); + + let full = file_state.content.as_str(); + let (content, truncated) = match content_limit { + Some(limit) if full.len() > limit => { + let cut = + markdown::content_cut_point(full, limit, extension, ts_config.truncate.as_str()); + (&full[..cut], cut < full.len()) + } + _ => (full, false), + }; + + // When --signatures is active, only suppress content for supported code files + let signatures_only = + ts_config.signatures && crate::tree_sitter::is_supported_extension(extension); + + if !signatures_only { + // `output` is a String (`fmt::Write`); the shared writer speaks `io::Write`. + let mut rendered = Vec::new(); + markdown::write_text_content(&mut rendered, content, language, args.line_numbers)?; + let rendered = std::str::from_utf8(&rendered) + .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?; + output.push_str(rendered); + } + // The marker is about the file, not about whether its body is shown: with + // --signatures the signatures below are computed from the clamped content. + if truncated { + let (shown, total, mode) = (content.len(), full.len(), &ts_config.truncate); + output.push_str(&format!( + "\n_⚠️ File content truncated: {shown} of {total} bytes (mode: {mode})._\n\n" + )); + } + + // Tree-sitter enrichment (same as standard path), over what is displayed + let mut enrichment_buf = Vec::new(); + markdown::write_tree_sitter_enrichment(&mut enrichment_buf, content, extension, ts_config)?; + if !enrichment_buf.is_empty() { + output.push_str(&String::from_utf8_lossy(&enrichment_buf)); + } + + output.push('\n'); Ok(output) } diff --git a/src/markdown.rs b/src/markdown.rs index 60ca0f9..063ec9f 100644 --- a/src/markdown.rs +++ b/src/markdown.rs @@ -814,15 +814,17 @@ pub fn process_file_with_content_limit( if !signatures_only { write_text_content(output, content, language, line_numbers)?; - if truncated { - writeln!( - output, - "\n_⚠️ File content truncated: {} of {} bytes (mode: {})._\n", - content.len(), - original_len, - ts_config.truncate - )?; - } + } + // About the file, not about whether its body is shown: with + // --signatures the block below is computed from the clamped content. + if truncated { + writeln!( + output, + "\n_⚠️ File content truncated: {} of {} bytes (mode: {})._\n", + content.len(), + original_len, + ts_config.truncate + )?; } // Tree-sitter enrichment: signatures and/or structure @@ -939,27 +941,46 @@ pub fn render_truncated_to_fit( token_encoding: TokenEncoding, remaining_tokens: usize, ) -> Option { + let (chunk, tokens) = fit_chunk_to_budget(remaining_tokens, token_encoding, |limit| { + let mut probe: Vec = Vec::new(); + process_file_with_content_limit( + base_path, + file_path, + &mut probe, + line_numbers, + encoding_strategy, + ts_config, + limit, + ) + .ok()?; + Some(String::from_utf8_lossy(&probe).into_owned()) + })?; + *buf = chunk.into_bytes(); + Some(tokens) +} + +/// The budget search shared by every renderer: find a rendering of one file +/// section whose real token count fits `remaining_tokens`. +/// +/// `render(limit)` returns the section with its content clamped to `limit` +/// bytes (`None` = no clamp) or `None` if it cannot be rendered. The fixed +/// header cost is measured with a zero-byte probe, the content allowance is +/// what is left at ~4 bytes/token, and it shrinks proportionally (with +/// margin) while the tokenizer says the chunk overshoots. Gives up (omit the +/// file) when even an empty section does not fit. +pub(crate) fn fit_chunk_to_budget( + remaining_tokens: usize, + token_encoding: TokenEncoding, + mut render: impl FnMut(Option) -> Option, +) -> Option<(String, usize)> { if remaining_tokens == 0 { return None; } // 1. Measure the fixed section cost (header + empty fence + marker) by // rendering with a zero-byte content allowance. - let mut base_probe: Vec = Vec::new(); - if process_file_with_content_limit( - base_path, - file_path, - &mut base_probe, - line_numbers, - encoding_strategy, - ts_config, - Some(0), - ) - .is_err() - { - return None; - } - let base_tokens = estimate_tokens(token_encoding, &String::from_utf8_lossy(&base_probe)); + let base_probe = render(Some(0))?; + let base_tokens = estimate_tokens(token_encoding, &base_probe); // Even an empty section would overdraw — omit the file instead. if base_tokens >= remaining_tokens { @@ -973,25 +994,10 @@ pub fn render_truncated_to_fit( // tokenizer; shrink proportionally on overshoot. Bounded attempts — // if it still doesn't fit, omission is safer than an over-budget emit. for _ in 0..5 { - let mut probe: Vec = Vec::new(); - if process_file_with_content_limit( - base_path, - file_path, - &mut probe, - line_numbers, - encoding_strategy, - ts_config, - Some(byte_allowance), - ) - .is_err() - { - return None; - } - - let chunk_tokens = estimate_tokens(token_encoding, &String::from_utf8_lossy(&probe)); + let probe = render(Some(byte_allowance))?; + let chunk_tokens = estimate_tokens(token_encoding, &probe); if chunk_tokens <= remaining_tokens { - *buf = probe; - return Some(chunk_tokens); + return Some((probe, chunk_tokens)); } if byte_allowance == 0 { diff --git a/tests/test_autodiff_budget.rs b/tests/test_autodiff_budget.rs new file mode 100644 index 0000000..ecc249e --- /dev/null +++ b/tests/test_autodiff_budget.rs @@ -0,0 +1,259 @@ +//! `--max-tokens` with auto-diff budgets file by file with the real tokenizer, +//! like the standard renderer (#20 B14): no cut mid-fence, an overrunning file +//! is clamped and says so, omitted files are counted. + +use context_builder::cli::Args; +use context_builder::config::Config; +use context_builder::token_count::{Encoding, estimate_tokens}; +use context_builder::{Prompter, run_with_args}; +use serial_test::serial; +use std::fs; +use std::path::{Path, PathBuf}; +use std::process::Command; +use tempfile::tempdir; + +/// The auto-diff cache lives under the current directory, so each test runs +/// inside its own project (and serially). +struct CwdGuard(PathBuf); + +impl CwdGuard { + fn enter(path: &Path) -> Self { + let prev = std::env::current_dir().unwrap(); + std::env::set_current_dir(path).unwrap(); + Self(prev) + } +} + +impl Drop for CwdGuard { + fn drop(&mut self) { + let _ = std::env::set_current_dir(&self.0); + } +} + +struct Yes; + +impl Prompter for Yes { + fn confirm_overwrite(&self, _file_path: &str) -> std::io::Result { + Ok(true) + } +} + +fn args(input: &Path, output: &Path, max_tokens: Option, encoding: &str) -> Args { + Args { + input: input.to_string_lossy().to_string(), + output: output.to_string_lossy().to_string(), + filter: vec![], + ignore: vec![], + line_numbers: false, + preview: false, + token_count: false, + yes: true, + diff_only: false, + clear_cache: false, + encoding: encoding.to_string(), + init: false, + max_tokens, + signatures: false, + structure: false, + truncate: "smart".to_string(), + visibility: "all".to_string(), + file_metadata: false, + max_file_size: "256K".to_string(), + hidden: false, + include_secrets: false, + include_lockfiles: false, + } +} + +fn project() -> tempfile::TempDir { + let dir = tempdir().unwrap(); + for i in 1..=4 { + let mut body = format!("fn f{i}() {{\n"); + for j in 0..300 { + body.push_str(&format!(" let v{i}_{j} = {j};\n")); + } + body.push_str("}\n"); + fs::write(dir.path().join(format!("f{i}.rs")), body).unwrap(); + } + dir +} + +fn auto_diff() -> Config { + Config { + auto_diff: Some(true), + ..Default::default() + } +} + +fn fence_lines(doc: &str) -> usize { + doc.lines().filter(|l| l.starts_with("```")).count() +} + +#[test] +#[serial] +fn budget_is_met_by_the_real_tokenizer_and_fences_stay_closed() { + let dir = project(); + let _cwd = CwdGuard::enter(dir.path()); + let out = dir.path().join("out.md"); + run_with_args( + args(dir.path(), &out, Some(3000), "o200k_base"), + auto_diff(), + &Yes, + ) + .unwrap(); + let doc = fs::read_to_string(&out).unwrap(); + + let tokens = estimate_tokens(Encoding::O200kBase, &doc); + assert!(tokens <= 3000, "{tokens} tokens exceed the 3000 budget"); + assert!(tokens > 1500, "budget should be used, got {tokens}"); + assert_eq!(fence_lines(&doc) % 2, 0, "unbalanced fences:\n{doc}"); + assert!(doc.contains("Token budget (3000) reached"), "{doc}"); +} + +#[test] +#[serial] +fn crossing_file_is_clamped_and_the_rest_counted() { + let dir = project(); + let _cwd = CwdGuard::enter(dir.path()); + let out = dir.path().join("out.md"); + run_with_args( + args(dir.path(), &out, Some(5000), "o200k_base"), + auto_diff(), + &Yes, + ) + .unwrap(); + let doc = fs::read_to_string(&out).unwrap(); + + assert!(doc.contains("### File: `f1.rs`"), "first file kept"); + assert!( + doc.contains("File content truncated"), + "clamped file says so:\n{doc}" + ); + // f1 fits, f2 is clamped; f3 and f4 are omitted and reported as such. + assert!(doc.contains("2 remaining files omitted"), "{doc}"); + assert!(!doc.contains("### File: `f3.rs`") && !doc.contains("### File: `f4.rs`")); +} + +#[test] +#[serial] +fn a_large_budget_changes_nothing() { + // Separate projects: a second run in one would diff against the first. + let (a, b) = (project(), project()); + let full = a.path().join("full.md"); + let capped = b.path().join("capped.md"); + run_with_args(args(a.path(), &full, None, "o200k_base"), auto_diff(), &Yes).unwrap(); + run_with_args( + args(b.path(), &capped, Some(1_000_000), "o200k_base"), + auto_diff(), + &Yes, + ) + .unwrap(); + let strip = |p: &Path| { + fs::read_to_string(p) + .unwrap() + .lines() + .filter(|l| { + !l.starts_with("**Generated:**") + && !l.starts_with("**Project:**") + && !l.starts_with("Content hash") + }) + .collect::>() + .join("\n") + }; + assert_eq!(strip(&full), strip(&capped)); +} + +#[test] +#[serial] +fn encoding_is_honored_for_the_budget() { + let dir = project(); + let _cwd = CwdGuard::enter(dir.path()); + let out = dir.path().join("out.md"); + run_with_args( + args(dir.path(), &out, Some(3000), "cl100k_base"), + auto_diff(), + &Yes, + ) + .unwrap(); + let doc = fs::read_to_string(&out).unwrap(); + let tokens = estimate_tokens(Encoding::Cl100kBase, &doc); + assert!(tokens <= 3000, "{tokens} cl100k tokens exceed the budget"); +} + +#[test] +#[serial] +fn header_larger_than_the_budget_still_ends_cleanly() { + let dir = project(); + let _cwd = CwdGuard::enter(dir.path()); + let out = dir.path().join("out.md"); + run_with_args( + args(dir.path(), &out, Some(40), "o200k_base"), + auto_diff(), + &Yes, + ) + .unwrap(); + let doc = fs::read_to_string(&out).unwrap(); + assert!( + doc.contains("Output truncated: exceeded 40 token budget"), + "{doc}" + ); + assert_eq!(fence_lines(&doc) % 2, 0); +} + +#[test] +fn binary_with_auto_diff_config_respects_the_budget() { + let dir = project(); + fs::write( + dir.path().join("context-builder.toml"), + "auto_diff = true\n", + ) + .unwrap(); + let out = Command::new(env!("CARGO_BIN_EXE_context-builder")) + .args(["-d", ".", "-y", "-o", "-", "--max-tokens", "3000"]) + .current_dir(dir.path()) + .output() + .unwrap(); + assert!(out.status.success()); + let doc = String::from_utf8_lossy(&out.stdout); + let tokens = estimate_tokens(Encoding::O200kBase, &doc); + assert!(tokens <= 3000, "{tokens}"); + assert!(doc.contains("remaining files omitted"), "{doc}"); +} + +#[test] +#[serial] +fn file_metadata_is_rendered_in_budgeted_sections() { + let dir = project(); + let _cwd = CwdGuard::enter(dir.path()); + let out = dir.path().join("out.md"); + let mut a = args(dir.path(), &out, Some(100_000), "o200k_base"); + a.file_metadata = true; + run_with_args(a, auto_diff(), &Yes).unwrap(); + let doc = fs::read_to_string(&out).unwrap(); + assert!(doc.contains("- Size: "), "{doc}"); + assert!(doc.contains("- Modified: "), "{doc}"); + assert!(!doc.contains("Token budget"), "plenty of room, nothing cut"); +} + +#[test] +#[serial] +#[cfg(feature = "tree-sitter-rust")] +fn truncation_marker_is_kept_when_signatures_replace_the_body() { + let dir = tempdir().unwrap(); + let mut body = String::new(); + for j in 0..300 { + body.push_str(&format!( + "fn function_number_{j}(a: u32) -> u32 {{ a + {j} }}\n" + )); + } + fs::write(dir.path().join("many.rs"), body).unwrap(); + let _cwd = CwdGuard::enter(dir.path()); + let out = dir.path().join("out.md"); + let mut a = args(dir.path(), &out, Some(600), "o200k_base"); + a.signatures = true; + run_with_args(a, auto_diff(), &Yes).unwrap(); + let doc = fs::read_to_string(&out).unwrap(); + assert!(doc.contains("File content truncated"), "{doc}"); + assert!(doc.contains("**Signatures:**"), "{doc}"); + assert!(estimate_tokens(Encoding::O200kBase, &doc) <= 600); +}