diff --git a/crates/csp/src/bin/csp/main.rs b/crates/csp/src/bin/csp/main.rs index f61c9ae..1934708 100644 --- a/crates/csp/src/bin/csp/main.rs +++ b/crates/csp/src/bin/csp/main.rs @@ -14,8 +14,10 @@ use csp::indexing::cache::{clear_index_cache, CacheLocation}; use csp::indexing::index::{ load_or_build_index, CspIndex, LoadOptions, LoadOrBuildOptions, QueryOptions, }; -use csp::stats::{clear_savings, default_stats_file, format_savings_report, now_secs}; -use csp::types::ContentType; +use csp::stats::{ + clear_savings, default_stats_file, format_savings_report, now_secs, save_search_stats, +}; +use csp::types::{CallType, ContentType}; use csp::utils::{format_results, is_git_url, resolve_chunk, resolve_snippet_lines}; #[derive(Parser)] @@ -218,12 +220,14 @@ fn load_index( } } -/// JSON output for `search` (pure — testable without stdout capture). +/// JSON output for `search`. `stats_file` records token-savings telemetry when +/// `Some`; tests pass `None` to stay off the real `~/.csp` file. fn search_output( index: &CspIndex, query: &str, top_k: usize, max_snippet_lines: Option, + stats_file: Option<&Path>, ) -> String { let results = index.search( query, @@ -232,6 +236,15 @@ fn search_output( ..Default::default() }, ); + if let Some(stats_file) = stats_file { + save_search_stats( + stats_file, + &results, + CallType::Search, + &index.file_sizes, + max_snippet_lines, + ); + } let out = if results.is_empty() { serde_json::json!({ "error": "No results found." }) } else { @@ -247,6 +260,7 @@ fn find_related_output( line: &str, top_k: usize, max_snippet_lines: Option, + stats_file: Option<&Path>, ) -> Result { let Ok(line_num) = line.parse::() else { return Err(format!("line must be an integer, got: {line}")); @@ -268,6 +282,15 @@ fn find_related_output( ..Default::default() }, ); + if let Some(stats_file) = stats_file { + save_search_stats( + stats_file, + &related, + CallType::FindRelated, + &index.file_sizes, + max_snippet_lines, + ); + } let out = if related.is_empty() { serde_json::json!({ "error": format!("No related chunks found for {file}:{line_num}.") }) } else { @@ -350,6 +373,13 @@ fn run() -> ExitCode { /// `ExitCode` — so the dispatch logic is unit-testable without going through /// `Cli::parse` (which reads argv) or an opaque, non-comparable `ExitCode`. fn dispatch(command: Command) -> u8 { + dispatch_with_stats(command, &default_stats_file()) +} + +/// [`dispatch`] with the token-savings file injected so tests can redirect the +/// telemetry that `search` / `find-related` append (keeping the real `~/.csp` +/// untouched). +fn dispatch_with_stats(command: Command, stats_file: &Path) -> u8 { match command { Command::Init { agent, force } => { let agent = agent.unwrap_or(Agent::Claude); @@ -420,6 +450,7 @@ fn dispatch(command: Command) -> u8 { &query, top_k.unwrap_or(5), resolve_snippet_lines(max_snippet_lines), + Some(stats_file), ) ); EXIT_SUCCESS @@ -459,6 +490,7 @@ fn dispatch(command: Command) -> u8 { &line, top_k.unwrap_or(5), resolve_snippet_lines(max_snippet_lines), + Some(stats_file), ) { Ok(out) => { println!("{out}"); @@ -553,7 +585,7 @@ mod tests { fn search_output_shapes_results() { let dir = build_index_dir(); let idx = CspIndex::from_path(dir.path(), &LoadOptions::default()).unwrap(); - let out = search_output(&idx, "greet", 5, None); + let out = search_output(&idx, "greet", 5, None, None); let value: serde_json::Value = serde_json::from_str(&out).unwrap(); assert!(value.get("results").is_some() || value.get("error").is_some()); if let Some(results) = value.get("results").and_then(|r| r.as_array()) { @@ -572,7 +604,7 @@ mod tests { fn search_output_caps_snippet_lines() { let dir = build_index_dir(); let idx = CspIndex::from_path(dir.path(), &LoadOptions::default()).unwrap(); - let out = search_output(&idx, "greet", 5, Some(0)); + let out = search_output(&idx, "greet", 5, Some(0), None); let value: serde_json::Value = serde_json::from_str(&out).unwrap(); if let Some(first) = value["results"].as_array().and_then(|r| r.first()) { assert!(first.get("content").is_none()); @@ -580,11 +612,29 @@ mod tests { } } + #[test] + fn search_output_records_savings_when_stats_file_given() { + let dir = build_index_dir(); + let idx = CspIndex::from_path(dir.path(), &LoadOptions::default()).unwrap(); + // file_sizes is captured at build time from the source tree. + assert!(!idx.file_sizes.is_empty()); + + let stats = tempdir().unwrap(); + let stats_file = stats.path().join("savings.jsonl"); + let _ = search_output(&idx, "greet", 5, None, Some(&stats_file)); + + let content = std::fs::read_to_string(&stats_file).unwrap(); + let lines: Vec<&str> = content.lines().filter(|l| !l.is_empty()).collect(); + assert_eq!(lines.len(), 1); + assert!(lines[0].contains("\"call\":\"search\"")); + assert!(lines[0].contains("file_chars")); + } + #[test] fn find_related_rejects_non_integer_line() { let dir = build_index_dir(); let idx = CspIndex::from_path(dir.path(), &LoadOptions::default()).unwrap(); - let err = find_related_output(&idx, "sample.ts", "abc", 5, None).unwrap_err(); + let err = find_related_output(&idx, "sample.ts", "abc", 5, None, None).unwrap_err(); assert!(err.contains("line must be an integer")); } @@ -592,7 +642,7 @@ mod tests { fn find_related_no_chunk_at_location() { let dir = build_index_dir(); let idx = CspIndex::from_path(dir.path(), &LoadOptions::default()).unwrap(); - let err = find_related_output(&idx, "nope.ts", "1", 5, None).unwrap_err(); + let err = find_related_output(&idx, "nope.ts", "1", 5, None, None).unwrap_err(); assert!(err.contains("No chunk found")); } @@ -648,49 +698,64 @@ mod tests { #[test] fn dispatch_index_then_search_and_find_related() { // Keep everything on an explicit --index path so the test never writes - // to the global ~/.csp auto-cache. + // to the global ~/.csp auto-cache, and redirect savings telemetry to a + // temp file so it never touches the real ~/.csp/savings.jsonl. let src = build_index_dir(); let out = tempdir().unwrap(); + let stats_file = out.path().join("savings.jsonl"); assert_eq!(index_to(out.path(), src.path()), EXIT_SUCCESS); let idx_path = out.path().to_string_lossy().into_owned(); - let search = dispatch(Command::Search { - query: "greet".to_string(), - path: None, - top_k: Some(5), - max_snippet_lines: None, - content: vec![], - index: Some(idx_path.clone()), - git_ref: None, - }); + let search = dispatch_with_stats( + Command::Search { + query: "greet".to_string(), + path: None, + top_k: Some(5), + max_snippet_lines: None, + content: vec![], + index: Some(idx_path.clone()), + git_ref: None, + }, + &stats_file, + ); assert_eq!(search, EXIT_SUCCESS); // sample.ts:1 has an indexable chunk → find-related succeeds. - let related = dispatch(Command::FindRelated { - file: "sample.ts".to_string(), - line: "1".to_string(), - path: None, - top_k: Some(5), - max_snippet_lines: None, - content: vec![], - index: Some(idx_path.clone()), - git_ref: None, - }); + let related = dispatch_with_stats( + Command::FindRelated { + file: "sample.ts".to_string(), + line: "1".to_string(), + path: None, + top_k: Some(5), + max_snippet_lines: None, + content: vec![], + index: Some(idx_path.clone()), + git_ref: None, + }, + &stats_file, + ); assert_eq!(related, EXIT_SUCCESS); // A non-integer line is a caller error → failure exit. - let bad = dispatch(Command::FindRelated { - file: "sample.ts".to_string(), - line: "abc".to_string(), - path: None, - top_k: Some(5), - max_snippet_lines: None, - content: vec![], - index: Some(idx_path), - git_ref: None, - }); + let bad = dispatch_with_stats( + Command::FindRelated { + file: "sample.ts".to_string(), + line: "abc".to_string(), + path: None, + top_k: Some(5), + max_snippet_lines: None, + content: vec![], + index: Some(idx_path), + git_ref: None, + }, + &stats_file, + ); assert_eq!(bad, EXIT_FAILURE); + + // The two successful calls appended one savings record each. + let recorded = std::fs::read_to_string(&stats_file).unwrap(); + assert_eq!(recorded.lines().filter(|l| !l.is_empty()).count(), 2); } #[test] diff --git a/crates/csp/src/bin/csp/mcp_server.rs b/crates/csp/src/bin/csp/mcp_server.rs index f920707..44d8c61 100644 --- a/crates/csp/src/bin/csp/mcp_server.rs +++ b/crates/csp/src/bin/csp/mcp_server.rs @@ -16,6 +16,7 @@ use rmcp::{tool, tool_handler, tool_router, ErrorData as McpError, ServerHandler use tokio::sync::Mutex; use csp::mcp::{find_related_tool, search_tool, IndexCache, SERVER_INSTRUCTIONS}; +use csp::stats::default_stats_file; use csp::types::ContentType; use csp::utils::resolve_snippet_lines; @@ -65,6 +66,9 @@ pub struct CspMcpServer { cache: Arc>, default_source: Option, default_ref: Option, + /// Where token-savings telemetry is appended; `None` disables recording + /// (used by tests so they don't touch the real `~/.csp/savings.jsonl`). + stats_file: Option, tool_router: ToolRouter, } @@ -79,6 +83,7 @@ impl CspMcpServer { cache: Arc::new(Mutex::new(IndexCache::new(content))), default_source, default_ref, + stats_file: Some(default_stats_file()), tool_router: Self::tool_router(), } } @@ -99,6 +104,7 @@ impl CspMcpServer { p.repo.as_deref(), p.top_k.unwrap_or(5) as usize, resolve_snippet_lines(p.max_snippet_lines), + self.stats_file.as_deref(), ); Ok(CallToolResult::success(vec![Content::text(out)])) } @@ -120,6 +126,7 @@ impl CspMcpServer { p.repo.as_deref(), p.top_k.unwrap_or(5) as usize, resolve_snippet_lines(p.max_snippet_lines), + self.stats_file.as_deref(), ); Ok(CallToolResult::success(vec![Content::text(out)])) } @@ -231,11 +238,13 @@ mod tests { #[tokio::test] async fn search_tool_call_returns_json_payload() { let dir = sample_source(); - let server = CspMcpServer::new( + let mut server = CspMcpServer::new( Some(dir.path().to_string_lossy().into_owned()), None, vec![ContentType::Code], ); + // Don't append telemetry to the developer's real ~/.csp during tests. + server.stats_file = None; let result = server .search(Parameters(SearchParams { query: "greet".to_string(), diff --git a/crates/csp/src/indexing/index.rs b/crates/csp/src/indexing/index.rs index e11e028..c3fa7e1 100644 --- a/crates/csp/src/indexing/index.rs +++ b/crates/csp/src/indexing/index.rs @@ -1,7 +1,7 @@ //! `CspIndex` — the hybrid (dense + BM25) search orchestrator. Port of //! `src/indexing/index.ts` (← semble `index/index.py`). -use std::collections::{BTreeMap, HashSet}; +use std::collections::{BTreeMap, HashMap, HashSet}; use std::fmt::Write as _; use std::path::Path; use std::process::Command; @@ -80,6 +80,11 @@ pub struct CspIndex { pub model_path: String, pub root: Option, pub content: Vec, + /// Per-file character counts (repo-relative path → UTF-16 length) captured + /// at build time from the source tree, for token-savings telemetry. Empty + /// when the source files aren't available (e.g. a git index loaded from + /// cache). Derived metadata, not part of [`CspIndexState`]. + pub file_sizes: HashMap, } pub(crate) fn normalize_content(content: Option>) -> Vec { @@ -96,6 +101,7 @@ impl CspIndex { model_path: state.model_path, root: state.root, content: state.content, + file_sizes: HashMap::new(), } } @@ -120,15 +126,25 @@ impl CspIndex { }, )?; - Ok(Self::new(CspIndexState { + let mut index = Self::new(CspIndexState { model, bm25_index: result.bm25_index, semantic_index: result.semantic_index, chunks: result.chunks, model_path, - root: Some(path.to_string_lossy().into_owned()), + // Absolute, like upstream's `path.resolve()`, so an index built from + // `.` still finds its source tree when loaded from another cwd. + root: Some( + std::path::absolute(path) + .unwrap_or_else(|_| path.to_path_buf()) + .to_string_lossy() + .into_owned(), + ), content, - })) + }); + // Capture file sizes now, while the source tree is on disk. + index.file_sizes = compute_file_sizes(path, &index.chunks); + Ok(index) } /// Build an index from a remote git URL (shallow clone into a temp dir). @@ -149,9 +165,12 @@ impl CspIndex { clone_shallow(url, dir.path(), git_ref)?; let index = Self::from_path(dir.path(), options)?; + // `from_path` already captured file sizes from the checkout; carry them + // over since the temp dir is removed when `dir` drops. + let file_sizes = index.file_sizes.clone(); // Re-root at the URL so a persisted manifest records a stable sourceId // (the temp checkout is removed when `dir` drops). - Ok(Self::new(CspIndexState { + let mut rerooted = Self::new(CspIndexState { model: index.model, bm25_index: index.bm25_index, semantic_index: index.semantic_index, @@ -159,7 +178,9 @@ impl CspIndex { model_path: index.model_path, root: Some(url.to_string()), content: index.content, - })) + }); + rerooted.file_sizes = file_sizes; + Ok(rerooted) } /// Aggregate index statistics. @@ -348,7 +369,7 @@ impl CspIndex { make_stub_model(semantic_index.dim) }; - Ok(Self::new(CspIndexState { + let mut index = Self::new(CspIndexState { model, bm25_index, semantic_index, @@ -356,8 +377,66 @@ impl CspIndex { model_path, root: manifest.source_id, content: manifest.content, - })) + }); + // Recompute file sizes from the source when it's a still-present local + // directory (mirrors semble reading sizes off `root` on load). A git URL + // or a moved source leaves this empty → `file_chars` is simply 0. + if let Some(root) = index.root.as_deref() { + let root_path = Path::new(root); + if root_path.is_dir() { + index.file_sizes = compute_file_sizes(root_path, &index.chunks); + } + } + Ok(index) + } +} + +/// Per-file UTF-16 character counts for the unique files referenced by `chunks`, +/// read from `root`. Mirrors semble `_compute_file_sizes` (unreadable files are +/// skipped). Feeds the `file_chars` side of token-savings telemetry; UTF-16 keeps +/// it consistent with `stats::save_search_stats`'s snippet accounting. +/// +/// Chunk paths are repo-relative by construction; a path that is absolute or +/// escapes `root` via `..` can only come from a tampered on-disk index, so it is +/// skipped rather than resolved (path traversal guard — a deliberate addition +/// over upstream, which joins the path unchecked). Only regular files are read: +/// the file walker never follows symlinks, and a path that has since become a +/// symlink, FIFO, or device must not be able to redirect or stall the read. +fn compute_file_sizes(root: &Path, chunks: &[Chunk]) -> HashMap { + let mut sizes: HashMap = HashMap::new(); + for chunk in chunks { + if sizes.contains_key(&chunk.file_path) { + continue; + } + let rel = Path::new(&chunk.file_path); + if !is_safe_relative_path(rel) { + continue; + } + let full = root.join(rel); + let is_regular_file = std::fs::symlink_metadata(&full) + .map(|m| m.is_file()) + .unwrap_or(false); + if !is_regular_file { + continue; + } + if let Ok(text) = std::fs::read_to_string(&full) { + sizes.insert(chunk.file_path.clone(), text.encode_utf16().count() as u64); + } } + sizes +} + +/// `true` when `path` is relative and contains no `..` or root component, so +/// joining it onto an index root cannot resolve outside that root. +fn is_safe_relative_path(path: &Path) -> bool { + use std::path::Component; + !path.is_absolute() + && !path.components().any(|c| { + matches!( + c, + Component::ParentDir | Component::RootDir | Component::Prefix(_) + ) + }) } /// Shallow-clone `url` into `dir`, non-interactively. Rejects a ref starting diff --git a/crates/csp/src/indexing/index/tests.rs b/crates/csp/src/indexing/index/tests.rs index de9577a..96f9959 100644 --- a/crates/csp/src/indexing/index/tests.rs +++ b/crates/csp/src/indexing/index/tests.rs @@ -241,6 +241,20 @@ fn from_path_builds_index() { assert_eq!(idx.content, DEFAULT_CONTENT.to_vec()); } +#[test] +fn from_path_stores_absolute_root_for_relative_path() { + // Built from a cwd-relative path; `root` must still come out absolute so a + // reload from another cwd can find the source tree (upstream `path.resolve()`). + let dir = tempfile::Builder::new().tempdir_in(".").unwrap(); + std::fs::write(dir.path().join("sample.ts"), "export const x = 1\n").unwrap(); + let relative = Path::new(dir.path().file_name().unwrap()); + assert!(relative.is_relative()); + let idx = CspIndex::from_path(relative, &LoadOptions::default()).unwrap(); + let root = idx.root.as_deref().unwrap(); + assert!(Path::new(root).is_absolute(), "root was {root}"); + assert!(root.ends_with(relative.to_str().unwrap())); +} + // --- from_git --- #[test] @@ -297,3 +311,47 @@ fn from_git_clones_and_builds() { assert!(!idx.chunks.is_empty()); assert_eq!(idx.root.as_deref(), Some(url.as_str())); } + +#[test] +fn compute_file_sizes_skips_paths_that_escape_root() { + let outer = tempdir().unwrap(); + let root = outer.path().join("repo"); + std::fs::create_dir(&root).unwrap(); + std::fs::write(root.join("inside.ts"), "abc").unwrap(); + std::fs::write(outer.path().join("secret.txt"), "top secret").unwrap(); + let abs = root.join("inside.ts").to_string_lossy().into_owned(); + + let chunks = vec![ + make_chunk("inside.ts", 1, 1, None, "abc"), + make_chunk("../secret.txt", 1, 1, None, "x"), + make_chunk(&abs, 1, 1, None, "x"), + ]; + let sizes = compute_file_sizes(&root, &chunks); + + assert_eq!(sizes.get("inside.ts"), Some(&3)); + assert!(!sizes.contains_key("../secret.txt")); + assert!(!sizes.contains_key(abs.as_str())); +} + +#[cfg(unix)] +#[test] +fn compute_file_sizes_skips_symlinks_and_non_regular_files() { + let outer = tempdir().unwrap(); + let root = outer.path().join("repo"); + std::fs::create_dir(&root).unwrap(); + std::fs::write(root.join("real.ts"), "abcd").unwrap(); + std::fs::write(outer.path().join("secret.txt"), "top secret").unwrap(); + std::os::unix::fs::symlink(outer.path().join("secret.txt"), root.join("link.ts")).unwrap(); + std::fs::create_dir(root.join("dir.ts")).unwrap(); + + let chunks = vec![ + make_chunk("real.ts", 1, 1, None, "abcd"), + make_chunk("link.ts", 1, 1, None, "x"), + make_chunk("dir.ts", 1, 1, None, "x"), + ]; + let sizes = compute_file_sizes(&root, &chunks); + + assert_eq!(sizes.get("real.ts"), Some(&4)); + assert!(!sizes.contains_key("link.ts")); + assert!(!sizes.contains_key("dir.ts")); +} diff --git a/crates/csp/src/mcp.rs b/crates/csp/src/mcp.rs index f58dbad..ff80b78 100644 --- a/crates/csp/src/mcp.rs +++ b/crates/csp/src/mcp.rs @@ -9,6 +9,7 @@ //! testable. [`IndexCache`] holds `Arc` so it can be shared across the //! async server's tokio tasks. +use std::path::Path; use std::sync::Arc; use std::time::{Duration, Instant}; @@ -18,7 +19,8 @@ use serde_json::json; use crate::indexing::index::{ load_or_build_index, source_fingerprint, CspIndex, LoadOrBuildOptions, QueryOptions, }; -use crate::types::ContentType; +use crate::stats::save_search_stats; +use crate::types::{CallType, ContentType}; use crate::utils::{format_results, is_git_url, resolve_chunk}; /// Server instructions advertised to MCP clients (preserved for the transport). @@ -239,6 +241,10 @@ pub fn get_index( /// `search` tool handler. Returns a JSON string (results or `{error}`), or an /// error message string on failure (mirroring the TS handler's catch). +/// `stats_file`, when `Some`, records token-savings telemetry (tests pass `None`). +// Positional transport params mirror the MCP tool signature; a struct would just +// move the plumbing without clarifying it (same call as `find_related_tool`). +#[allow(clippy::too_many_arguments)] pub fn search_tool( cache: &mut IndexCache, default_source: Option<&str>, @@ -247,6 +253,7 @@ pub fn search_tool( repo: Option<&str>, top_k: usize, max_snippet_lines: Option, + stats_file: Option<&Path>, ) -> String { let index = match get_index(repo, default_source, default_ref, cache) { Ok(idx) => idx, @@ -259,6 +266,15 @@ pub fn search_tool( ..Default::default() }, ); + if let Some(stats_file) = stats_file { + save_search_stats( + stats_file, + &results, + CallType::Search, + &index.file_sizes, + max_snippet_lines, + ); + } if results.is_empty() { json!({ "error": "No results found." }).to_string() } else { @@ -279,6 +295,7 @@ pub fn find_related_tool( repo: Option<&str>, top_k: usize, max_snippet_lines: Option, + stats_file: Option<&Path>, ) -> String { let index = match get_index(repo, default_source, default_ref, cache) { Ok(idx) => idx, @@ -304,6 +321,15 @@ pub fn find_related_tool( ..Default::default() }, ); + if let Some(stats_file) = stats_file { + save_search_stats( + stats_file, + &results, + CallType::FindRelated, + &index.file_sizes, + max_snippet_lines, + ); + } if results.is_empty() { json!({ "error": format!("No related chunks found for {file_path}:{line}.") }).to_string() } else { @@ -541,6 +567,7 @@ mod tests { None, 5, None, + None, ); assert_eq!(out, json!({ "error": "No results found." }).to_string()); } @@ -564,7 +591,16 @@ mod tests { #[test] fn search_tool_returns_results_json() { let mut cache = IndexCache::with_seam(vec![ContentType::Code], OneChunkSeam); - let out = search_tool(&mut cache, Some("/tmp/repo"), None, "main", None, 5, None); + let out = search_tool( + &mut cache, + Some("/tmp/repo"), + None, + "main", + None, + 5, + None, + None, + ); let value: serde_json::Value = serde_json::from_str(&out).unwrap(); assert!(value.get("query").is_some()); assert!(value["results"].as_array().is_some()); @@ -572,6 +608,27 @@ mod tests { assert!(value["results"][0].get("content").is_some()); } + #[test] + fn search_tool_records_savings_when_stats_file_given() { + let mut cache = IndexCache::with_seam(vec![ContentType::Code], OneChunkSeam); + let dir = tempfile::tempdir().unwrap(); + let stats_file = dir.path().join("savings.jsonl"); + let _ = search_tool( + &mut cache, + Some("/tmp/repo"), + None, + "main", + None, + 5, + None, + Some(&stats_file), + ); + let content = std::fs::read_to_string(&stats_file).unwrap(); + let lines: Vec<&str> = content.lines().filter(|l| !l.is_empty()).collect(); + assert_eq!(lines.len(), 1); + assert!(lines[0].contains("\"call\":\"search\"")); + } + #[test] fn search_tool_respects_max_snippet_lines_zero() { let mut cache = IndexCache::with_seam(vec![ContentType::Code], OneChunkSeam); @@ -583,6 +640,7 @@ mod tests { None, 5, Some(0), + None, ); let value: serde_json::Value = serde_json::from_str(&out).unwrap(); let entry = &value["results"][0]; @@ -603,6 +661,7 @@ mod tests { None, 5, None, + None, ); assert!(out.contains("No chunk found at nope.ts:1")); } @@ -619,6 +678,7 @@ mod tests { None, 5, None, + None, ); // Either related results or the no-related error — both valid JSON. let value: serde_json::Value = serde_json::from_str(&out).unwrap(); diff --git a/crates/csp/src/stats.rs b/crates/csp/src/stats.rs index 6be0370..89d56ec 100644 --- a/crates/csp/src/stats.rs +++ b/crates/csp/src/stats.rs @@ -82,14 +82,41 @@ fn utf16_len(s: &str) -> u64 { s.encode_utf16().count() as u64 } +/// Characters actually delivered to the caller for one result, honoring the +/// `max_snippet_lines` cap (mirrors semble#206): `None` → the whole chunk, +/// `Some(0)` → nothing, `Some(n)` → the first `n` lines. +fn delivered_chars(content: &str, max_snippet_lines: Option) -> u64 { + match max_snippet_lines { + None => utf16_len(content), + Some(0) => 0, + Some(n) => { + // Sum of the first `n` lines plus one `\n` between each pair, without + // materializing the joined snippet. + let (len, count) = content + .lines() + .take(n) + .fold((0u64, 0u64), |(len, count), line| { + (len + utf16_len(line), count + 1) + }); + len + count.saturating_sub(1) + } + } +} + /// Append one telemetry record. Best-effort: any I/O error is swallowed. +/// `max_snippet_lines` matches the cap applied to the returned snippet so the +/// recorded `snippet_chars` reflects what the caller actually received. pub fn save_search_stats( stats_file: &Path, results: &[SearchResult], call_type: CallType, file_sizes: &HashMap, + max_snippet_lines: Option, ) { - let snippet_chars: u64 = results.iter().map(|r| utf16_len(&r.chunk.content)).sum(); + let snippet_chars: u64 = results + .iter() + .map(|r| delivered_chars(&r.chunk.content, max_snippet_lines)) + .sum(); let mut unique_paths: Vec<&str> = Vec::new(); for r in results { if !unique_paths.contains(&r.chunk.file_path.as_str()) { @@ -476,6 +503,7 @@ mod tests { &results, CallType::Search, &sizes(&[("a.ts", 100), ("b.ts", 200)]), + None, ); let content = std::fs::read_to_string(&file).unwrap(); @@ -488,12 +516,55 @@ mod tests { assert_eq!(record.file_chars, 300); } + #[test] + fn save_caps_snippet_chars_by_max_snippet_lines() { + let dir = tempdir().unwrap(); + let file = dir.path().join("savings.jsonl"); + let results = vec![result("line1\nline2\nline3", "a.ts")]; + // The first 2 lines "line1\nline2" are 11 UTF-16 units (semble#206). + save_search_stats( + &file, + &results, + CallType::Search, + &sizes(&[("a.ts", 100)]), + Some(2), + ); + let content = std::fs::read_to_string(&file).unwrap(); + let record: StatsRecord = serde_json::from_str(content.lines().next().unwrap()).unwrap(); + assert_eq!(record.snippet_chars, 11); + assert_eq!(record.file_chars, 100); + } + + #[test] + fn save_zero_snippet_lines_records_no_snippet_chars() { + let dir = tempdir().unwrap(); + let file = dir.path().join("savings.jsonl"); + let results = vec![result("anything at all", "a.ts")]; + save_search_stats( + &file, + &results, + CallType::Search, + &sizes(&[("a.ts", 100)]), + Some(0), + ); + let content = std::fs::read_to_string(&file).unwrap(); + let record: StatsRecord = serde_json::from_str(content.lines().next().unwrap()).unwrap(); + assert_eq!(record.snippet_chars, 0); + assert_eq!(record.file_chars, 100); + } + #[test] fn save_dedups_file_chars_per_path() { let dir = tempdir().unwrap(); let file = dir.path().join("savings.jsonl"); let results = vec![result("abc", "a.ts"), result("def", "a.ts")]; - save_search_stats(&file, &results, CallType::Search, &sizes(&[("a.ts", 100)])); + save_search_stats( + &file, + &results, + CallType::Search, + &sizes(&[("a.ts", 100)]), + None, + ); let content = std::fs::read_to_string(&file).unwrap(); let record: StatsRecord = serde_json::from_str(content.lines().next().unwrap()).unwrap(); assert_eq!(record.file_chars, 100); @@ -505,7 +576,13 @@ mod tests { let dir = tempdir().unwrap(); let file = dir.path().join("savings.jsonl"); let results = vec![result("x", "a.ts"), result("y", "missing.ts")]; - save_search_stats(&file, &results, CallType::Search, &sizes(&[("a.ts", 100)])); + save_search_stats( + &file, + &results, + CallType::Search, + &sizes(&[("a.ts", 100)]), + None, + ); let content = std::fs::read_to_string(&file).unwrap(); let record: StatsRecord = serde_json::from_str(content.lines().next().unwrap()).unwrap(); assert_eq!(record.file_chars, 100); @@ -520,12 +597,14 @@ mod tests { &[result("a", "a.ts")], CallType::Search, &sizes(&[("a.ts", 10)]), + None, ); save_search_stats( &file, &[result("b", "b.ts")], CallType::FindRelated, &sizes(&[("b.ts", 10)]), + None, ); let content = std::fs::read_to_string(&file).unwrap(); let lines: Vec<&str> = content.lines().filter(|l| !l.is_empty()).collect();