1pub mod proxy;
69use std::io::{BufRead, Write};
70use std::path::{Path, PathBuf};
71use std::sync::{Arc, Mutex};
72
73use serde::{Deserialize, Serialize};
74use serde_json::Value;
75use sqz_engine::{SqzEngine, ToolDefinition, ToolSelector};
76use sqz_engine::error::{Result, SqzError};
77use sqz_engine::preset::{Preset, PresetParser};
78
79#[derive(Debug, Clone, Serialize, Deserialize)]
84pub struct ToolCallRequest {
85 pub tool_id: String,
86 pub input: Value,
87 pub intent: Option<String>,
88}
89
90#[derive(Debug, Clone, Serialize, Deserialize)]
93pub struct ToolCallResponse {
94 pub tool_id: String,
95 pub output: String,
97 pub tokens_original: u32,
98 pub tokens_compressed: u32,
99}
100
101#[derive(Debug, Clone)]
103pub enum McpTransport {
104 Stdio,
105 Sse { port: u16 },
106}
107
108#[derive(Debug, Deserialize)]
111struct JsonRpcRequest {
112 #[allow(dead_code)]
113 jsonrpc: String,
114 id: Option<Value>,
115 method: String,
116 params: Option<Value>,
117}
118
119#[derive(Debug, Serialize)]
120struct JsonRpcResponse {
121 jsonrpc: String,
122 id: Option<Value>,
123 #[serde(skip_serializing_if = "Option::is_none")]
124 result: Option<Value>,
125 #[serde(skip_serializing_if = "Option::is_none")]
126 error: Option<JsonRpcError>,
127}
128
129#[derive(Debug, Serialize)]
130struct JsonRpcError {
131 code: i32,
132 message: String,
133}
134
135impl JsonRpcResponse {
136 fn ok(id: Option<Value>, result: Value) -> Self {
137 Self {
138 jsonrpc: "2.0".to_string(),
139 id,
140 result: Some(result),
141 error: None,
142 }
143 }
144
145 fn err(id: Option<Value>, code: i32, message: impl Into<String>) -> Self {
146 Self {
147 jsonrpc: "2.0".to_string(),
148 id,
149 result: None,
150 error: Some(JsonRpcError { code, message: message.into() }),
151 }
152 }
153}
154
155struct SharedState {
160 pending_preset: Mutex<Option<String>>,
163 tool_selector: Mutex<ToolSelector>,
165 registered_tools: Mutex<Vec<ToolDefinition>>,
167}
168
169pub struct McpServer {
173 engine: SqzEngine,
174 shared: Arc<SharedState>,
175 preset_dir: PathBuf,
176}
177
178impl McpServer {
179 pub fn new(preset_dir: &Path) -> Result<Self> {
181 let engine = SqzEngine::new()?;
182 Self::with_engine(preset_dir, engine)
183 }
184
185 #[cfg(test)]
191 fn new_with_store(preset_dir: &Path, store_path: &Path) -> Result<Self> {
192 let engine = SqzEngine::with_preset_and_store(Preset::default(), store_path)?;
193 Self::with_engine(preset_dir, engine)
194 }
195
196 fn with_engine(preset_dir: &Path, engine: SqzEngine) -> Result<Self> {
200 let preset = Preset::default();
201 let model_path = Path::new("");
202 let mut tool_selector = ToolSelector::new(model_path, &preset)?;
203
204 let default_tools = default_tool_definitions();
206 tool_selector.register_tools(&default_tools)?;
207
208 let shared = Arc::new(SharedState {
209 pending_preset: Mutex::new(None),
210 tool_selector: Mutex::new(tool_selector),
211 registered_tools: Mutex::new(default_tools),
212 });
213
214 Ok(McpServer {
215 engine,
216 shared,
217 preset_dir: preset_dir.to_owned(),
218 })
219 }
220
221 fn apply_pending_preset(&mut self) {
223 let pending = {
224 let mut guard = self.shared.pending_preset.lock()
225 .unwrap_or_else(|e| e.into_inner());
226 guard.take()
227 };
228 if let Some(toml_str) = pending {
229 match self.engine.reload_preset(&toml_str) {
230 Ok(()) => {
231 if let Ok(new_preset) = PresetParser::parse(&toml_str) {
233 if let Ok(mut sel) = self.shared.tool_selector.lock() {
234 if let Ok(mut new_sel) = ToolSelector::new(Path::new(""), &new_preset) {
235 if let Ok(tools) = self.shared.registered_tools.lock() {
236 let _ = new_sel.register_tools(&tools);
237 }
238 *sel = new_sel;
239 }
240 }
241 }
242 eprintln!("[sqz-mcp] preset applied from hot-reload");
243 }
244 Err(e) => eprintln!("[sqz-mcp] engine reload error: {e}"),
245 }
246 }
247 }
248
249 pub fn handle_tool_call(&mut self, request: ToolCallRequest) -> Result<ToolCallResponse> {
274 self.apply_pending_preset();
275
276 match request.tool_id.as_str() {
277 "passthrough" => self.handle_passthrough(request),
278 "expand" => self.handle_expand(request),
279 "sqz_recall" => self.handle_sqz_recall(request),
280 "sqz_read_file" => self.handle_sqz_read_file(request),
281 "sqz_grep" => self.handle_sqz_grep(request),
282 "sqz_list_dir" => self.handle_sqz_list_dir(request),
283 _ => self.handle_compress(request),
285 }
286 }
287
288 fn log_compression(&self, tool_id: &str, tokens_original: u32, tokens_compressed: u32) {
291 let project = std::env::current_dir().ok();
292 let project_str = project.as_ref().map(|p| p.to_string_lossy().to_string());
293 let _ = self.engine.session_store().log_compression_with_project(
294 tokens_original,
295 tokens_compressed,
296 &[],
297 tool_id,
298 project_str.as_deref(),
299 );
300 }
301
302 fn handle_compress(&mut self, request: ToolCallRequest) -> Result<ToolCallResponse> {
310 let raw_input = match request.input.get("text").and_then(|v| v.as_str()) {
312 Some(text) => text.to_string(),
313 None => serde_json::to_string(&request.input)
314 .map_err(|e| SqzError::Other(format!("input serialization error: {e}")))?,
315 };
316
317 let compressed = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
318 let tokens_original = self.engine.count_tokens(&raw_input);
319 self.compress_cached(&raw_input, tokens_original, false, true)
320 .map(|(output, tokens_compressed)| (output, tokens_original, tokens_compressed))
321 }));
322 let (output, tokens_original, tokens_compressed) = match compressed {
323 Ok(result) => result?,
324 Err(_) => {
325 self.engine.clear_poison();
326 let tokens = estimate_tokens(&raw_input);
327 (raw_input, tokens, tokens)
328 }
329 };
330 self.log_compression(&request.tool_id, tokens_original, tokens_compressed);
331
332 Ok(ToolCallResponse {
333 tool_id: request.tool_id,
334 output,
335 tokens_original,
336 tokens_compressed,
337 })
338 }
339
340 fn handle_sqz_recall(&mut self, request: ToolCallRequest) -> Result<ToolCallResponse> {
347 let query = request
348 .input
349 .get("query")
350 .and_then(|v| v.as_str())
351 .ok_or_else(|| {
352 SqzError::Other("sqz_recall: input must be { \"query\": \"<terms>\" }".to_string())
353 })?;
354 let limit = request
355 .input
356 .get("limit")
357 .and_then(|v| v.as_u64())
358 .map(|n| n.clamp(1, 25) as u32)
359 .unwrap_or(5);
360
361 let hits = self.engine.session_store().recall_search(query, limit)?;
362 let mut output = if hits.is_empty() {
363 format!("[sqz:recall no matches for \"{query}\"]")
364 } else {
365 let mut out = format!("[sqz:recall query=\"{query}\" hits={}]\n", hits.len());
366 let mut any_output_kind = false;
367 for (i, hit) in hits.iter().enumerate() {
368 let snippet = hit.snippet.replace('\n', " ");
369 if hit.kind == "output" {
370 any_output_kind = true;
371 out.push_str(&format!(
372 "{}. [{} {}] ref={} — {}\n",
373 i + 1,
374 hit.kind,
375 hit.created_at,
376 &hit.ref_hash[..hit.ref_hash.len().min(16)],
377 snippet,
378 ));
379 } else {
380 out.push_str(&format!(
381 "{}. [{} {}] id={} — {}\n",
382 i + 1,
383 hit.kind,
384 hit.created_at,
385 hit.ref_hash,
386 snippet,
387 ));
388 }
389 }
390 if any_output_kind {
391 out.push_str("Pass a ref to the expand tool for the full original content.\n");
392 }
393 out
394 };
395 if output.ends_with('\n') {
396 output.pop();
397 }
398 let tokens = estimate_tokens(&output);
399 Ok(ToolCallResponse {
400 tool_id: request.tool_id,
401 output,
402 tokens_original: tokens,
403 tokens_compressed: tokens,
404 })
405 }
406
407 fn handle_passthrough(&mut self, request: ToolCallRequest) -> Result<ToolCallResponse> {
420 let text = match request.input.get("text").and_then(|v| v.as_str()) {
421 Some(s) => s.to_string(),
422 None => {
423 serde_json::to_string(&request.input)
425 .map_err(|e| SqzError::Other(format!("input serialization error: {e}")))?
426 }
427 };
428 let tokens = estimate_tokens(&text);
429 Ok(ToolCallResponse {
430 tool_id: request.tool_id,
431 output: text,
432 tokens_original: tokens,
433 tokens_compressed: tokens,
436 })
437 }
438
439 fn handle_expand(&mut self, request: ToolCallRequest) -> Result<ToolCallResponse> {
445 let raw = request
446 .input
447 .get("prefix")
448 .and_then(|v| v.as_str())
449 .ok_or_else(|| {
450 SqzError::Other("expand: input must be { \"prefix\": \"<hex>\" }".to_string())
451 })?;
452 let (prefix, range) = sqz_engine::parse_ref_token(raw);
455
456 let result = self.engine.cache_manager().expand_ref(raw)?;
457 let output = match result {
458 Some(sqz_engine::ExpandResult::Original { bytes, hash }) => {
459 let as_text = String::from_utf8_lossy(&bytes).into_owned();
466 match range {
467 Some((a, b)) => format!("[sqz:expand hash={hash} lines={a}-{b}]\n{as_text}"),
468 None => format!("[sqz:expand hash={hash}]\n{as_text}"),
469 }
470 }
471 Some(sqz_engine::ExpandResult::CompressedOnly { compressed, hash }) => {
472 format!(
473 "[sqz:expand hash={hash} note=compressed-only (predates original-capture migration)]\n{compressed}"
474 )
475 }
476 None => {
477 format!("[sqz:expand hash-not-found prefix={prefix}]")
478 }
479 };
480
481 let tokens = estimate_tokens(&output);
482 Ok(ToolCallResponse {
483 tool_id: request.tool_id,
484 output,
485 tokens_original: tokens,
486 tokens_compressed: tokens,
487 })
488 }
489
490 fn compress_cached(&self, text: &str, tokens: u32, lossless: bool, allow_delta: bool) -> Result<(String, u32)> {
510 use sqz_engine::CacheResult;
511 #[cfg(test)]
512 tests::maybe_panic();
513 let result = if lossless {
516 self.engine.compress_with_cache_lossless(text)?
517 } else {
518 self.engine.compress_with_cache(text)?
519 };
520 let tiny = tokens <= REF_TOKENS;
521 Ok(match result {
522 CacheResult::Dedup { inline_ref, token_cost } if !tiny => (inline_ref, token_cost),
523 CacheResult::Delta { delta_text, token_cost, .. } if allow_delta && !tiny => (delta_text, token_cost),
524 CacheResult::Dedup { .. } | CacheResult::Delta { .. } => return self.compress_fresh(text, lossless),
525 CacheResult::Fresh { output } => {
526 let truncated = output.stages_applied.iter().any(|s| s == "entropy_truncate");
527 if truncated && !self.engine.never_store(text) {
528 let hash = sqz_engine::CacheManager::sha256_hex(text.as_bytes());
529 let data = format!("{}\n[full output: call expand with prefix \"{}\"]", output.data, &hash[..16]);
530 let tokens = self.engine.count_tokens(&data);
531 (data, tokens)
532 } else {
533 (output.data, output.tokens_compressed)
534 }
535 }
536 })
537 }
538
539 fn compress_fresh(&self, text: &str, lossless: bool) -> Result<(String, u32)> {
541 let output = if lossless {
542 self.engine.compress_lossless(text)?
543 } else {
544 self.engine.compress(text)?
545 };
546 Ok((output.data, output.tokens_compressed))
547 }
548
549 fn handle_sqz_read_file(&mut self, request: ToolCallRequest) -> Result<ToolCallResponse> {
582 let path_str = request
583 .input
584 .get("path")
585 .and_then(|v| v.as_str())
586 .ok_or_else(|| {
587 SqzError::Other(
588 "sqz_read_file: input must be { \"path\": \"<file>\" }".to_string(),
589 )
590 })?;
591
592 let max_bytes = cap_from(&request.input, "max_bytes", READ_DEFAULT_MAX_BYTES);
596
597 let path = std::path::PathBuf::from(path_str);
598 if let Ok(meta) = std::fs::metadata(&path) {
599 if meta.len() > READ_HARD_LIMIT_BYTES {
600 return Err(SqzError::Other(format!(
601 "sqz_read_file: '{}' is {} MB, too large to load; use head/sed/tail through the shell for files this size",
602 path.display(),
603 meta.len() / (1024 * 1024)
604 )));
605 }
606 }
607 let bytes = match std::fs::read(&path) {
608 Ok(b) => b,
609 Err(e) => {
610 return Err(SqzError::Other(format!(
611 "sqz_read_file: could not read '{}': {e}",
612 path.display()
613 )))
614 }
615 };
616
617 let full_text = String::from_utf8_lossy(&bytes).into_owned();
621
622 let offset = request.input.get("offset").and_then(|v| v.as_u64()).map(|v| v.max(1) as usize);
627 let limit = request.input.get("limit").and_then(|v| v.as_u64()).map(|v| v as usize);
628 let total_lines = full_text.lines().count();
629 let (raw_text, range) = match (offset, limit) {
630 (None, None) => (full_text, None),
631 (offset, limit) => {
632 let start = offset.unwrap_or(1);
633 let end = match limit {
634 Some(n) => start.saturating_add(n).saturating_sub(1).min(total_lines),
635 None => total_lines,
636 };
637 if start > total_lines {
638 return Err(SqzError::Other(format!(
639 "sqz_read_file: offset {start} is past the end of '{}' ({total_lines} lines)",
640 path.display()
641 )));
642 }
643 (line_range(&full_text, start, end).to_string(), Some((start, end)))
644 }
645 };
646
647 let first_line = range.map(|(a, _)| a).unwrap_or(1);
650 let mut shown = range;
651 let mut cut = None;
652 let raw_text = match truncate_to_lines(&raw_text, max_bytes) {
653 Some(kept) => {
654 let last = first_line + kept.lines().count().max(1) - 1;
655 shown = Some((first_line, last));
656 cut = Some((last, !kept.ends_with('\n')));
657 kept.to_string()
658 }
659 None => raw_text,
660 };
661
662 let tokens_original = self.engine.count_tokens(&raw_text);
663 let (compressed_data, tokens_compressed) =
664 self.compress_cached(&raw_text, tokens_original, true, range.is_none())?;
665
666 let mut header = format!("[sqz_read_file path={} size={}", path.display(), bytes.len());
670 if let Some((a, b)) = shown {
671 header.push_str(&format!(" lines={a}-{b} of {total_lines}"));
672 }
673 if let Some((last, mid_line)) = cut {
674 header.push_str(&format!(" truncated_to={max_bytes}"));
675 if mid_line {
676 header.push_str(&format!(" line_cut={last}"));
677 }
678 if last < total_lines {
679 header.push_str(&format!(" continue_with_offset={}", last + 1));
680 }
681 }
682 let output = format!("{header}]\n{compressed_data}");
683
684 self.log_compression(&request.tool_id, tokens_original, tokens_compressed);
685
686 Ok(ToolCallResponse {
687 tool_id: request.tool_id,
688 output,
689 tokens_original,
690 tokens_compressed,
691 })
692 }
693
694 fn handle_sqz_list_dir(&mut self, request: ToolCallRequest) -> Result<ToolCallResponse> {
707 let path_str = request
708 .input
709 .get("path")
710 .and_then(|v| v.as_str())
711 .unwrap_or(".");
712
713 let max_depth = request
716 .input
717 .get("max_depth")
718 .and_then(|v| v.as_u64())
719 .unwrap_or(1)
720 .max(1) as usize;
721
722 let root = std::path::PathBuf::from(path_str);
723 let mut listing = Listing {
724 lines: Vec::new(),
725 max_depth,
726 max_entries: cap_from(&request.input, "max_entries", LIST_DEFAULT_MAX_ENTRIES),
727 capped: false,
728 };
729 list_dir_recursive(&root, &root, 1, &mut listing)?;
730
731 let raw = listing.lines.join("\n");
732 let tokens_original = self.engine.count_tokens(&raw);
733 let (compressed_data, tokens_compressed) = self.compress_cached(&raw, tokens_original, true, true)?;
734 self.log_compression(&request.tool_id, tokens_original, tokens_compressed);
735
736 let mut header = format!(
737 "[sqz_list_dir path={} entries={}",
738 root.display(),
739 listing.lines.len()
740 );
741 if listing.capped {
742 header.push_str(&format!(" stopped_at_max_entries={}", listing.max_entries));
743 }
744 let output = format!("{header}]\n{compressed_data}");
745
746 Ok(ToolCallResponse {
747 tool_id: request.tool_id,
748 output,
749 tokens_original,
750 tokens_compressed,
751 })
752 }
753
754 fn handle_sqz_grep(&mut self, request: ToolCallRequest) -> Result<ToolCallResponse> {
772 let pattern = request
773 .input
774 .get("pattern")
775 .and_then(|v| v.as_str())
776 .ok_or_else(|| {
777 SqzError::Other(
778 "sqz_grep: input must include { \"pattern\": \"<text>\" }".to_string(),
779 )
780 })?;
781
782 let path_str = request
783 .input
784 .get("path")
785 .and_then(|v| v.as_str())
786 .unwrap_or(".");
787
788 let limits = GrepLimits {
789 max_matches: request
790 .input
791 .get("max_matches")
792 .and_then(|v| v.as_u64())
793 .unwrap_or(200) as usize,
794 max_line_chars: cap_from(&request.input, "max_line_chars", GREP_DEFAULT_MAX_LINE_CHARS),
795 max_bytes: cap_from(&request.input, "max_bytes", GREP_DEFAULT_MAX_BYTES),
796 };
797
798 let use_regex = request
799 .input
800 .get("regex")
801 .and_then(|v| v.as_bool())
802 .unwrap_or(false);
803
804 let root = std::path::PathBuf::from(path_str);
805
806 let regex = if use_regex {
809 match regex::Regex::new(pattern) {
810 Ok(r) => Some(r),
811 Err(e) => {
812 return Err(SqzError::Other(format!(
813 "sqz_grep: invalid regex: {e}"
814 )))
815 }
816 }
817 } else {
818 None
819 };
820
821 let mut found = GrepResults::default();
822 grep_walk(&root, pattern, regex.as_ref(), &limits, &mut found)?;
823
824 let raw = found.lines.join("\n");
825 let tokens_original = self.engine.count_tokens(&raw);
826 let (compressed_data, tokens_compressed) = self.compress_cached(&raw, tokens_original, true, true)?;
827 self.log_compression(&request.tool_id, tokens_original, tokens_compressed);
828
829 let mut header = format!(
830 "[sqz_grep pattern={:?} root={} matches={}",
831 pattern,
832 root.display(),
833 found.lines.len(),
834 );
835 if found.lines.len() >= limits.max_matches {
836 header.push_str(&format!(" max_matches_reached={}", limits.max_matches));
837 }
838 if found.byte_capped {
839 header.push_str(&format!(" stopped_at_max_bytes={}", limits.max_bytes));
840 }
841 if found.clipped > 0 {
842 header.push_str(&format!(
843 " long_lines_clipped={} (max_line_chars={})",
844 found.clipped, limits.max_line_chars
845 ));
846 }
847 let output = format!("{header}]\n{compressed_data}");
848
849 Ok(ToolCallResponse {
850 tool_id: request.tool_id,
851 output,
852 tokens_original,
853 tokens_compressed,
854 })
855 }
856
857 pub fn list_tools(&self, intent: Option<&str>) -> Result<Vec<ToolDefinition>> {
859 let tools = self.shared.registered_tools.lock()
860 .unwrap_or_else(|e| e.into_inner());
861
862 match intent {
863 Some(intent_str) if !intent_str.is_empty() => {
864 let selector = self.shared.tool_selector.lock()
865 .unwrap_or_else(|e| e.into_inner());
866 let selected_ids = selector.select(intent_str, 5)?;
867 let filtered: Vec<ToolDefinition> = tools
868 .iter()
869 .filter(|t| selected_ids.contains(&t.id))
870 .cloned()
871 .collect();
872 Ok(filtered)
873 }
874 _ => Ok(tools.clone()),
875 }
876 }
877
878 pub fn start(self, transport: McpTransport) -> Result<()> {
883 match transport {
884 McpTransport::Stdio => self.run_stdio(),
885 McpTransport::Sse { port } => self.run_sse(port),
886 }
887 }
888
889 pub fn watch_presets(&self) -> Result<notify::RecommendedWatcher> {
896 use notify::{Event, EventKind, RecursiveMode, Watcher};
897
898 let shared = Arc::clone(&self.shared);
899
900 let mut watcher = notify::recommended_watcher(move |res: notify::Result<Event>| {
901 if let Ok(event) = res {
902 if !matches!(event.kind, EventKind::Modify(_) | EventKind::Create(_)) {
903 return;
904 }
905 for path in &event.paths {
906 if path.extension().and_then(|e| e.to_str()) != Some("toml") {
907 continue;
908 }
909 match std::fs::read_to_string(path) {
910 Ok(toml_str) => {
911 match PresetParser::parse(&toml_str) {
913 Ok(_) => {
914 if let Ok(mut pending) = shared.pending_preset.lock() {
916 *pending = Some(toml_str);
917 }
918 eprintln!("[sqz-mcp] preset change detected: {}", path.display());
919 }
920 Err(e) => {
921 eprintln!("[sqz-mcp] invalid preset TOML in {}: {e}", path.display());
923 }
924 }
925 }
926 Err(e) => eprintln!("[sqz-mcp] preset file read error: {e}"),
927 }
928 }
929 }
930 })
931 .map_err(|e| SqzError::Other(format!("watcher init error: {e}")))?;
932
933 watcher
934 .watch(&self.preset_dir, RecursiveMode::NonRecursive)
935 .map_err(|e| SqzError::Other(format!("watcher watch error: {e}")))?;
936
937 Ok(watcher)
938 }
939
940 fn run_stdio(mut self) -> Result<()> {
943 let stdin = std::io::stdin();
944 let stdout = std::io::stdout();
945 let mut out = stdout.lock();
946 let mut input = stdin.lock();
947 let mut buf = Vec::new();
948
949 while read_line_bytes(&mut input, &mut buf)
950 .map_err(|e| SqzError::Other(format!("stdin read error: {e}")))?
951 {
952 let response = match std::str::from_utf8(&buf) {
953 Ok(line) if line.trim().is_empty() => continue,
954 Ok(line) => match self.handle_jsonrpc_line(line) {
956 Some(response) => response,
957 None => continue,
958 },
959 Err(e) => JsonRpcResponse::err(None, -32700, format!("parse error: {e}")),
960 };
961 let serialized = serde_json::to_string(&response)
962 .unwrap_or_else(|_| r#"{"jsonrpc":"2.0","error":{"code":-32700,"message":"serialize error"}}"#.to_string());
963 writeln!(out, "{serialized}")
964 .map_err(|e| SqzError::Other(format!("stdout write error: {e}")))?;
965 out.flush()
966 .map_err(|e| SqzError::Other(format!("stdout flush error: {e}")))?;
967 }
968 Ok(())
969 }
970
971 fn run_sse(mut self, port: u16) -> Result<()> {
972 use std::net::TcpListener;
973 use std::io::BufReader;
974
975 let listener = TcpListener::bind(format!("127.0.0.1:{port}"))
976 .map_err(|e| SqzError::Other(format!("SSE bind error on port {port}: {e}")))?;
977 eprintln!("[sqz-mcp] SSE server listening on http://127.0.0.1:{port}");
978
979 for stream in listener.incoming() {
980 match stream {
981 Ok(mut stream) => {
982 let mut reader = BufReader::new(stream.try_clone()
983 .map_err(|e| SqzError::Other(format!("stream clone error: {e}")))?);
984 let mut request_line = String::new();
985 let _ = reader.read_line(&mut request_line);
986
987 let mut content_length = 0usize;
989 loop {
990 let mut header = String::new();
991 let _ = reader.read_line(&mut header);
992 if header == "\r\n" || header.is_empty() {
993 break;
994 }
995 let lower = header.to_lowercase();
996 if lower.starts_with("content-length:") {
997 if let Some(v) = lower.split(':').nth(1) {
998 content_length = v.trim().parse().unwrap_or(0);
999 }
1000 }
1001 }
1002
1003 let mut body = vec![0u8; content_length];
1005 use std::io::Read;
1006 let _ = reader.read_exact(&mut body);
1007 let body_str = String::from_utf8_lossy(&body);
1008
1009 let (status, json) = match self.handle_jsonrpc_line(body_str.trim()) {
1013 Some(resp) => ("200 OK", serde_json::to_string(&resp).unwrap_or_default()),
1014 None => ("204 No Content", String::new()),
1015 };
1016
1017 let http_response = format!(
1018 "HTTP/1.1 {}\r\nContent-Type: application/json\r\nContent-Length: {}\r\nAccess-Control-Allow-Origin: *\r\n\r\n{}",
1019 status,
1020 json.len(),
1021 json
1022 );
1023 let _ = stream.write_all(http_response.as_bytes());
1024 }
1025 Err(e) => eprintln!("[sqz-mcp] connection error: {e}"),
1026 }
1027 }
1028 Ok(())
1029 }
1030
1031 fn handle_jsonrpc_line(&mut self, line: &str) -> Option<JsonRpcResponse> {
1032 self.apply_pending_preset();
1034
1035 let req: JsonRpcRequest = match serde_json::from_str(line) {
1036 Ok(r) => r,
1037 Err(e) => return Some(JsonRpcResponse::err(None, -32700, format!("parse error: {e}"))),
1038 };
1039
1040 if req.id.is_none() {
1044 return None;
1045 }
1046
1047 let id = req.id.clone();
1048 let response = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| match req.method.as_str() {
1049 "tools/list" => {
1050 let intent = req.params
1051 .as_ref()
1052 .and_then(|p| p.get("intent"))
1053 .and_then(|v| v.as_str())
1054 .map(|s| s.to_string());
1055
1056 match self.list_tools(intent.as_deref()) {
1057 Ok(tools) => {
1058 let tool_list: Vec<Value> = tools.iter().map(|t| {
1059 let mut tool_json = serde_json::json!({
1070 "name": t.id,
1071 "description": t.description,
1072 "inputSchema": t.input_schema,
1073 "sqz:transforms": t.compression_transforms,
1074 });
1075 if !t.output_schema.is_null() {
1076 if let Some(obj) = tool_json.as_object_mut() {
1077 obj.insert(
1078 "outputSchema".to_string(),
1079 t.output_schema.clone(),
1080 );
1081 }
1082 }
1083 tool_json
1084 }).collect();
1085 JsonRpcResponse::ok(req.id, serde_json::json!({ "tools": tool_list }))
1086 }
1087 Err(e) => JsonRpcResponse::err(req.id, -32603, e.to_string()),
1088 }
1089 }
1090
1091 "tools/call" => {
1092 let params = match req.params {
1093 Some(p) => p,
1094 None => return JsonRpcResponse::err(req.id, -32602, "missing params"),
1095 };
1096 let tool_id = match params.get("name").and_then(|v| v.as_str()) {
1097 Some(id) => id.to_string(),
1098 None => return JsonRpcResponse::err(req.id, -32602, "missing params.name"),
1099 };
1100 let input = params.get("arguments").cloned().unwrap_or(Value::Null);
1101 let intent = params.get("intent").and_then(|v| v.as_str()).map(|s| s.to_string());
1102
1103 let call_req = ToolCallRequest { tool_id, input, intent };
1104 match self.handle_tool_call(call_req) {
1105 Ok(resp) => JsonRpcResponse::ok(req.id, serde_json::json!({
1106 "content": [{ "type": "text", "text": resp.output }],
1107 "tokens_original": resp.tokens_original,
1108 "tokens_compressed": resp.tokens_compressed,
1109 })),
1110 Err(e) => JsonRpcResponse::err(req.id, -32603, e.to_string()),
1111 }
1112 }
1113
1114 "initialize" => {
1115 JsonRpcResponse::ok(req.id, serde_json::json!({
1126 "protocolVersion": "2024-11-05",
1127 "capabilities": {
1128 "tools": { "listChanged": false }
1129 },
1130 "serverInfo": { "name": "sqz-mcp", "version": env!("CARGO_PKG_VERSION") }
1131 }))
1132 }
1133
1134 "ping" => JsonRpcResponse::ok(req.id, serde_json::json!({})),
1135
1136 _ => JsonRpcResponse::err(req.id, -32601, format!("method not found: {}", req.method)),
1137 }));
1138 Some(response.unwrap_or_else(|_| {
1139 self.engine.clear_poison();
1140 JsonRpcResponse::err(id, -32603, format!("internal error: {} panicked", req.method))
1141 }))
1142 }
1143}
1144
1145fn read_line_bytes(reader: &mut impl BufRead, buf: &mut Vec<u8>) -> std::io::Result<bool> {
1147 buf.clear();
1148 if reader.read_until(b'\n', buf)? == 0 {
1149 return Ok(false);
1150 }
1151 if buf.last() == Some(&b'\n') {
1152 buf.pop();
1153 if buf.last() == Some(&b'\r') {
1154 buf.pop();
1155 }
1156 }
1157 Ok(true)
1158}
1159
1160pub fn default_tool_definitions() -> Vec<ToolDefinition> {
1179 vec![
1180 ToolDefinition {
1181 id: "compress".to_string(),
1182 name: "Compress Text".to_string(),
1183 description: "Compress text or JSON you already have, such as a \
1184 long tool result, through the sqz pipeline. It reads no files \
1185 and runs no commands."
1186 .to_string(),
1187 input_schema: serde_json::json!({
1188 "type": "object",
1189 "properties": {
1190 "text": {
1191 "type": "string",
1192 "description": "Text or JSON to compress."
1193 }
1194 },
1195 "required": ["text"]
1196 }),
1197 compression_transforms: vec![
1198 "sha256_cache: repeat inputs within the session return a ~13-token §ref:HASH§ token".to_string(),
1199 "ast_extract: recognised source code collapses to signatures only".to_string(),
1200 "ansi_strip: removes color/formatting codes".to_string(),
1201 "condense: repeated identical lines collapsed to max 3 occurrences".to_string(),
1202 "git_diff_fold: diff output has unchanged context lines folded".to_string(),
1203 "log_fold: repeated log lines with timestamps folded to [xN]".to_string(),
1204 "path_shorten: common path prefixes replaced with ~/".to_string(),
1205 "truncate_strings: strings > 500 chars are truncated with '...'".to_string(),
1206 "entropy_truncate: low-information segments are dropped and marked; the output then ends with the expand prefix of the full original".to_string(),
1207 "safe_fallback: error/warning lines always preserved verbatim".to_string(),
1208 "preservation_verifier: path-like and identifier tokens are \
1209 checked for byte-exact survival; compression is discarded if \
1210 coverage drops below 85%"
1211 .to_string(),
1212 ],
1213 ..Default::default()
1214 },
1215 ToolDefinition {
1221 id: "passthrough".to_string(),
1222 name: "Passthrough (No Compression)".to_string(),
1223 description: "Return the text unchanged, for when you need it \
1224 raw rather than compressed."
1225 .to_string(),
1226 input_schema: serde_json::json!({
1227 "type": "object",
1228 "properties": {
1229 "text": {
1230 "type": "string",
1231 "description": "Text to return as is."
1232 }
1233 },
1234 "required": ["text"]
1235 }),
1236 compression_transforms: vec![
1237 "none: input is returned byte-for-byte".to_string(),
1238 ],
1239 ..Default::default()
1240 },
1241 ToolDefinition {
1247 id: "expand".to_string(),
1248 name: "Expand Dedup Ref".to_string(),
1249 description: "Return the original content behind a ref. \
1250 `§ref:HASH§` stands for content sent earlier in the session; \
1251 `§ref:HASH:L40-80§` for lines 40-80 of content sent in full, \
1252 and expands to just those lines."
1253 .to_string(),
1254 input_schema: serde_json::json!({
1255 "type": "object",
1256 "properties": {
1257 "prefix": {
1258 "type": "string",
1259 "description": "The ref token pasted whole, or its hex prefix."
1260 }
1261 },
1262 "required": ["prefix"]
1263 }),
1264 compression_transforms: vec![
1265 "none: returns cached original bytes".to_string(),
1266 ],
1267 ..Default::default()
1268 },
1269 ToolDefinition {
1275 id: "sqz_recall".to_string(),
1276 name: "Recall (Search Session Memory)".to_string(),
1277 description: "Search everything sqz has seen, this session and \
1278 earlier, before re-running a command or re-reading a file. \
1279 Pass a hit's ref to `expand` for the full text."
1280 .to_string(),
1281 input_schema: serde_json::json!({
1282 "type": "object",
1283 "properties": {
1284 "query": {
1285 "type": "string",
1286 "description": "Search terms; all must match."
1287 },
1288 "limit": {
1289 "type": "integer",
1290 "description": "Maximum number of hits, 1-25.",
1291 "default": 5
1292 }
1293 },
1294 "required": ["query"]
1295 }),
1296 compression_transforms: vec![
1297 "bm25_rank: hits ordered by FTS5 relevance".to_string(),
1298 "snippet: ~24 tokens of context per hit with >>match<< markers".to_string(),
1299 ],
1300 ..Default::default()
1301 },
1302 ToolDefinition {
1321 id: "sqz_read_file".to_string(),
1322 name: "Read File (Compressed)".to_string(),
1323 description: "Read a file faithfully, stripping only ANSI \
1324 escapes. Prefer it over the built-in read for files over \
1325 2KB or ones you may read again: a repeat or ranged re-read \
1326 returns a short ref (see `expand`) and a re-read after an \
1327 edit returns only the changed lines."
1328 .to_string(),
1329 input_schema: serde_json::json!({
1330 "type": "object",
1331 "properties": {
1332 "path": {
1333 "type": "string",
1334 "description": "File path, absolute or relative to the working directory."
1335 },
1336 "offset": {
1337 "type": "integer",
1338 "description": "First line to return, 1-based."
1339 },
1340 "limit": {
1341 "type": "integer",
1342 "description": "Number of lines to return."
1343 },
1344 "max_bytes": {
1345 "type": "integer",
1346 "description": "Byte cap, cut at a line boundary; 0 for no cap. When cut, the header gives continue_with_offset.",
1347 "default": 262144
1348 }
1349 },
1350 "required": ["path"]
1351 }),
1352 compression_transforms: vec![
1353 "sha256_cache: repeat reads of unchanged content return a ~13-token §ref:HASH§ token".to_string(),
1354 "slice_ref: a line range of content already read in full returns §ref:HASH:L<a>-<b>§".to_string(),
1355 "delta: a re-read after a small edit returns only the changed lines".to_string(),
1356 "ansi_strip: removes color/formatting codes".to_string(),
1357 "max_bytes: output past the cap is cut at a line boundary and the header says where to continue".to_string(),
1358 "lossless: no lines are folded or summarized".to_string(),
1359 ],
1360 ..Default::default()
1361 },
1362 ToolDefinition {
1363 id: "sqz_list_dir".to_string(),
1364 name: "List Directory (Compressed)".to_string(),
1365 description: "List a directory, skipping hidden entries and \
1366 node_modules, target, dist, build, vendor and __pycache__. \
1367 Prefer it over `ls -la` for a project layout."
1368 .to_string(),
1369 input_schema: serde_json::json!({
1370 "type": "object",
1371 "properties": {
1372 "path": {
1373 "type": "string",
1374 "description": "Directory to list.",
1375 "default": "."
1376 },
1377 "max_depth": {
1378 "type": "integer",
1379 "description": "Levels to recurse; 1 lists direct children only.",
1380 "default": 1
1381 },
1382 "max_entries": {
1383 "type": "integer",
1384 "description": "Entry cap; 0 for no cap.",
1385 "default": 1000
1386 }
1387 }
1388 }),
1389 compression_transforms: vec![
1390 "skip_bulk_dirs: .git, node_modules, target, dist, build, vendor, __pycache__ omitted".to_string(),
1391 "max_entries: listing stops at the cap (default 1000); the header says so".to_string(),
1392 "sha256_cache: repeat listings dedupe via §ref§".to_string(),
1393 "lossless: every listed entry survives verbatim".to_string(),
1394 ],
1395 ..Default::default()
1396 },
1397 ToolDefinition {
1398 id: "sqz_grep".to_string(),
1399 name: "Grep Files (Compressed)".to_string(),
1400 description: "Search files for a literal string or regex, \
1401 returning `path:lineno:text` lines. Prefer it over the \
1402 built-in grep when a search may match more than a few \
1403 lines: output is capped and a repeat search returns a \
1404 short ref."
1405 .to_string(),
1406 input_schema: serde_json::json!({
1407 "type": "object",
1408 "properties": {
1409 "pattern": {
1410 "type": "string",
1411 "description": "Text to find, a literal substring unless `regex` is true."
1412 },
1413 "path": {
1414 "type": "string",
1415 "description": "File or directory to search.",
1416 "default": "."
1417 },
1418 "regex": {
1419 "type": "boolean",
1420 "description": "Treat `pattern` as a regex.",
1421 "default": false
1422 },
1423 "max_matches": {
1424 "type": "integer",
1425 "description": "Stop after this many matches.",
1426 "default": 200
1427 },
1428 "max_line_chars": {
1429 "type": "integer",
1430 "description": "Clip longer match lines around the hit; 0 for no cap.",
1431 "default": 400
1432 },
1433 "max_bytes": {
1434 "type": "integer",
1435 "description": "Stop at this many bytes of output; 0 for no cap.",
1436 "default": 100000
1437 }
1438 },
1439 "required": ["pattern"]
1440 }),
1441 compression_transforms: vec![
1442 "max_matches: search stops after the cap (default 200); the header reports the match count".to_string(),
1443 "max_line_chars: a match line over 400 chars is clipped around the hit and marked".to_string(),
1444 "max_bytes: search stops at 100000 bytes of output; the header says so".to_string(),
1445 "sha256_cache: repeat searches dedupe via §ref§".to_string(),
1446 "lossless: match lines under the caps survive verbatim".to_string(),
1447 ],
1448 ..Default::default()
1449 },
1450 ]
1451}
1452
1453fn estimate_tokens(text: &str) -> u32 {
1457 ((text.len() as f64) / 4.0).ceil() as u32
1458}
1459
1460const READ_DEFAULT_MAX_BYTES: usize = 256 * 1024;
1464const READ_HARD_LIMIT_BYTES: u64 = 64 * 1024 * 1024;
1466const GREP_DEFAULT_MAX_BYTES: usize = 100_000;
1468const GREP_DEFAULT_MAX_LINE_CHARS: usize = 400;
1471const LIST_DEFAULT_MAX_ENTRIES: usize = 1000;
1473const REF_TOKENS: u32 = 13;
1475
1476fn cap_from(input: &serde_json::Value, key: &str, default: usize) -> usize {
1479 match input.get(key).and_then(|v| v.as_u64()) {
1480 Some(0) => usize::MAX,
1481 Some(v) => usize::try_from(v).unwrap_or(usize::MAX),
1482 None => default,
1483 }
1484}
1485
1486fn truncate_to_lines(text: &str, max_bytes: usize) -> Option<&str> {
1491 if text.len() <= max_bytes {
1492 return None;
1493 }
1494 let mut cut = max_bytes;
1495 while !text.is_char_boundary(cut) {
1496 cut -= 1;
1497 }
1498 match text[..cut].rfind('\n') {
1499 Some(nl) => Some(&text[..=nl]),
1500 None => Some(&text[..cut]),
1501 }
1502}
1503
1504fn clip_line(line: &str, hit: (usize, usize), max_chars: usize) -> Option<String> {
1508 if line.len() <= max_chars {
1509 return None;
1510 }
1511 let total = line.chars().count();
1512 if total <= max_chars {
1513 return None;
1514 }
1515 let (hit_start, hit_end) = hit;
1516 let hit_start_char = line[..hit_start].chars().count();
1517 let hit_chars = line[hit_start..hit_end].chars().count();
1518 let lead = max_chars.saturating_sub(hit_chars) / 2;
1519 let mut from = hit_start_char.saturating_sub(lead);
1520 let to = (from + max_chars).min(total);
1521 from = to.saturating_sub(max_chars);
1522 let window: String = line.chars().skip(from).take(to - from).collect();
1523 Some(format!(
1524 "{}{}{} [line clipped: {total} chars]",
1525 if from > 0 { "…" } else { "" },
1526 window,
1527 if to < total { "…" } else { "" },
1528 ))
1529}
1530
1531fn line_range(text: &str, start: usize, end: usize) -> &str {
1534 let mut line = 1;
1535 let mut begin = None;
1536 for (i, ch) in text.char_indices() {
1537 if line == start && begin.is_none() {
1538 begin = Some(i);
1539 }
1540 if ch == '\n' {
1541 if line == end {
1542 return &text[begin.unwrap_or(i)..=i];
1543 }
1544 line += 1;
1545 }
1546 }
1547 begin.map(|b| &text[b..]).unwrap_or("")
1548}
1549
1550fn list_dir_recursive(
1563 root: &std::path::Path,
1564 current: &std::path::Path,
1565 depth: usize,
1566 out: &mut Listing,
1567) -> Result<()> {
1568 if depth > out.max_depth || out.capped {
1569 return Ok(());
1570 }
1571
1572 let entries = match std::fs::read_dir(current) {
1573 Ok(e) => e,
1574 Err(e) => {
1575 return Err(SqzError::Other(format!(
1576 "sqz_list_dir: could not read '{}': {e}",
1577 current.display()
1578 )))
1579 }
1580 };
1581
1582 let mut sorted: Vec<_> = entries
1583 .filter_map(|e| e.ok())
1584 .collect();
1585 sorted.sort_by_key(|e| e.file_name());
1588
1589 for entry in sorted {
1590 if out.capped {
1591 break;
1592 }
1593 let name = entry.file_name();
1594 let name_str = name.to_string_lossy();
1595
1596 if name_str.starts_with('.') {
1598 continue;
1599 }
1600 if matches!(
1601 name_str.as_ref(),
1602 "node_modules" | "target" | "dist" | "build" | "__pycache__"
1603 | "vendor" | ".next" | ".nuxt"
1604 ) {
1605 continue;
1606 }
1607
1608 let path = entry.path();
1609 let rel = path.strip_prefix(root).unwrap_or(&path);
1610
1611 let type_char = match entry.file_type() {
1612 Ok(ft) if ft.is_dir() => "d",
1613 Ok(ft) if ft.is_symlink() => "l",
1614 Ok(_) => "f",
1615 Err(_) => "?",
1616 };
1617
1618 if out.lines.len() >= out.max_entries {
1619 out.capped = true;
1620 break;
1621 }
1622 out.lines.push(format!("{type_char} {}", rel.display()));
1623
1624 if entry.file_type().map(|ft| ft.is_dir()).unwrap_or(false) && depth < out.max_depth {
1625 let _ = list_dir_recursive(root, &path, depth + 1, out);
1626 }
1627 }
1628
1629 Ok(())
1630}
1631
1632struct Listing {
1635 lines: Vec<String>,
1636 max_depth: usize,
1637 max_entries: usize,
1638 capped: bool,
1639}
1640
1641fn grep_walk(
1649 root: &std::path::Path,
1650 needle: &str,
1651 regex: Option<®ex::Regex>,
1652 limits: &GrepLimits,
1653 out: &mut GrepResults,
1654) -> Result<()> {
1655 if out.full(limits) {
1656 return Ok(());
1657 }
1658
1659 if !root.exists() {
1660 return Err(SqzError::Other(format!(
1661 "sqz_grep: path '{}' does not exist",
1662 root.display()
1663 )));
1664 }
1665
1666 if root.is_file() {
1668 grep_one_file(root, needle, regex, limits, out)?;
1669 return Ok(());
1670 }
1671
1672 let entries = match std::fs::read_dir(root) {
1675 Ok(e) => e,
1676 Err(e) => {
1677 return Err(SqzError::Other(format!(
1678 "sqz_grep: could not read '{}': {e}",
1679 root.display()
1680 )))
1681 }
1682 };
1683
1684 let mut sorted: Vec<_> = entries.filter_map(|e| e.ok()).collect();
1685 sorted.sort_by_key(|e| e.file_name());
1686
1687 for entry in sorted {
1688 if out.full(limits) {
1689 break;
1690 }
1691 let name = entry.file_name();
1692 let name_str = name.to_string_lossy();
1693 if name_str.starts_with('.') {
1694 continue;
1695 }
1696 if matches!(
1697 name_str.as_ref(),
1698 "node_modules" | "target" | "dist" | "build" | "__pycache__"
1699 | "vendor" | ".next" | ".nuxt"
1700 ) {
1701 continue;
1702 }
1703
1704 let path = entry.path();
1705 match entry.file_type() {
1706 Ok(ft) if ft.is_dir() => {
1707 grep_walk(&path, needle, regex, limits, out)?;
1708 }
1709 Ok(_) => {
1710 grep_one_file(&path, needle, regex, limits, out)?;
1711 }
1712 Err(_) => continue,
1713 }
1714 }
1715
1716 Ok(())
1717}
1718
1719struct GrepLimits {
1721 max_matches: usize,
1722 max_line_chars: usize,
1723 max_bytes: usize,
1724}
1725
1726#[derive(Default)]
1728struct GrepResults {
1729 lines: Vec<String>,
1730 bytes: usize,
1731 clipped: usize,
1732 byte_capped: bool,
1733}
1734
1735impl GrepResults {
1736 fn full(&self, limits: &GrepLimits) -> bool {
1737 self.byte_capped || self.lines.len() >= limits.max_matches
1738 }
1739
1740 fn push(&mut self, line: String, clipped: bool, limits: &GrepLimits) {
1744 let cost = line.len() + 1;
1745 if !self.lines.is_empty() && self.bytes.saturating_add(cost) > limits.max_bytes {
1746 self.byte_capped = true;
1747 return;
1748 }
1749 self.bytes += cost;
1750 self.lines.push(line);
1751 if clipped {
1752 self.clipped += 1;
1753 }
1754 }
1755}
1756
1757fn grep_one_file(
1761 file: &std::path::Path,
1762 needle: &str,
1763 regex: Option<®ex::Regex>,
1764 limits: &GrepLimits,
1765 out: &mut GrepResults,
1766) -> Result<()> {
1767 if let Ok(meta) = std::fs::metadata(file) {
1770 if meta.len() > 50 * 1024 * 1024 {
1771 return Ok(());
1772 }
1773 }
1774
1775 let bytes = match std::fs::read(file) {
1776 Ok(b) => b,
1777 Err(_) => return Ok(()), };
1779 let text = match std::str::from_utf8(&bytes) {
1780 Ok(t) => t,
1781 Err(_) => return Ok(()), };
1783
1784 for (lineno, line) in text.lines().enumerate() {
1785 if out.full(limits) {
1786 break;
1787 }
1788 let hit = match regex {
1789 Some(r) => r.find(line).map(|m| (m.start(), m.end())),
1790 None => line.find(needle).map(|i| (i, i + needle.len())),
1791 };
1792 if let Some(hit) = hit {
1793 let clipped = clip_line(line, hit, limits.max_line_chars);
1794 let was_clipped = clipped.is_some();
1795 let shown = clipped.unwrap_or_else(|| line.to_string());
1796 out.push(
1797 format!("{}:{}:{}", file.display(), lineno + 1, shown),
1798 was_clipped,
1799 limits,
1800 );
1801 }
1802 }
1803
1804 Ok(())
1805}
1806
1807#[cfg(test)]
1810mod tests {
1811 use super::*;
1812 use std::time::{Duration, Instant};
1813 use tempfile::TempDir;
1814
1815 impl McpServer {
1821 pub(crate) fn handle_jsonrpc_line_unwrap(&mut self, line: &str) -> JsonRpcResponse {
1822 self.handle_jsonrpc_line(line)
1823 .expect("expected response; got None (notification). Use handle_jsonrpc_line directly if that's intended.")
1824 }
1825 }
1826
1827 fn make_server() -> (McpServer, TempDir) {
1828 let dir = TempDir::new().expect("tempdir");
1829 let store_path = dir.path().join("sessions.db");
1834 let server = McpServer::new_with_store(dir.path(), &store_path)
1835 .expect("McpServer::new_with_store");
1836 (server, dir)
1837 }
1838
1839 thread_local! {
1840 static PANIC_NEXT_ENGINE_CALL: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
1841 }
1842
1843 pub(crate) fn panic_next_engine_call() {
1845 PANIC_NEXT_ENGINE_CALL.with(|p| p.set(true));
1846 }
1847
1848 pub(crate) fn maybe_panic() {
1849 if PANIC_NEXT_ENGINE_CALL.with(|p| p.replace(false)) {
1850 panic!("injected engine panic");
1851 }
1852 }
1853
1854 #[test]
1859 fn test_handle_tool_call_compresses_output() {
1860 let (mut server, _dir) = make_server();
1861
1862 let input = serde_json::json!({
1863 "status": "ok",
1864 "data": {
1865 "id": 1,
1866 "name": "test",
1867 "debug_info": null,
1868 "trace_id": null,
1869 "metadata": {
1870 "internal_id": "abc123",
1871 "created_at": "2025-01-01T00:00:00Z"
1872 },
1873 "items": ["a", "b", "c", "d", "e", "f", "g", "h"]
1874 }
1875 });
1876
1877 let req = ToolCallRequest {
1878 tool_id: "compress".to_string(),
1879 input: input.clone(),
1880 intent: None,
1881 };
1882
1883 let resp = server.handle_tool_call(req).expect("handle_tool_call");
1884 assert_eq!(resp.tool_id, "compress");
1885 assert!(!resp.output.is_empty(), "output should not be empty");
1886 assert!(resp.tokens_original > 0, "tokens_original should be > 0");
1887 }
1888
1889 #[test]
1892 fn test_handle_tool_call_preserves_tool_id() {
1893 let (mut server, _dir) = make_server();
1894 let req = ToolCallRequest {
1895 tool_id: "compress".to_string(),
1896 input: serde_json::json!({ "text": "ls -la output here" }),
1897 intent: None,
1898 };
1899 let resp = server.handle_tool_call(req).expect("handle_tool_call");
1900 assert_eq!(resp.tool_id, "compress");
1901 }
1902
1903 #[test]
1904 fn test_compress_tool_compresses_the_text_field() {
1905 let (mut server, _dir) = make_server();
1906 let mut text: String = (0..300)
1907 .map(|i| format!("2026-10-01T12:00:{:02}Z INFO worker {i} finished job {i} in {} ms\n", i % 60, i % 97))
1908 .collect();
1909 text.push_str("ERROR worker 299 failed: disk full\n");
1910 let resp = server
1911 .handle_tool_call(ToolCallRequest {
1912 tool_id: "compress".to_string(),
1913 input: serde_json::json!({ "text": text }),
1914 intent: None,
1915 })
1916 .unwrap();
1917 let head: String = resp.output.chars().take(200).collect();
1918 assert!(!resp.output.starts_with("TOON:{text:"), "{head}");
1919 assert!(resp.output.contains("disk full"), "{head}");
1920 }
1921
1922 #[test]
1923 fn test_compress_tool_names_the_expand_prefix_when_it_drops_segments() {
1924 let (mut server, _dir) = make_server();
1925 let mut text = String::from("run marker for the compress spill test\n\n");
1926 for i in 0..12u8 {
1927 let word: String = std::iter::repeat((b'a' + i) as char).take(10).collect();
1928 text.push_str(&format!("{} filler-{i}\n\n", format!("{word} ").repeat(30)));
1929 }
1930 text.push_str("The quick brown fox jumps over the lazy dog by silver rivers today.\n");
1931 let call = |server: &mut McpServer, tool: &str, input: serde_json::Value| {
1932 server
1933 .handle_tool_call(ToolCallRequest { tool_id: tool.to_string(), input, intent: None })
1934 .unwrap()
1935 };
1936
1937 let resp = call(&mut server, "compress", serde_json::json!({ "text": text }));
1938 let hash = sqz_engine::CacheManager::sha256_hex(text.as_bytes());
1939 assert!(resp.output.contains("low-information segments omitted"), "{}", resp.output);
1940 assert!(
1941 resp.output.ends_with(&format!("\n[full output: call expand with prefix \"{}\"]", &hash[..16])),
1942 "{}",
1943 resp.output
1944 );
1945 assert_eq!(resp.tokens_compressed, server.engine.count_tokens(&resp.output));
1946
1947 let expanded = call(&mut server, "expand", serde_json::json!({ "prefix": &hash[..16] }));
1948 assert_eq!(expanded.output, format!("[sqz:expand hash={hash}]\n{text}"));
1949
1950 let compress = default_tool_definitions().into_iter().find(|t| t.id == "compress").unwrap();
1951 assert!(compress.compression_transforms.iter().any(|t| t.starts_with("entropy_truncate:")));
1952 }
1953
1954 fn high_risk_texts() -> Vec<(&'static str, String)> {
1955 let hex = |s: &str| sqz_engine::CacheManager::sha256_hex(s.as_bytes());
1956 let frames: String = (0..40)
1957 .map(|i| format!(" {i:>2}: myapp::scheduler::step_{i}\n at ./src/scheduler/mod_{}.rs:{}:5\n", i % 7, 40 + 7 * i))
1958 .collect();
1959 let tables: String = ["customers", "orders", "refunds", "payments", "shipments", "invoices"]
1960 .iter()
1961 .map(|t| format!("CREATE TABLE {t} (\n id BIGSERIAL PRIMARY KEY,\n status TEXT NOT NULL DEFAULT 'pending',\n created_at TIMESTAMPTZ NOT NULL DEFAULT now()\n);\nCREATE INDEX idx_{t}_created_at ON {t} (created_at);\n"))
1962 .collect();
1963 let pem: String = (0..20).map(|i| hex(&format!("pem {i}")) + "\n").collect();
1964 let creds: String = ["stripe", "sendgrid", "github", "datadog", "sentry"]
1965 .iter()
1966 .map(|s| format!("{s}_api_key: sk_test_{}\n{s}_endpoint: https://{s}.example.com/v1\n", &hex(s)[..32]))
1967 .collect();
1968 vec![
1969 ("panic", format!("thread 'main' panicked at src/scheduler/queue.rs:217:31:\ncalled `Option::unwrap()` on a `None` value\nstack backtrace:\n{frames}")),
1970 ("migration", format!("BEGIN;\n{tables}ALTER TABLE refunds ADD COLUMN reason TEXT;\nCOMMIT;\n")),
1971 ("keys", format!("-----BEGIN RSA PRIVATE KEY-----\n{pem}-----END RSA PRIVATE KEY-----\n{creds}")),
1972 ]
1973 }
1974
1975 #[test]
1976 fn test_compress_tool_returns_high_risk_text_unchanged() {
1977 let (mut server, _dir) = make_server();
1978 for (name, text) in high_risk_texts() {
1979 assert!(text.len() > 500, "{name}");
1980 for run in ["first", "second"] {
1981 let resp = server
1982 .handle_tool_call(ToolCallRequest {
1983 tool_id: "compress".to_string(),
1984 input: serde_json::json!({ "text": text }),
1985 intent: None,
1986 })
1987 .unwrap();
1988 assert_eq!(resp.output, text, "{name}, {run} call");
1989 }
1990 let stored = server.engine.cache_manager().check_dedup_with_meta(text.as_bytes()).unwrap();
1991 assert!(stored.is_none(), "{name} was stored");
1992 }
1993 }
1994
1995 #[test]
1998 fn test_list_tools_no_intent_returns_all() {
1999 let (server, _dir) = make_server();
2000 let tools = server.list_tools(None).expect("list_tools");
2001 assert_eq!(tools.len(), default_tool_definitions().len());
2002 }
2003
2004 #[test]
2015 fn test_list_tools_with_intent_filters() {
2016 let (server, _dir) = make_server();
2017 let registered = default_tool_definitions().len();
2018
2019 let tools = server
2020 .list_tools(Some("compress arbitrary text through the sqz pipeline"))
2021 .expect("list_tools with intent should not error");
2022 assert!(
2023 tools.len() <= registered,
2024 "filtered list must not exceed registered count ({registered})"
2025 );
2026
2027 let tools = server.list_tools(Some("")).expect("empty intent = all tools");
2028 assert_eq!(
2029 tools.len(),
2030 registered,
2031 "empty intent is treated as `no intent` and returns every tool"
2032 );
2033 }
2034
2035 #[test]
2038 fn test_tool_selector_latency_under_500ms() {
2039 let (server, _dir) = make_server();
2040
2041 let start = Instant::now();
2042 for _ in 0..10 {
2043 let _ = server.list_tools(Some("search for files matching a pattern"));
2044 }
2045 let elapsed = start.elapsed();
2046
2047 assert!(
2048 elapsed < Duration::from_millis(500),
2049 "10 tool selections took {:?}, expected < 500ms",
2050 elapsed
2051 );
2052 }
2053
2054 #[test]
2057 fn test_preset_hot_reload_latency() {
2058 let dir = TempDir::new().expect("tempdir");
2059 let store = TempDir::new().expect("tempdir");
2060 let server = McpServer::new_with_store(dir.path(), &store.path().join("sessions.db"))
2061 .expect("McpServer::new_with_store");
2062
2063 let _watcher = server.watch_presets().expect("watch_presets");
2065
2066 let preset_path = dir.path().join("test.toml");
2068 let toml_content = r#"
2069[preset]
2070name = "hot-reload-test"
2071version = "1.0"
2072
2073[compression]
2074stages = []
2075
2076[tool_selection]
2077max_tools = 5
2078similarity_threshold = 0.3
2079
2080[budget]
2081warning_threshold = 0.70
2082ceiling_threshold = 0.85
2083default_window_size = 200000
2084
2085[terse_mode]
2086enabled = false
2087level = "moderate"
2088
2089[model]
2090family = "anthropic"
2091primary = "claude-sonnet-4-20250514"
2092complexity_threshold = 0.4
2093"#;
2094 std::fs::write(&preset_path, toml_content).expect("write preset");
2095
2096 let deadline = Instant::now() + Duration::from_secs(2);
2098 while Instant::now() < deadline {
2099 std::thread::sleep(Duration::from_millis(50));
2100 if let Ok(guard) = server.shared.pending_preset.lock() {
2102 if guard.is_some() {
2103 break;
2104 }
2105 }
2106 }
2107
2108 let has_pending = server.shared.pending_preset.lock()
2110 .map(|g| g.is_some())
2111 .unwrap_or(false);
2112 assert!(has_pending, "preset should have been hot-reloaded within 2 seconds");
2113 }
2114
2115 #[test]
2118 fn test_invalid_toml_keeps_previous_preset() {
2119 let dir = TempDir::new().expect("tempdir");
2120 let store = TempDir::new().expect("tempdir");
2121 let server = McpServer::new_with_store(dir.path(), &store.path().join("sessions.db"))
2122 .expect("McpServer::new_with_store");
2123 let _watcher = server.watch_presets().expect("watch_presets");
2124
2125 let bad_path = dir.path().join("bad.toml");
2127 std::fs::write(&bad_path, "this is not valid toml ][[[").expect("write bad preset");
2128
2129 std::thread::sleep(Duration::from_millis(200));
2131
2132 let has_pending = server.shared.pending_preset.lock()
2134 .map(|g| g.is_some())
2135 .unwrap_or(false);
2136 assert!(!has_pending, "invalid TOML should not be stored as pending preset");
2137
2138 let tools = server.list_tools(None).expect("list_tools after bad preset");
2140 assert!(!tools.is_empty(), "tools should still be available after invalid preset");
2141 }
2142
2143 #[test]
2145 fn test_jsonrpc_initialize() {
2146 let (mut server, _dir) = make_server();
2147 let line = r#"{"jsonrpc":"2.0","id":1,"method":"initialize","params":{}}"#;
2148 let resp = server.handle_jsonrpc_line_unwrap(line);
2149 assert!(resp.error.is_none(), "initialize should not error");
2150 let result = resp.result.expect("initialize should have result");
2151 assert!(result.get("protocolVersion").is_some());
2152 }
2153
2154 #[test]
2159 fn test_initialize_advertises_tools_capability() {
2160 let (mut server, _dir) = make_server();
2161 let line = r#"{"jsonrpc":"2.0","id":1,"method":"initialize","params":{}}"#;
2162 let resp = server.handle_jsonrpc_line_unwrap(line);
2163 let result = resp.result.expect("initialize should have result");
2164
2165 let caps = result.get("capabilities")
2166 .expect("initialize result must include capabilities");
2167 let tools_cap = caps.get("tools")
2168 .expect("capabilities must include 'tools' key");
2169
2170 let tools_obj = tools_cap.as_object()
2172 .expect("'tools' capability must be an object");
2173 assert!(
2176 !tools_obj.is_empty(),
2177 "'tools' capability must not be empty {{}} — some MCP clients \
2178 interpret that as no tools available. Got: {tools_cap:?}"
2179 );
2180 assert!(
2183 tools_obj.contains_key("listChanged"),
2184 "'tools' capability should include listChanged per MCP 2024-11-05 \
2185 spec. Got: {tools_cap:?}"
2186 );
2187 }
2188
2189 #[test]
2191 fn test_jsonrpc_tools_list() {
2192 let (mut server, _dir) = make_server();
2193 let line = r#"{"jsonrpc":"2.0","id":2,"method":"tools/list","params":{}}"#;
2194 let resp = server.handle_jsonrpc_line_unwrap(line);
2195 assert!(resp.error.is_none(), "tools/list should not error");
2196 let result = resp.result.expect("tools/list should have result");
2197 let tools = result.get("tools").expect("result should have tools");
2198 assert!(tools.as_array().map(|a| !a.is_empty()).unwrap_or(false));
2199 }
2200
2201 #[test]
2203 fn test_jsonrpc_tools_call() {
2204 let (mut server, _dir) = make_server();
2205 let line = r#"{"jsonrpc":"2.0","id":3,"method":"tools/call","params":{"name":"compress","arguments":{"text":"lorem ipsum dolor sit amet"}}}"#;
2206 let resp = server.handle_jsonrpc_line_unwrap(line);
2207 assert!(resp.error.is_none(), "tools/call should not error: {:?}", resp.error);
2208 let result = resp.result.expect("tools/call should have result");
2209 assert!(result.get("content").is_some());
2210 }
2211
2212 #[test]
2214 fn test_jsonrpc_compress_panic_returns_the_text() {
2215 let (mut server, _dir) = make_server();
2216 let text = "2026-10-06T12:00:00Z INFO worker finished job in 3 ms\n".repeat(40);
2217 let call = serde_json::json!({
2218 "jsonrpc": "2.0", "id": 1, "method": "tools/call",
2219 "params": { "name": "compress", "arguments": { "text": text } }
2220 })
2221 .to_string();
2222 panic_next_engine_call();
2223 let resp = server.handle_jsonrpc_line_unwrap(&call);
2224 assert!(resp.error.is_none(), "{:?}", resp.error);
2225 assert_eq!(resp.result.unwrap()["content"][0]["text"], text.as_str());
2226
2227 let resp = server.handle_jsonrpc_line_unwrap(&call);
2228 let out = resp.result.unwrap()["content"][0]["text"].as_str().unwrap().to_string();
2229 assert!(out.len() < text.len(), "compression works again: {out}");
2230 }
2231
2232 #[test]
2234 fn test_jsonrpc_tool_panic_returns_an_error() {
2235 let (mut server, dir) = make_server();
2236 let file = dir.path().join("notes.txt");
2237 std::fs::write(&file, "first line of notes\nsecond line of notes\n").unwrap();
2238 let call = serde_json::json!({
2239 "jsonrpc": "2.0", "id": 2, "method": "tools/call",
2240 "params": { "name": "sqz_read_file", "arguments": { "path": file.to_string_lossy() } }
2241 })
2242 .to_string();
2243 panic_next_engine_call();
2244 let resp = server.handle_jsonrpc_line_unwrap(&call);
2245 assert_eq!(resp.id, Some(serde_json::json!(2)));
2246 assert_eq!(resp.error.expect("an error").code, -32603);
2247
2248 let resp = server.handle_jsonrpc_line_unwrap(&call);
2249 assert!(resp.error.is_none(), "{:?}", resp.error);
2250 let text = resp.result.unwrap()["content"][0]["text"].as_str().unwrap().to_string();
2251 assert!(text.contains("second line of notes"), "{text}");
2252 }
2253
2254 #[test]
2256 fn test_jsonrpc_ping() {
2257 let (mut server, _dir) = make_server();
2258 let resp = server.handle_jsonrpc_line_unwrap(r#"{"jsonrpc":"2.0","id":7,"method":"ping"}"#);
2259 assert!(resp.error.is_none(), "ping should not error: {:?}", resp.error);
2260 assert_eq!(resp.result, Some(serde_json::json!({})));
2261 }
2262
2263 #[test]
2265 fn test_jsonrpc_unknown_method() {
2266 let (mut server, _dir) = make_server();
2267 let line = r#"{"jsonrpc":"2.0","id":4,"method":"unknown/method","params":{}}"#;
2268 let resp = server.handle_jsonrpc_line_unwrap(line);
2269 assert!(resp.error.is_some());
2270 assert_eq!(resp.error.unwrap().code, -32601);
2271 }
2272
2273 #[test]
2275 fn test_read_line_bytes_splits_like_lines() {
2276 let text = "a\nb\r\nc\r\r\n\r\n\n\rd\ne";
2277 let mut reader = text.as_bytes();
2278 let mut buf = Vec::new();
2279 let mut got = Vec::new();
2280 while read_line_bytes(&mut reader, &mut buf).unwrap() {
2281 got.push(String::from_utf8(buf.clone()).unwrap());
2282 }
2283 let want: Vec<String> = text.as_bytes().lines().map(|l| l.unwrap()).collect();
2284 assert_eq!(got, want);
2285
2286 let mut reader: &[u8] = b"{\"text\":\"caf\xe9\"}\nnext\n";
2287 assert!(read_line_bytes(&mut reader, &mut buf).unwrap());
2288 assert_eq!(buf, b"{\"text\":\"caf\xe9\"}");
2289 assert!(read_line_bytes(&mut reader, &mut buf).unwrap());
2290 assert_eq!(buf, b"next");
2291 assert!(!read_line_bytes(&mut reader, &mut buf).unwrap());
2292 }
2293
2294 #[test]
2296 fn test_jsonrpc_parse_error() {
2297 let (mut server, _dir) = make_server();
2298 let resp = server.handle_jsonrpc_line_unwrap("not json at all {{{");
2299 assert!(resp.error.is_some());
2300 assert_eq!(resp.error.unwrap().code, -32700);
2301 }
2302
2303 #[test]
2315 fn test_tools_list_outputschema_is_valid_object_or_absent() {
2316 let (mut server, _dir) = make_server();
2317 let line = r#"{"jsonrpc":"2.0","id":1,"method":"tools/list","params":{}}"#;
2318 let resp = server.handle_jsonrpc_line_unwrap(line);
2319 assert!(resp.error.is_none(), "tools/list errored: {:?}", resp.error);
2320
2321 let tools = resp.result
2322 .expect("tools/list must have result")
2323 .get("tools")
2324 .cloned()
2325 .expect("result must have tools array");
2326 let tools = tools.as_array().expect("tools must be an array");
2327 assert!(!tools.is_empty(), "no tools registered");
2328
2329 for tool in tools {
2330 let name = tool.get("name").and_then(|v| v.as_str()).unwrap_or("?");
2331
2332 let input_type = tool
2334 .get("inputSchema")
2335 .and_then(|s| s.get("type"))
2336 .and_then(|t| t.as_str());
2337 assert_eq!(
2338 input_type,
2339 Some("object"),
2340 "tool {name}: inputSchema.type must be \"object\", got {input_type:?}"
2341 );
2342
2343 if let Some(out) = tool.get("outputSchema") {
2348 if !out.is_null() {
2352 let out_type = out.get("type").and_then(|t| t.as_str());
2353 assert_eq!(
2354 out_type,
2355 Some("object"),
2356 "tool {name}: outputSchema.type must be \"object\" \
2357 per MCP spec; got {out_type:?}. This is the \
2358 exact bug OpenCode reported in issue #5."
2359 );
2360 }
2361 }
2362 }
2363 }
2364
2365 #[test]
2371 fn test_default_tools_omit_outputschema() {
2372 let (mut server, _dir) = make_server();
2373 let line = r#"{"jsonrpc":"2.0","id":1,"method":"tools/list","params":{}}"#;
2374 let resp = server.handle_jsonrpc_line_unwrap(line);
2375 let tools = resp.result.unwrap().get("tools").cloned().unwrap();
2376 for tool in tools.as_array().unwrap() {
2377 let name = tool.get("name").and_then(|v| v.as_str()).unwrap_or("?");
2378 assert!(
2379 tool.get("outputSchema").is_none(),
2380 "default tool {name} unexpectedly has outputSchema: \
2381 {:?}. Remove it, or make tools/call also emit \
2382 structuredContent matching the schema (MCP 2025-06-18).",
2383 tool.get("outputSchema")
2384 );
2385 }
2386 }
2387
2388 #[test]
2404 fn test_tools_list_has_no_io_impostor_tools() {
2405 let (mut server, _dir) = make_server();
2406 let line = r#"{"jsonrpc":"2.0","id":1,"method":"tools/list","params":{}}"#;
2407 let resp = server.handle_jsonrpc_line_unwrap(line);
2408 let tools = resp.result.unwrap().get("tools").cloned().unwrap();
2409 let names: Vec<String> = tools
2410 .as_array()
2411 .unwrap()
2412 .iter()
2413 .filter_map(|t| t.get("name").and_then(|v| v.as_str()).map(String::from))
2414 .collect();
2415
2416 const FORBIDDEN: &[&str] = &[
2420 "read_file",
2421 "write_file",
2422 "edit_file",
2423 "execute_command",
2424 "list_directory",
2425 "search_files",
2426 "create_directory",
2427 "delete_file",
2428 ];
2429
2430 for forbidden in FORBIDDEN {
2431 assert!(
2432 !names.iter().any(|n| n == forbidden),
2433 "sqz-mcp must not advertise {forbidden} — that name implies \
2434 I/O we cannot perform and shadows the host's real tool. \
2435 See the silent-write bug follow-up to issue #5. \
2436 Tools registered: {names:?}"
2437 );
2438 }
2439 }
2440
2441 #[test]
2446 fn test_default_tools_advertise_compress_tool() {
2447 let (mut server, _dir) = make_server();
2448 let line = r#"{"jsonrpc":"2.0","id":1,"method":"tools/list","params":{}}"#;
2449 let resp = server.handle_jsonrpc_line_unwrap(line);
2450 let tools = resp.result.unwrap().get("tools").cloned().unwrap();
2451 let names: Vec<String> = tools
2452 .as_array()
2453 .unwrap()
2454 .iter()
2455 .filter_map(|t| t.get("name").and_then(|v| v.as_str()).map(String::from))
2456 .collect();
2457
2458 assert!(
2459 names.iter().any(|n| n == "compress"),
2460 "default tools must include `compress`; got {names:?}"
2461 );
2462 }
2463
2464 #[test]
2469 fn test_default_tools_advertise_passthrough_and_expand() {
2470 let (mut server, _dir) = make_server();
2471 let line = r#"{"jsonrpc":"2.0","id":1,"method":"tools/list","params":{}}"#;
2472 let resp = server.handle_jsonrpc_line_unwrap(line);
2473 let tools = resp.result.unwrap().get("tools").cloned().unwrap();
2474 let names: Vec<String> = tools
2475 .as_array()
2476 .unwrap()
2477 .iter()
2478 .filter_map(|t| t.get("name").and_then(|v| v.as_str()).map(String::from))
2479 .collect();
2480 assert!(
2481 names.iter().any(|n| n == "passthrough"),
2482 "passthrough tool must be advertised; got {names:?}"
2483 );
2484 assert!(
2485 names.iter().any(|n| n == "expand"),
2486 "expand tool must be advertised; got {names:?}"
2487 );
2488 }
2489
2490 #[test]
2491 fn test_tool_descriptions_stay_short_and_explain_refs_once() {
2492 let mut total = 0;
2493 for tool in default_tool_definitions() {
2494 let params: Vec<String> = tool.input_schema["properties"]
2495 .as_object()
2496 .map(|props| {
2497 props
2498 .values()
2499 .filter_map(|p| p["description"].as_str().map(str::to_string))
2500 .collect()
2501 })
2502 .unwrap_or_default();
2503 for text in params.iter().chain(std::iter::once(&tool.description)) {
2504 assert!(!text.contains('\u{2014}') && !text.contains('\u{2013}'), "{}: {text}", tool.id);
2505 total += text.len();
2506 }
2507 assert_eq!(
2508 tool.description.contains("§ref:HASH:L40-80§"),
2509 tool.id == "expand",
2510 "{}",
2511 tool.id
2512 );
2513 }
2514 assert!(total < 2000, "tool and parameter descriptions are {total} bytes");
2515 }
2516
2517 #[test]
2518 fn test_passthrough_returns_input_unchanged() {
2519 let (mut server, _dir) = make_server();
2522 let text = "ls -la\ntotal 42\n-rw-r--r-- 1 root root 17 Jan 1 00:00 readme.md\n";
2523 let req = ToolCallRequest {
2524 tool_id: "passthrough".to_string(),
2525 input: serde_json::json!({ "text": text }),
2526 intent: None,
2527 };
2528 let resp = server.handle_tool_call(req).unwrap();
2529 assert_eq!(resp.output, text, "passthrough must return byte-exact input");
2530 assert_eq!(
2531 resp.tokens_original, resp.tokens_compressed,
2532 "passthrough is 1:1 so token counts must match"
2533 );
2534 }
2535
2536 #[test]
2537 fn test_passthrough_falls_back_to_serialising_if_no_text_field() {
2538 let (mut server, _dir) = make_server();
2544 let req = ToolCallRequest {
2545 tool_id: "passthrough".to_string(),
2546 input: serde_json::json!({ "foo": 1, "bar": "baz" }),
2547 intent: None,
2548 };
2549 let resp = server.handle_tool_call(req).unwrap();
2550 assert!(resp.output.contains("foo"));
2551 assert!(resp.output.contains("bar"));
2552 }
2553
2554 #[test]
2555 fn test_expand_tool_returns_not_found_marker_on_miss() {
2556 let (mut server, _dir) = make_server();
2561 let req = ToolCallRequest {
2562 tool_id: "expand".to_string(),
2563 input: serde_json::json!({ "prefix": "deadbeef00000000" }),
2564 intent: None,
2565 };
2566 let resp = server.handle_tool_call(req).unwrap();
2567 assert!(resp.output.contains("hash-not-found"));
2568 assert!(resp.output.contains("deadbeef00000000"));
2569 }
2570
2571 #[test]
2572 fn recall_tool_finds_previously_compressed_content() {
2573 let (mut server, _dir) = make_server();
2574 let text = format!(
2576 "unique-recall-marker deployment rollout failed with status 503\n{}",
2577 "surrounding output line with ordinary words\n".repeat(20)
2578 );
2579 let req = ToolCallRequest {
2580 tool_id: "compress".to_string(),
2581 input: serde_json::json!({ "text": text }),
2582 intent: None,
2583 };
2584 server.handle_tool_call(req).unwrap();
2585
2586 let req = ToolCallRequest {
2587 tool_id: "sqz_recall".to_string(),
2588 input: serde_json::json!({ "query": "rollout 503" }),
2589 intent: None,
2590 };
2591 let resp = server.handle_tool_call(req).unwrap();
2592 assert!(
2593 resp.output.contains("hits=") && resp.output.contains("ref="),
2594 "expected a ranked hit with a ref: {}",
2595 resp.output
2596 );
2597 assert!(resp.output.contains(">>rollout<<"), "snippet markers missing: {}", resp.output);
2598
2599 let ref_prefix = resp
2601 .output
2602 .split("ref=")
2603 .nth(1)
2604 .unwrap()
2605 .chars()
2606 .take_while(|c| c.is_ascii_hexdigit())
2607 .collect::<String>();
2608 let req = ToolCallRequest {
2609 tool_id: "expand".to_string(),
2610 input: serde_json::json!({ "prefix": ref_prefix }),
2611 intent: None,
2612 };
2613 let resp = server.handle_tool_call(req).unwrap();
2614 assert!(
2615 resp.output.contains("unique-recall-marker"),
2616 "expand must recover the indexed original: {}",
2617 resp.output
2618 );
2619 }
2620
2621 #[test]
2622 fn recall_tool_reports_no_matches_cleanly() {
2623 let (mut server, _dir) = make_server();
2624 let req = ToolCallRequest {
2625 tool_id: "sqz_recall".to_string(),
2626 input: serde_json::json!({ "query": "nothing-indexed-yet-zzz" }),
2627 intent: None,
2628 };
2629 let resp = server.handle_tool_call(req).unwrap();
2630 assert!(resp.output.contains("no matches"), "{}", resp.output);
2631 }
2632
2633 #[test]
2634 fn test_default_tools_advertise_recall() {
2635 let (mut server, _dir) = make_server();
2636 let line = r#"{"jsonrpc":"2.0","id":1,"method":"tools/list","params":{}}"#;
2637 let resp = server.handle_jsonrpc_line_unwrap(line);
2638 let tools = resp.result.unwrap().get("tools").cloned().unwrap();
2639 let ids: Vec<String> = tools
2640 .as_array()
2641 .unwrap()
2642 .iter()
2643 .filter_map(|t| t.get("name").and_then(|v| v.as_str()).map(String::from))
2644 .collect();
2645 assert!(ids.iter().any(|n| n == "sqz_recall"), "tools: {ids:?}");
2646 }
2647
2648 #[test]
2649 fn test_expand_tool_strips_ref_token_wrapper() {
2650 let (mut server, _dir) = make_server();
2654 for prefix_input in [
2655 "§ref:deadbeef00000000§",
2656 "ref:deadbeef00000000",
2657 "deadbeef00000000",
2658 " deadbeef00000000 ",
2659 ] {
2660 let req = ToolCallRequest {
2661 tool_id: "expand".to_string(),
2662 input: serde_json::json!({ "prefix": prefix_input }),
2663 intent: None,
2664 };
2665 let resp = server.handle_tool_call(req).unwrap();
2666 assert!(
2667 resp.output.contains("deadbeef00000000"),
2668 "input {prefix_input:?} did not yield expected prefix in output: {}",
2669 resp.output
2670 );
2671 }
2672 }
2673
2674 #[test]
2680 fn test_notification_returns_none() {
2681 let (mut server, _dir) = make_server();
2682 let line = r#"{"jsonrpc":"2.0","method":"notifications/initialized","params":{}}"#;
2683 assert!(server.handle_jsonrpc_line(line).is_none());
2684 }
2685
2686 #[test]
2690 fn test_unknown_notification_returns_none() {
2691 let (mut server, _dir) = make_server();
2692 let line = r#"{"jsonrpc":"2.0","method":"some/unknown/notif"}"#;
2693 assert!(server.handle_jsonrpc_line(line).is_none());
2694 }
2695
2696 #[test]
2698 fn test_request_still_responds() {
2699 let (mut server, _dir) = make_server();
2700 let line = r#"{"jsonrpc":"2.0","id":1,"method":"initialize","params":{}}"#;
2701 assert!(server.handle_jsonrpc_line(line).is_some());
2702 }
2703
2704 #[test]
2714 fn test_sqz_file_tools_are_advertised() {
2715 let (mut server, _dir) = make_server();
2716 let line = r#"{"jsonrpc":"2.0","id":1,"method":"tools/list","params":{}}"#;
2717 let resp = server.handle_jsonrpc_line_unwrap(line);
2718 let tools = resp.result.unwrap().get("tools").cloned().unwrap();
2719 let names: Vec<String> = tools
2720 .as_array()
2721 .unwrap()
2722 .iter()
2723 .filter_map(|t| t.get("name").and_then(|v| v.as_str()).map(String::from))
2724 .collect();
2725
2726 for expected in ["sqz_read_file", "sqz_list_dir", "sqz_grep"] {
2727 assert!(
2728 names.iter().any(|n| n == expected),
2729 "{expected} must be advertised; got {names:?}"
2730 );
2731 }
2732 }
2733
2734 #[test]
2735 fn test_sqz_file_tools_have_sqz_prefix() {
2736 let (mut server, _dir) = make_server();
2741 let line = r#"{"jsonrpc":"2.0","id":1,"method":"tools/list","params":{}}"#;
2742 let resp = server.handle_jsonrpc_line_unwrap(line);
2743 let tools = resp.result.unwrap().get("tools").cloned().unwrap();
2744
2745 for tool in tools.as_array().unwrap() {
2746 let name = tool.get("name").and_then(|v| v.as_str()).unwrap_or("");
2747 if matches!(
2751 name,
2752 "read_file" | "grep" | "list_dir" | "list_directory"
2753 | "search_files" | "write_file" | "delete_file"
2754 | "edit_file" | "execute_command" | "create_directory"
2755 ) {
2756 panic!(
2757 "tool `{name}` shadows a host-native tool. \
2758 I/O tools must be prefixed `sqz_` — see issue #5 \
2759 follow-up."
2760 );
2761 }
2762 }
2763 }
2764
2765 #[test]
2766 fn test_sqz_read_file_reads_real_file() {
2767 let (mut server, dir) = make_server();
2773 let file_path = dir.path().join("read_me.txt");
2774 let original = "hello from sqz_read_file\nline two\nline three\n";
2775 std::fs::write(&file_path, original).expect("write test file");
2776
2777 let req = ToolCallRequest {
2778 tool_id: "sqz_read_file".to_string(),
2779 input: serde_json::json!({ "path": file_path.to_string_lossy() }),
2780 intent: None,
2781 };
2782 let resp = server.handle_tool_call(req).expect("read should succeed");
2783
2784 assert!(resp.output.contains("[sqz_read_file"));
2787 assert!(resp.output.contains(&format!("size={}", original.len())));
2788
2789 assert!(
2793 resp.output.contains("line two") || resp.output.contains("§ref:"),
2794 "content must survive (possibly as a dedup ref); got: {}",
2795 resp.output
2796 );
2797 }
2798
2799 #[test]
2800 fn test_sqz_read_file_reports_missing_file_clearly() {
2801 let (mut server, _dir) = make_server();
2805 let req = ToolCallRequest {
2806 tool_id: "sqz_read_file".to_string(),
2807 input: serde_json::json!({ "path": "/nonexistent/path/xyz.txt" }),
2808 intent: None,
2809 };
2810 let result = server.handle_tool_call(req);
2811 assert!(result.is_err());
2812 let err = result.unwrap_err().to_string();
2813 assert!(err.contains("sqz_read_file"), "error should name the tool: {err}");
2814 assert!(err.contains("xyz.txt"), "error should name the path: {err}");
2815 }
2816
2817 fn numbered_lines(n: usize) -> String {
2818 (1..=n).map(|i| format!("line {i:06} of the big file\n")).collect()
2819 }
2820
2821 fn read(server: &mut McpServer, input: serde_json::Value) -> String {
2822 server
2823 .handle_tool_call(ToolCallRequest {
2824 tool_id: "sqz_read_file".to_string(),
2825 input,
2826 intent: None,
2827 })
2828 .expect("read should succeed")
2829 .output
2830 }
2831
2832 #[test]
2833 fn test_sqz_read_file_respects_max_bytes() {
2834 let (mut server, dir) = make_server();
2837 let file_path = dir.path().join("big.txt");
2838 let content = numbered_lines(100);
2839 std::fs::write(&file_path, &content).expect("write big file");
2840 let line_len = "line 000001 of the big file\n".len();
2841
2842 let out = read(&mut server, serde_json::json!({ "path": file_path.to_string_lossy(), "max_bytes": line_len * 3 + 5 }));
2843 let (header, body) = out.split_once('\n').unwrap();
2844 assert!(
2845 header.contains(&format!(" lines=1-3 of 100 truncated_to={} continue_with_offset=4]", line_len * 3 + 5)),
2846 "{header}"
2847 );
2848 assert_eq!(body, line_range(&content, 1, 3));
2849
2850 let next = read(&mut server, serde_json::json!({ "path": file_path.to_string_lossy(), "offset": 4, "limit": 2 }));
2851 assert!(next.ends_with("line 000004 of the big file\nline 000005 of the big file\n"), "{next}");
2852 }
2853
2854 #[test]
2855 fn test_sqz_read_file_default_cap_and_ranged_read_past_it() {
2856 let (mut server, dir) = make_server();
2857 let file_path = dir.path().join("huge.log");
2858 let content = numbered_lines(20_000);
2859 assert!(content.len() > READ_DEFAULT_MAX_BYTES);
2860 std::fs::write(&file_path, &content).unwrap();
2861
2862 let out = read(&mut server, serde_json::json!({ "path": file_path.to_string_lossy() }));
2863 let (header, body) = out.split_once('\n').unwrap();
2864 assert!(header.contains(&format!("truncated_to={READ_DEFAULT_MAX_BYTES}")), "{header}");
2865 assert!(body.len() <= READ_DEFAULT_MAX_BYTES && body.ends_with('\n'));
2866 let shown = body.lines().count();
2867 assert!(header.contains(&format!(" lines=1-{shown} of 20000 ")), "{header}");
2868 assert!(header.contains(&format!("continue_with_offset={}", shown + 1)), "{header}");
2869
2870 let deep = read(&mut server, serde_json::json!({ "path": file_path.to_string_lossy(), "offset": 19_999 }));
2872 assert!(deep.contains(" lines=19999-20000 of 20000]"), "{deep}");
2873 assert!(!deep.contains("truncated_to"), "{deep}");
2874 assert!(deep.ends_with("line 019999 of the big file\nline 020000 of the big file\n"));
2875
2876 let all = read(&mut server, serde_json::json!({ "path": file_path.to_string_lossy(), "max_bytes": 0 }));
2878 assert!(!all.contains("truncated_to"));
2879 }
2880
2881 #[test]
2882 fn test_sqz_read_file_reports_a_cut_line() {
2883 let (mut server, dir) = make_server();
2886 let file_path = dir.path().join("min.js");
2887 std::fs::write(&file_path, format!("{}\nsecond line\n", "é".repeat(1000))).unwrap();
2888
2889 let out = read(&mut server, serde_json::json!({ "path": file_path.to_string_lossy(), "max_bytes": 101 }));
2890 let (header, body) = out.split_once('\n').unwrap();
2891 assert!(header.contains(" lines=1-1 of 2 truncated_to=101 line_cut=1 continue_with_offset=2]"), "{header}");
2892 assert_eq!(body, "é".repeat(50), "cut on a char boundary");
2893 }
2894
2895 #[test]
2896 fn test_sqz_read_file_refuses_files_over_the_hard_limit() {
2897 let (mut server, dir) = make_server();
2898 let file_path = dir.path().join("sparse.bin");
2899 let f = std::fs::File::create(&file_path).unwrap();
2900 f.set_len(READ_HARD_LIMIT_BYTES + 1).unwrap();
2901 let err = server
2902 .handle_tool_call(ToolCallRequest {
2903 tool_id: "sqz_read_file".to_string(),
2904 input: serde_json::json!({ "path": file_path.to_string_lossy() }),
2905 intent: None,
2906 })
2907 .unwrap_err()
2908 .to_string();
2909 assert!(err.contains("too large") && err.contains("sparse.bin"), "{err}");
2910 }
2911
2912 #[test]
2921 fn test_sqz_read_file_dedup_fires_on_repeat_read() {
2922 let (mut server, dir) = make_server();
2923 let file_path = dir.path().join("dedup_me.txt");
2924 let content = "sqz dedup regression — issue #12 follow-up\n\
2927 this file should collapse to a 13-token ref on \
2928 the second read, proving the cache is wired up.\n"
2929 .repeat(10);
2930 std::fs::write(&file_path, &content).expect("write test file");
2931
2932 let req = || ToolCallRequest {
2933 tool_id: "sqz_read_file".to_string(),
2934 input: serde_json::json!({ "path": file_path.to_string_lossy() }),
2935 intent: None,
2936 };
2937
2938 let first = server.handle_tool_call(req()).expect("first read");
2940 assert!(
2941 !first.output.contains("§ref:"),
2942 "first read must emit full content, not a ref; got: {}",
2943 first.output
2944 );
2945
2946 let second = server.handle_tool_call(req()).expect("second read");
2948 assert!(
2949 second.output.contains("§ref:"),
2950 "second read must emit a §ref:HASH§ token (the whole point \
2951 of the cache). If this fails, handle_sqz_read_file is \
2952 calling engine.compress() instead of compress_with_cache. \
2953 Got: {}",
2954 second.output
2955 );
2956 assert!(
2957 second.tokens_compressed < 30,
2958 "dedup ref should be ~13 tokens, got {}",
2959 second.tokens_compressed
2960 );
2961 }
2962
2963 #[test]
2964 fn test_sqz_read_file_never_stores_credential_mentions() {
2965 let (mut server, dir) = make_server();
2966 let file_path = dir.path().join("app.log");
2967 let mut content: String = (0..40)
2968 .map(|i| format!("2026-10-05T12:00:{i:02}Z INFO GET /healthz 200 1ms\n"))
2969 .collect();
2970 content.push_str("2026-10-05T12:30:00Z DEBUG db config host=db.internal user=app password=hunter2\n");
2971 std::fs::write(&file_path, &content).expect("write test file");
2972 for run in ["first", "second"] {
2973 let resp = server
2974 .handle_tool_call(ToolCallRequest {
2975 tool_id: "sqz_read_file".to_string(),
2976 input: serde_json::json!({ "path": file_path.to_string_lossy() }),
2977 intent: None,
2978 })
2979 .unwrap();
2980 assert!(resp.output.contains(&content), "{run} read: {}", resp.output);
2981 }
2982 let stored = server.engine.cache_manager().check_dedup_with_meta(content.as_bytes()).unwrap();
2983 assert!(stored.is_none());
2984 }
2985
2986 fn python_module(fns: usize) -> String {
2987 (1..=fns)
2988 .map(|i| format!("def handler_{i:03}(request):\n token = request.headers.get('Authorization')\n return verify(token, scope='handler_{i:03}')\n\n"))
2989 .collect()
2990 }
2991
2992 #[test]
2993 fn test_sqz_read_file_ranged_read_returns_lines_with_header() {
2994 let (mut server, dir) = make_server();
2995 let file_path = dir.path().join("auth.py");
2996 let content = python_module(50);
2997 std::fs::write(&file_path, &content).unwrap();
2998
2999 let resp = server
3000 .handle_tool_call(ToolCallRequest {
3001 tool_id: "sqz_read_file".to_string(),
3002 input: serde_json::json!({ "path": file_path.to_string_lossy(), "offset": 5, "limit": 4 }),
3003 intent: None,
3004 })
3005 .expect("ranged read");
3006 let (header, body) = resp.output.split_once('\n').unwrap();
3007 assert!(header.ends_with(" lines=5-8 of 200]"), "{header}");
3008 assert_eq!(body, line_range(&content, 5, 8));
3009 assert_eq!(body.lines().count(), 4);
3010 assert!(body.starts_with("def handler_002"), "{body}");
3011 assert!(body.ends_with("\n\n"), "line 8 is the blank line after handler_002");
3012
3013 let tail = server
3015 .handle_tool_call(ToolCallRequest {
3016 tool_id: "sqz_read_file".to_string(),
3017 input: serde_json::json!({ "path": file_path.to_string_lossy(), "offset": 197 }),
3018 intent: None,
3019 })
3020 .unwrap();
3021 assert!(tail.output.contains(" lines=197-200 of 200]"), "{}", tail.output);
3022 let head = server
3023 .handle_tool_call(ToolCallRequest {
3024 tool_id: "sqz_read_file".to_string(),
3025 input: serde_json::json!({ "path": file_path.to_string_lossy(), "limit": 3 }),
3026 intent: None,
3027 })
3028 .unwrap();
3029 assert!(head.output.contains(" lines=1-3 of 200]"), "{}", head.output);
3030 assert!(head.output.ends_with(" return verify(token, scope='handler_001')\n"), "{}", head.output);
3031
3032 let past_end = server.handle_tool_call(ToolCallRequest {
3033 tool_id: "sqz_read_file".to_string(),
3034 input: serde_json::json!({ "path": file_path.to_string_lossy(), "offset": 500 }),
3035 intent: None,
3036 });
3037 assert!(past_end.is_err());
3038 }
3039
3040 #[test]
3041 fn test_sqz_read_file_slice_after_full_read_returns_line_range_ref() {
3042 let (mut server, dir) = make_server();
3043 let file_path = dir.path().join("auth.py");
3044 let content = python_module(50);
3045 std::fs::write(&file_path, &content).unwrap();
3046 let hash = sqz_engine::CacheManager::sha256_hex(content.as_bytes());
3047
3048 let full = server
3049 .handle_tool_call(ToolCallRequest {
3050 tool_id: "sqz_read_file".to_string(),
3051 input: serde_json::json!({ "path": file_path.to_string_lossy() }),
3052 intent: None,
3053 })
3054 .unwrap();
3055 assert!(!full.output.contains("§ref:"));
3056
3057 let slice = server
3061 .handle_tool_call(ToolCallRequest {
3062 tool_id: "sqz_read_file".to_string(),
3063 input: serde_json::json!({ "path": file_path.to_string_lossy(), "offset": 41, "limit": 40 }),
3064 intent: None,
3065 })
3066 .unwrap();
3067 let expected = format!("§ref:{}:L41-80§", &hash[..16]);
3068 assert!(slice.output.ends_with(&expected), "{}", slice.output);
3069 assert!(slice.output.contains(" lines=41-80 of 200]"));
3070 assert!(slice.tokens_compressed < 25 && slice.tokens_original > 200, "{} / {}", slice.tokens_compressed, slice.tokens_original);
3071
3072 let expanded = server
3074 .handle_tool_call(ToolCallRequest {
3075 tool_id: "expand".to_string(),
3076 input: serde_json::json!({ "prefix": expected }),
3077 intent: None,
3078 })
3079 .unwrap();
3080 let (header, body) = expanded.output.split_once('\n').unwrap();
3081 assert!(header.contains("lines=41-80"), "{header}");
3082 assert_eq!(body, line_range(&content, 41, 80));
3083
3084 let tiny = server
3086 .handle_tool_call(ToolCallRequest {
3087 tool_id: "sqz_read_file".to_string(),
3088 input: serde_json::json!({ "path": file_path.to_string_lossy(), "offset": 1, "limit": 2 }),
3089 intent: None,
3090 })
3091 .unwrap();
3092 assert!(!tiny.output.contains("§ref:"), "{}", tiny.output);
3093 assert!(tiny.output.contains("def handler_001"));
3094 }
3095
3096 #[test]
3097 fn test_sqz_read_file_delta_after_an_edit_numbers_lines_from_one() {
3098 let (mut server, dir) = make_server();
3099 let file_path = dir.path().join("rows.txt");
3100 let rows: Vec<String> = (1..=60).map(|i| format!("row {i:02}\n")).collect();
3101 std::fs::write(&file_path, rows.concat()).unwrap();
3102 let path = file_path.to_string_lossy().to_string();
3103 read(&mut server, serde_json::json!({ "path": path }));
3104 let hash = sqz_engine::CacheManager::sha256_hex(rows.concat().as_bytes());
3105
3106 std::fs::write(&file_path, [&rows[..2], &rows[3..]].concat().concat()).unwrap();
3107 let out = read(&mut server, serde_json::json!({ "path": path }));
3108 let expected = format!(
3109 "§delta:{}§\n @@ skip 1 unchanged lines @@\n row 02\n-[1 lines removed at L3]\n row 04\n @@ skip 56 unchanged lines @@",
3110 &hash[..16]
3111 );
3112 assert!(out.ends_with(&expected), "{out}");
3113 }
3114
3115 #[test]
3116 fn test_sqz_read_file_ranged_reread_is_never_a_delta() {
3117 let (mut server, dir) = make_server();
3118 let file_path = dir.path().join("lines.txt");
3119 let content: String = (1..=500)
3120 .map(|i| format!("line {i:04} the quick brown fox needle={}\n", i % 7 == 0))
3121 .collect();
3122 std::fs::write(&file_path, &content).unwrap();
3123 let path = file_path.to_string_lossy().to_string();
3124 read(&mut server, serde_json::json!({ "path": path, "offset": 1, "limit": 50 }));
3125
3126 let next = read(&mut server, serde_json::json!({ "path": path, "offset": 11, "limit": 50 }));
3127 let (header, body) = next.split_once('\n').unwrap();
3128 assert!(header.ends_with(" lines=11-60 of 500]"), "{header}");
3129 assert_eq!(body, line_range(&content, 11, 60));
3130 }
3131
3132 #[test]
3133 fn test_content_smaller_than_a_ref_is_never_a_ref_or_delta() {
3134 let (mut server, dir) = make_server();
3135 let root = project_dir(&dir);
3136 std::fs::write(root.join("empty.txt"), "").unwrap();
3137 std::fs::write(root.join("ok.txt"), "\x1b[32mok\x1b[0m\n").unwrap();
3138 let compress = |server: &mut McpServer, text: &str| {
3139 server
3140 .handle_tool_call(ToolCallRequest {
3141 tool_id: "compress".to_string(),
3142 input: serde_json::json!({ "text": text }),
3143 intent: None,
3144 })
3145 .unwrap()
3146 .output
3147 };
3148 for run in ["first", "second"] {
3149 assert_eq!(compress(&mut server, ""), "", "{run}");
3150 let empty = read(&mut server, serde_json::json!({ "path": root.join("empty.txt").to_string_lossy() }));
3151 assert!(empty.ends_with("size=0]\n"), "{run}: {empty}");
3152 let ok = read(&mut server, serde_json::json!({ "path": root.join("ok.txt").to_string_lossy() }));
3153 assert!(ok.ends_with("]\nok\n"), "{run}: {ok}");
3154 let none = grep(&mut server, serde_json::json!({ "pattern": "no such text", "path": root.to_string_lossy() }));
3155 assert!(none.ends_with("matches=0]\n"), "{run}: {none}");
3156 }
3157 }
3158
3159 #[test]
3160 fn test_content_smaller_than_a_ref_is_still_stored() {
3161 let (mut server, dir) = make_server();
3162 let root = project_dir(&dir);
3163 std::fs::write(root.join("ok.txt"), "zebra build ok\n").unwrap();
3164 let call = |server: &mut McpServer, tool: &str, input: serde_json::Value| {
3165 server
3166 .handle_tool_call(ToolCallRequest { tool_id: tool.to_string(), input, intent: None })
3167 .unwrap()
3168 .output
3169 };
3170
3171 read(&mut server, serde_json::json!({ "path": root.join("ok.txt").to_string_lossy() }));
3172 call(&mut server, "compress", serde_json::json!({ "text": "zebra crossing seven" }));
3173 let recalled = call(&mut server, "sqz_recall", serde_json::json!({ "query": "zebra" }));
3174 for text in ["zebra build ok\n", "zebra crossing seven"] {
3175 let hash = sqz_engine::CacheManager::sha256_hex(text.as_bytes());
3176 assert!(recalled.contains(&format!("ref={}", &hash[..16])), "{recalled}");
3177 let expanded = call(&mut server, "expand", serde_json::json!({ "prefix": &hash[..16] }));
3178 assert_eq!(expanded, format!("[sqz:expand hash={hash}]\n{text}"));
3179 }
3180 }
3181
3182 #[test]
3183 fn test_read_tools_count_tokens_with_one_tokenizer() {
3184 let (mut server, dir) = make_server();
3188 for i in 0..40 {
3189 std::fs::write(dir.path().join(format!("f{i}.rs")), "fn needle_fn() -> Result<(), Error> { Ok(()) }\n").unwrap();
3190 }
3191 let resp = server
3192 .handle_tool_call(ToolCallRequest {
3193 tool_id: "sqz_grep".to_string(),
3194 input: serde_json::json!({ "pattern": "needle_fn", "path": dir.path().to_string_lossy() }),
3195 intent: None,
3196 })
3197 .unwrap();
3198 assert!(
3199 resp.tokens_compressed <= resp.tokens_original,
3200 "verbatim result must never log as negative savings: {} -> {}",
3201 resp.tokens_original,
3202 resp.tokens_compressed
3203 );
3204 }
3205
3206 #[test]
3207 fn test_sqz_list_dir_lists_directory() {
3208 let (mut server, dir) = make_server();
3211 std::fs::write(dir.path().join("a.txt"), "").unwrap();
3212 std::fs::write(dir.path().join("b.rs"), "").unwrap();
3213 std::fs::write(dir.path().join(".secret"), "").unwrap();
3215 std::fs::create_dir_all(dir.path().join("node_modules/foo")).unwrap();
3217 std::fs::write(dir.path().join("node_modules/foo/pkg.json"), "").unwrap();
3218
3219 let req = ToolCallRequest {
3220 tool_id: "sqz_list_dir".to_string(),
3221 input: serde_json::json!({ "path": dir.path().to_string_lossy() }),
3222 intent: None,
3223 };
3224 let resp = server.handle_tool_call(req).expect("list should succeed");
3225
3226 assert!(resp.output.contains("a.txt"));
3227 assert!(resp.output.contains("b.rs"));
3228 assert!(
3229 !resp.output.contains(".secret"),
3230 "hidden files must be skipped; got: {}",
3231 resp.output
3232 );
3233 assert!(
3234 !resp.output.contains("node_modules"),
3235 "node_modules must be skipped; got: {}",
3236 resp.output
3237 );
3238 }
3239
3240 fn list(server: &mut McpServer, input: serde_json::Value) -> String {
3241 server
3242 .handle_tool_call(ToolCallRequest {
3243 tool_id: "sqz_list_dir".to_string(),
3244 input,
3245 intent: None,
3246 })
3247 .expect("list should succeed")
3248 .output
3249 }
3250
3251 fn project_dir(dir: &TempDir) -> std::path::PathBuf {
3254 let root = dir.path().join("proj");
3255 std::fs::create_dir_all(&root).unwrap();
3256 root
3257 }
3258
3259 #[test]
3260 fn test_sqz_list_dir_depth_one_is_immediate_children() {
3261 let (mut server, dir) = make_server();
3262 let root = project_dir(&dir);
3263 std::fs::create_dir_all(root.join("src/nested")).unwrap();
3264 std::fs::write(root.join("src/lib.rs"), "").unwrap();
3265 std::fs::write(root.join("src/nested/deep.rs"), "").unwrap();
3266 std::fs::write(root.join("README.md"), "").unwrap();
3267 let path = root.to_string_lossy().to_string();
3268
3269 let one = list(&mut server, serde_json::json!({ "path": path }));
3270 assert!(one.contains("entries=2]"), "{one}");
3271 assert!(one.contains("d src") && one.contains("f README.md"), "{one}");
3272 assert!(!one.contains("lib.rs"), "default depth lists immediate children only: {one}");
3273
3274 let two = list(&mut server, serde_json::json!({ "path": path, "max_depth": 2 }));
3275 assert!(two.contains("entries=4]"), "{two}");
3276 let nested = format!("d {}", Path::new("src").join("nested").display());
3277 assert!(two.contains("lib.rs") && two.contains(&nested), "{two}");
3278 assert!(!two.contains("deep.rs"), "{two}");
3279
3280 let zero = list(&mut server, serde_json::json!({ "path": path, "max_depth": 0 }));
3281 assert!(zero.contains("entries=2]"), "{zero}");
3282 }
3283
3284 #[test]
3285 fn test_sqz_list_dir_max_entries() {
3286 let (mut server, dir) = make_server();
3287 let root = project_dir(&dir);
3288 for i in 0..12 {
3289 std::fs::write(root.join(format!("f{i:02}.txt")), "").unwrap();
3290 }
3291 let path = root.to_string_lossy().to_string();
3292
3293 let capped = list(&mut server, serde_json::json!({ "path": path, "max_entries": 5 }));
3294 assert!(capped.contains("entries=5 stopped_at_max_entries=5]"), "{capped}");
3295 assert!(capped.contains("f04.txt") && !capped.contains("f05.txt"), "{capped}");
3296
3297 let exact = list(&mut server, serde_json::json!({ "path": path, "max_entries": 12 }));
3299 assert!(exact.contains("entries=12]"), "{exact}");
3300
3301 std::fs::create_dir(root.join("vendor")).unwrap();
3303 let skipped = list(&mut server, serde_json::json!({ "path": path, "max_entries": 12 }));
3304 assert!(skipped.contains("entries=12]"), "{skipped}");
3305 }
3306
3307 #[test]
3308 fn test_sqz_grep_finds_literal_matches() {
3309 let (mut server, dir) = make_server();
3312 std::fs::write(
3313 dir.path().join("code.rs"),
3314 "fn main() {\n // TODO: refactor this\n println!(\"hello\");\n}\n",
3315 )
3316 .unwrap();
3317 std::fs::write(
3318 dir.path().join("other.rs"),
3319 "fn nothing_to_do() {\n // just a comment\n}\n",
3320 )
3321 .unwrap();
3322
3323 let req = ToolCallRequest {
3324 tool_id: "sqz_grep".to_string(),
3325 input: serde_json::json!({
3326 "pattern": "TODO",
3327 "path": dir.path().to_string_lossy(),
3328 }),
3329 intent: None,
3330 };
3331 let resp = server.handle_tool_call(req).expect("grep should succeed");
3332
3333 assert!(resp.output.contains("matches=1"), "header must include match count; got: {}", resp.output);
3336 assert!(
3337 resp.output.contains("TODO"),
3338 "match content must be present; got: {}",
3339 resp.output
3340 );
3341 assert!(
3343 resp.output.contains("code.rs:2:"),
3344 "output must be grep-style path:lineno:text; got: {}",
3345 resp.output
3346 );
3347 }
3348
3349 #[test]
3350 fn test_sqz_grep_regex_mode() {
3351 let (mut server, dir) = make_server();
3352 std::fs::write(
3353 dir.path().join("t.rs"),
3354 "fn my_test() {}\nfn your_test() {}\nfn not_matching() {}\n",
3355 )
3356 .unwrap();
3357
3358 let req = ToolCallRequest {
3359 tool_id: "sqz_grep".to_string(),
3360 input: serde_json::json!({
3361 "pattern": r"fn \w+_test",
3362 "path": dir.path().to_string_lossy(),
3363 "regex": true,
3364 }),
3365 intent: None,
3366 };
3367 let resp = server.handle_tool_call(req).expect("grep should succeed");
3368 assert!(resp.output.contains("matches=2"));
3369 }
3370
3371 #[test]
3372 fn test_sqz_grep_invalid_regex_errors_cleanly() {
3373 let (mut server, dir) = make_server();
3377 let req = ToolCallRequest {
3378 tool_id: "sqz_grep".to_string(),
3379 input: serde_json::json!({
3380 "pattern": "fn (unclosed",
3381 "path": dir.path().to_string_lossy(),
3382 "regex": true,
3383 }),
3384 intent: None,
3385 };
3386 let result = server.handle_tool_call(req);
3387 assert!(result.is_err());
3388 let err = result.unwrap_err().to_string();
3389 assert!(err.contains("sqz_grep"));
3390 assert!(err.contains("invalid regex"));
3391 }
3392
3393 #[test]
3394 fn test_sqz_grep_caps_at_max_matches() {
3395 let (mut server, dir) = make_server();
3398 let content: String = (0..100).map(|i| format!("line {i}: match\n")).collect();
3399 std::fs::write(dir.path().join("many.txt"), content).unwrap();
3400
3401 let req = ToolCallRequest {
3402 tool_id: "sqz_grep".to_string(),
3403 input: serde_json::json!({
3404 "pattern": "match",
3405 "path": dir.path().to_string_lossy(),
3406 "max_matches": 10,
3407 }),
3408 intent: None,
3409 };
3410 let resp = server.handle_tool_call(req).expect("grep should succeed");
3411 assert!(resp.output.contains("matches=10 max_matches_reached=10]"),
3412 "should stop at max_matches; got: {}", resp.output);
3413 }
3414
3415 fn grep(server: &mut McpServer, input: serde_json::Value) -> String {
3416 server
3417 .handle_tool_call(ToolCallRequest {
3418 tool_id: "sqz_grep".to_string(),
3419 input,
3420 intent: None,
3421 })
3422 .expect("grep should succeed")
3423 .output
3424 }
3425
3426 #[test]
3427 fn test_sqz_grep_clips_long_lines_around_the_hit() {
3428 let (mut server, dir) = make_server();
3430 let root = project_dir(&dir);
3431 let line = format!("{}needle_here{}", "é".repeat(3000), "ü".repeat(3000));
3432 std::fs::write(root.join("bundle.min.js"), format!("{line}\nshort needle_here line\n")).unwrap();
3433 let input = serde_json::json!({ "pattern": "needle_here", "path": root.to_string_lossy() });
3434
3435 let out = grep(&mut server, input.clone());
3436 let (header, body) = out.split_once('\n').unwrap();
3437 assert!(header.contains("matches=2 long_lines_clipped=1 (max_line_chars=400)]"), "{header}");
3438 let first = body.lines().next().unwrap();
3439 let shown = first.split_once("bundle.min.js:1:").map(|(_, s)| s).unwrap();
3440 assert_eq!(
3441 shown,
3442 format!("…{}needle_here{}… [line clipped: 6011 chars]", "é".repeat(194), "ü".repeat(195))
3443 );
3444 assert!(body.contains("bundle.min.js:2:short needle_here line"), "{body}");
3445
3446 let (mut fresh, _store) = make_server();
3449 let mut uncapped = input;
3450 uncapped["max_line_chars"] = 0.into();
3451 uncapped["max_bytes"] = 0.into();
3452 let whole = grep(&mut fresh, uncapped);
3453 assert!(!whole.contains("clipped"), "{}", whole.chars().take(200).collect::<String>());
3454 assert!(whole.contains(&line));
3455 }
3456
3457 #[test]
3458 fn test_clip_line_windows() {
3459 let line: String = ('a'..='z').cycle().take(1000).collect();
3460 let start = clip_line(&line, (0, 3), 100).unwrap();
3462 assert!(start.starts_with("abc") && start.contains("… [line clipped: 1000 chars]"), "{start}");
3463 let end = clip_line(&line, (997, 1000), 100).unwrap();
3465 assert!(end.starts_with('…') && end.ends_with(&format!("{} [line clipped: 1000 chars]", &line[900..])), "{end}");
3466 let long = clip_line(&line, (500, 800), 100).unwrap();
3468 assert!(long.starts_with(&format!("…{}", &line[500..600])), "{long}");
3469 assert!(clip_line("short", (0, 1), 100).is_none());
3471 }
3472
3473 #[test]
3474 fn test_sqz_grep_caps_total_bytes() {
3475 let (mut server, dir) = make_server();
3476 let root = project_dir(&dir);
3477 let content: String = (0..2000).map(|i| format!("{i:05} match {}\n", "x".repeat(200))).collect();
3478 std::fs::write(root.join("wide.txt"), content).unwrap();
3479
3480 let out = grep(&mut server, serde_json::json!({
3481 "pattern": "match",
3482 "path": root.to_string_lossy(),
3483 "max_matches": 10_000
3484 }));
3485 let (header, body) = out.split_once('\n').unwrap();
3486 assert!(header.contains(&format!("stopped_at_max_bytes={GREP_DEFAULT_MAX_BYTES}")), "{header}");
3487 assert!(!header.contains("max_matches_reached"), "{header}");
3488 assert!(body.len() <= GREP_DEFAULT_MAX_BYTES, "{}", body.len());
3489 assert!(body.lines().count() > 100, "{}", body.lines().count());
3490 }
3491}