Skip to content

feat(stats): wire token-savings telemetry into search and find_related #82

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
amondnet wants to merge 1 commit into
base: amondnet/max-snippet-lines
Choose a base branch
from
Open

feat(stats): wire token-savings telemetry into search and find_related #82

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
139 changes: 102 additions & 37 deletions crates/csp/src/bin/csp/main.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,8 +14,10 @@
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)]
Expand Down Expand Up @@ -224,12 +226,14 @@
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<usize>,
stats_file: Option<&Path>,
) -> String {
let results = index.search(
query,
Expand All @@ -238,6 +242,15 @@
..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 {
Expand All @@ -253,6 +266,7 @@
line: &str,
top_k: usize,
max_snippet_lines: Option<usize>,
stats_file: Option<&Path>,
) -> Result<String, String> {
let Ok(line_num) = line.parse::<i64>() else {
return Err(format!("line must be an integer, got: {line}"));
Expand All @@ -274,6 +288,15 @@
..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 {
Expand Down Expand Up @@ -356,6 +379,13 @@
/// `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 {

Check failure on line 388 in crates/csp/src/bin/csp/main.rs

View check run for this annotation

SonarQubeCloud / SonarCloud Code Analysis

Refactor this function to reduce its Cognitive Complexity from 16 to the 15 allowed. 


See more on https://sonarcloud.io/project/issues?id=pleaseai_code-search&issues=AZ9RGLnIDf33n6NsTOE5&open=AZ9RGLnIDf33n6NsTOE5&pullRequest=82
match command {
Command::Init { agent, force } => {
let agent = agent.unwrap_or(Agent::Claude);
Expand Down Expand Up @@ -426,6 +456,7 @@
&query,
top_k.unwrap_or(5),
resolve_snippet_lines(max_snippet_lines),
Some(stats_file),
)
);
EXIT_SUCCESS
Expand Down Expand Up @@ -465,6 +496,7 @@
&line,
top_k.unwrap_or(5),
resolve_snippet_lines(max_snippet_lines),
Some(stats_file),
) {
Ok(out) => {
println!("{out}");
Expand Down Expand Up @@ -559,7 +591,7 @@
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()) {
Expand All @@ -578,27 +610,45 @@
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());
assert!(first.get("file_path").is_some());
}
}

#[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"));
}

#[test]
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"));
}

Expand Down Expand Up @@ -654,49 +704,64 @@
#[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]
Expand Down
11 changes: 10 additions & 1 deletion crates/csp/src/bin/csp/mcp_server.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -70,6 +71,9 @@ pub struct CspMcpServer {
cache: Arc<Mutex<IndexCache>>,
default_source: Option<String>,
default_ref: Option<String>,
/// 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<std::path::PathBuf>,
tool_router: ToolRouter<CspMcpServer>,
}

Expand All @@ -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(),
}
}
Expand All @@ -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)]))
}
Expand All @@ -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)]))
}
Expand Down Expand Up @@ -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(),
Expand Down
Loading
Loading