diff --git a/crates/csp/src/bin/csp/main.rs b/crates/csp/src/bin/csp/main.rs index e041706..e2a1f6f 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}; #[derive(Parser)] @@ -224,12 +226,14 @@ fn resolve_snippet_lines(value: Option) -> Option { value.map(|n| n.max(0) as usize) } -/// 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, @@ -238,6 +242,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 { @@ -253,6 +266,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}")); @@ -274,6 +288,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 { @@ -356,6 +379,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); @@ -426,6 +456,7 @@ fn dispatch(command: Command) -> u8 { &query, top_k.unwrap_or(5), resolve_snippet_lines(max_snippet_lines), + Some(stats_file), ) ); EXIT_SUCCESS @@ -465,6 +496,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}"); @@ -559,7 +591,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()) { @@ -578,7 +610,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()); @@ -586,11 +618,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")); } @@ -598,7 +648,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")); } @@ -654,49 +704,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 0b47156..0ece8dd 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; /// MCP default: signature + first body lines, enough to confirm a location @@ -70,6 +71,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, } @@ -84,6 +88,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(), } } @@ -104,6 +109,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)])) } @@ -125,6 +131,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)])) } @@ -236,11 +243,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 1de41c1..643ebbe 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,7 +126,7 @@ impl CspIndex { }, )?; - Ok(Self::new(CspIndexState { + let mut index = Self::new(CspIndexState { model, bm25_index: result.bm25_index, semantic_index: result.semantic_index, @@ -128,7 +134,10 @@ impl CspIndex { model_path, root: Some(path.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 +158,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 +171,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 +362,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 +370,35 @@ 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. +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; + } + if let Ok(text) = std::fs::read_to_string(root.join(&chunk.file_path)) { + sizes.insert(chunk.file_path.clone(), text.encode_utf16().count() as u64); + } } + sizes } /// Shallow-clone `url` into `dir`, non-interactively. Rejects a ref starting diff --git a/crates/csp/src/mcp.rs b/crates/csp/src/mcp.rs index 9447a17..2007bf1 100644 --- a/crates/csp/src/mcp.rs +++ b/crates/csp/src/mcp.rs @@ -9,13 +9,15 @@ //! 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 indexmap::IndexMap; use serde_json::json; use crate::indexing::index::{load_or_build_index, 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). @@ -165,6 +167,8 @@ 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`). +#[allow(clippy::too_many_arguments)] pub fn search_tool( cache: &mut IndexCache, default_source: Option<&str>, @@ -173,6 +177,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, @@ -185,6 +190,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 { @@ -205,6 +219,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, @@ -230,6 +245,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 { @@ -410,6 +434,7 @@ mod tests { None, 5, None, + None, ); assert_eq!(out, json!({ "error": "No results found." }).to_string()); } @@ -429,7 +454,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()); @@ -437,6 +471,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); @@ -448,6 +503,7 @@ mod tests { None, 5, Some(0), + None, ); let value: serde_json::Value = serde_json::from_str(&out).unwrap(); let entry = &value["results"][0]; @@ -468,6 +524,7 @@ mod tests { None, 5, None, + None, ); assert!(out.contains("No chunk found at nope.ts:1")); } @@ -484,6 +541,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..1f6250b 100644 --- a/crates/csp/src/stats.rs +++ b/crates/csp/src/stats.rs @@ -82,14 +82,34 @@ 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) => { + let snippet = content.lines().take(n).collect::>().join("\n"); + utf16_len(&snippet) + } + } +} + /// 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 +496,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 +509,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 +569,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 +590,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();