diff --git a/Cargo.lock b/Cargo.lock index 5a7777a8ec..5b7c07d054 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -12241,6 +12241,7 @@ dependencies = [ "async-trait", "chrono", "futures", + "regex", "serde", "serde_json", "tokio", diff --git a/Cargo.toml b/Cargo.toml index 3d1edd8246..36611a3bfd 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -96,6 +96,7 @@ libc = "0.2" gethostname = "1" tokio-tungstenite = { version = "0.24", features = ["rustls-tls-webpki-roots"] } futures = "0.3" +regex = "1" async-trait = "0.1" # serialization diff --git a/crates/mcp/Cargo.toml b/crates/mcp/Cargo.toml index 0442ceaf71..d9cb09aaa3 100644 --- a/crates/mcp/Cargo.toml +++ b/crates/mcp/Cargo.toml @@ -18,3 +18,4 @@ uuid.workspace = true chrono.workspace = true async-trait.workspace = true futures.workspace = true +regex.workspace = true diff --git a/crates/mcp/src/lib.rs b/crates/mcp/src/lib.rs index 0cf61f0d81..7a22840ed7 100644 --- a/crates/mcp/src/lib.rs +++ b/crates/mcp/src/lib.rs @@ -17,6 +17,7 @@ //! originating chat, and a chat can never message itself. mod jsonrpc; +mod search; mod tools; mod transcript; mod zeron; diff --git a/crates/mcp/src/search.rs b/crates/mcp/src/search.rs new file mode 100644 index 0000000000..c31a050826 --- /dev/null +++ b/crates/mcp/src/search.rs @@ -0,0 +1,531 @@ +//! Content search across rendered transcripts. +//! +//! Two modes share one shape. `words` tokenizes the query and scores every +//! message per term: an exact word beats a substring, which beats a +//! typo-tolerant match, and rare terms count for more than ones that appear in +//! most chats. `regex` is the ripgrep-style escape hatch. Both keep a +//! snippet around the strongest hit so an agent can decide whether to +//! `read_chat` the whole thing. + +use regex::{Regex, RegexBuilder}; + +use crate::transcript::RenderedMessage; + +const SNIPPET_BEFORE: usize = 100; +const SNIPPET_AFTER: usize = 160; +/// Words longer than this are never fuzzy-compared (a pasted hash or blob). +const MAX_FUZZY_WORD: usize = 40; + +/// A word in a text: byte range plus its lowercase form. +struct Word { + start: usize, + end: usize, + lower: String, +} + +fn words(text: &str) -> Vec { + let mut out = Vec::new(); + let mut start = None; + for (i, c) in text.char_indices() { + if c.is_alphanumeric() { + start.get_or_insert(i); + } else if let Some(s) = start.take() { + out.push(Word { + start: s, + end: i, + lower: text[s..i].to_lowercase(), + }); + } + } + if let Some(s) = start { + out.push(Word { + start: s, + end: text.len(), + lower: text[s..].to_lowercase(), + }); + } + out +} + +/// Distinct lowercase query terms, in order. +pub fn query_terms(query: &str) -> Vec { + let mut terms: Vec = Vec::new(); + for word in words(query) { + if !terms.contains(&word.lower) { + terms.push(word.lower); + } + } + terms +} + +/// How many single-character edits a term of this length may absorb. +fn edit_budget(len: usize) -> usize { + match len { + 0..=3 => 0, + 4..=7 => 1, + _ => 2, + } +} + +/// Levenshtein distance, giving up once it exceeds `max`. +fn edit_distance(a: &[char], b: &[char], max: usize) -> Option { + if a.len().abs_diff(b.len()) > max { + return None; + } + let mut prev: Vec = (0..=b.len()).collect(); + for (i, ca) in a.iter().enumerate() { + let mut cur = vec![i + 1]; + for (j, cb) in b.iter().enumerate() { + let cost = usize::from(ca != cb); + cur.push((prev[j] + cost).min(prev[j + 1] + 1).min(cur[j] + 1)); + } + if cur.iter().all(|&d| d > max) { + return None; + } + prev = cur; + } + (prev[b.len()] <= max).then_some(prev[b.len()]) +} + +/// Match strength of one word against one term: 0 for no match. +fn term_weight(term: &str, term_chars: &[char], word: &str) -> f32 { + if word == term { + return 1.0; + } + if term_chars.len() >= 3 && word.contains(term) { + return 0.8; + } + let budget = edit_budget(term_chars.len()); + if budget == 0 || word.chars().count() > MAX_FUZZY_WORD { + return 0.0; + } + let word_chars: Vec = word.chars().collect(); + match edit_distance(term_chars, &word_chars, budget) { + Some(1) => 0.6, + Some(_) => 0.5, + None => 0.0, + } +} + +/// What a message (or title) contributes for each query term. +struct TermHits { + weights: Vec, + /// Byte range of the single strongest matching word. + anchor: Option<(usize, usize)>, +} + +fn score_text(terms: &[(String, Vec)], text: &str) -> TermHits { + let mut weights = vec![0.0_f32; terms.len()]; + let mut anchor: Option<(usize, usize, f32)> = None; + for word in words(text) { + for (i, (term, chars)) in terms.iter().enumerate() { + let w = term_weight(term, chars, &word.lower); + if w > weights[i] { + weights[i] = w; + } + if w > 0.0 && anchor.is_none_or(|(_, _, best)| w > best) { + anchor = Some((word.start, word.end, w)); + } + } + } + TermHits { + weights, + anchor: anchor.map(|(s, e, _)| (s, e)), + } +} + +/// The text of a rendered message that search looks at. +fn searchable(message: &RenderedMessage) -> String { + let mut blob = message.text.clone(); + for extra in message + .reasoning + .iter() + .chain(&message.tools) + .chain(&message.errors) + { + if !blob.is_empty() { + blob.push('\n'); + } + blob.push_str(extra); + } + blob +} + +/// A window of `text` around the byte range `[start, end)`, on one line. +fn snippet(text: &str, start: usize, end: usize) -> String { + let head: Vec<(usize, char)> = text[..start].char_indices().collect(); + let from = head + .len() + .checked_sub(SNIPPET_BEFORE) + .map_or(0, |i| head[i].0); + let to = text[end..] + .char_indices() + .nth(SNIPPET_AFTER) + .map_or(text.len(), |(i, _)| end + i); + let flat = text[from..to] + .split_whitespace() + .collect::>() + .join(" "); + format!( + "{}{flat}{}", + if from > 0 { "…" } else { "" }, + if to < text.len() { "…" } else { "" } + ) +} + +/// One message worth showing. +pub struct MessageHit { + /// Index into the chat's rendered messages, oldest first. + pub index: usize, + pub score: f32, + /// Query terms this message matched (empty in regex mode). + pub matched: Vec, + pub snippet: String, +} + +/// Everything one chat's transcript contributes, before cross-chat weighting. +pub struct ChatScan { + /// Best weight per term across the title and every message. + pub coverage: Vec, + pub title_match: bool, + /// Messages with at least one hit, unsorted. + hits: Vec, + /// Regex mode: total matches, capped per message. + pub regex_matches: usize, +} + +struct RawHit { + index: usize, + weights: Vec, + /// Regex mode: matches in this message. + count: usize, + blob: String, + anchor: (usize, usize), +} + +pub enum Matcher { + Words(Vec<(String, Vec)>), + Regex(Regex), +} + +impl Matcher { + pub fn words(query: &str) -> Result { + let terms = query_terms(query); + if terms.is_empty() { + return Err("query has no searchable words".into()); + } + Ok(Self::Words( + terms + .into_iter() + .map(|t| { + let chars = t.chars().collect(); + (t, chars) + }) + .collect(), + )) + } + + pub fn regex(pattern: &str, case_sensitive: bool) -> Result { + RegexBuilder::new(pattern) + .case_insensitive(!case_sensitive) + .size_limit(1 << 20) + .build() + .map(Self::Regex) + .map_err(|e| format!("invalid regex: {e}")) + } + + pub fn term_count(&self) -> usize { + match self { + Self::Words(terms) => terms.len(), + Self::Regex(_) => 1, + } + } + + /// Scan one chat. `messages` pairs each message to search with its index + /// in the full transcript, so callers can filter without losing places. + pub fn scan<'a>( + &self, + title: Option<&str>, + messages: impl IntoIterator, + ) -> ChatScan { + let mut scan = ChatScan { + coverage: vec![0.0; self.term_count()], + title_match: false, + hits: Vec::new(), + regex_matches: 0, + }; + match self { + Self::Words(terms) => { + if let Some(title) = title { + let hits = score_text(terms, title); + scan.title_match = hits.weights.iter().any(|&w| w > 0.0); + merge_max(&mut scan.coverage, &hits.weights); + } + for (index, message) in messages { + let blob = searchable(message); + let hits = score_text(terms, &blob); + let Some(anchor) = hits.anchor else { continue }; + merge_max(&mut scan.coverage, &hits.weights); + scan.hits.push(RawHit { + index, + weights: hits.weights, + count: 0, + blob, + anchor, + }); + } + } + Self::Regex(re) => { + if let Some(title) = title + && re.is_match(title) + { + scan.title_match = true; + scan.coverage[0] = 1.0; + scan.regex_matches += 1; + } + for (index, message) in messages { + let blob = searchable(message); + let mut found = re.find_iter(&blob); + let Some(first) = found.next() else { continue }; + let count = 1 + found.take(4).count(); + scan.coverage[0] = 1.0; + scan.regex_matches += count; + scan.hits.push(RawHit { + index, + weights: vec![1.0], + count, + anchor: (first.start(), first.end()), + blob, + }); + } + } + } + scan + } +} + +fn merge_max(into: &mut [f32], from: &[f32]) { + for (a, b) in into.iter_mut().zip(from) { + *a = a.max(*b); + } +} + +/// A chat that made the cut. +pub struct ChatResult { + /// Index into the scans this came from. + pub scan: usize, + pub score: f32, + pub title_match: bool, + pub matched: Vec, + pub hits: Vec, +} + +/// Rank chats across a set of scans. `min_match` is the share of the query +/// (weighted by term rarity) a chat must cover to be listed. +pub fn rank( + matcher: &Matcher, + scans: &[ChatScan], + min_match: f32, + limit: usize, + snippets: usize, +) -> (usize, Vec) { + // Rare terms matter more: a word in every chat says little. + let idf: Vec = match matcher { + Matcher::Words(terms) => (0..terms.len()) + .map(|t| { + let df = scans.iter().filter(|s| s.coverage[t] > 0.0).count().max(1); + (1.0 + scans.len() as f32 / df as f32).ln() + }) + .collect(), + Matcher::Regex(_) => vec![1.0], + }; + let idf_sum: f32 = idf.iter().sum(); + let weighted = |weights: &[f32]| -> f32 { + weights.iter().zip(&idf).map(|(w, i)| w * i).sum::() / idf_sum + }; + + let mut results: Vec = Vec::new(); + for (i, scan) in scans.iter().enumerate() { + let coverage = weighted(&scan.coverage); + if coverage < min_match || coverage <= 0.0 { + continue; + } + let mut hits: Vec = scan + .hits + .iter() + .map(|raw| { + let score = match matcher { + Matcher::Words(_) => weighted(&raw.weights), + Matcher::Regex(_) => raw.count as f32, + }; + let matched = match matcher { + Matcher::Words(terms) => terms + .iter() + .zip(&raw.weights) + .filter(|(_, w)| **w > 0.0) + .map(|((t, _), _)| t.clone()) + .collect(), + Matcher::Regex(_) => Vec::new(), + }; + MessageHit { + index: raw.index, + score, + matched, + snippet: snippet(&raw.blob, raw.anchor.0, raw.anchor.1), + } + }) + .collect(); + // Best first; among equals the newer message. + hits.sort_by(|a, b| b.score.total_cmp(&a.score).then(b.index.cmp(&a.index))); + let best = hits.first().map_or(0.0, |h| h.score); + let score = match matcher { + Matcher::Words(_) => 0.5 * coverage + 0.5 * best + f32::from(scan.title_match) * 0.25, + Matcher::Regex(_) => scan.regex_matches as f32, + }; + hits.truncate(snippets); + let matched = match matcher { + Matcher::Words(terms) => terms + .iter() + .zip(&scan.coverage) + .filter(|(_, w)| **w > 0.0) + .map(|((t, _), _)| t.clone()) + .collect(), + Matcher::Regex(_) => Vec::new(), + }; + results.push(ChatResult { + scan: i, + score, + title_match: scan.title_match, + matched, + hits, + }); + } + let total = results.len(); + results.sort_by(|a, b| b.score.total_cmp(&a.score)); + results.truncate(limit); + (total, results) +} + +#[cfg(test)] +mod tests { + use super::*; + use zeron_doc::MessageRole; + + fn msg(text: &str) -> RenderedMessage { + RenderedMessage { + id: "m".into(), + role: MessageRole::Assistant, + created_at: 0, + created_at_iso: String::new(), + device_id: "d".into(), + status: None, + text: text.into(), + reasoning: None, + tools: Vec::new(), + pending_input: None, + errors: Vec::new(), + } + } + + fn scan(m: &Matcher, title: &str, texts: &[&str]) -> ChatScan { + let messages: Vec<_> = texts.iter().map(|t| msg(t)).collect(); + m.scan(Some(title), messages.iter().enumerate()) + } + + #[test] + fn terms_split_on_punctuation_and_dedupe() { + assert_eq!( + query_terms("read_chat, Read-Chat!"), + vec!["read".to_string(), "chat".to_string()] + ); + } + + #[test] + fn weights_rank_exact_over_substring_over_typo() { + let chars = |s: &str| s.chars().collect::>(); + let w = |term: &str, word: &str| term_weight(term, &chars(term), word); + assert_eq!(w("proot", "proot"), 1.0); + assert_eq!(w("proot", "prootfs"), 0.8); + assert_eq!(w("proot", "prot"), 0.6); + assert_eq!(w("environment", "enviromet"), 0.5); + assert_eq!(w("cat", "car"), 0.0, "short terms are exact-only"); + assert_eq!(w("rootfs", "kernel"), 0.0); + } + + #[test] + fn typos_and_partial_queries_find_the_chat() { + let m = Matcher::words("androd proot rootfs").unwrap(); + let scans = vec![ + scan( + &m, + "Android runtime", + &["mounting the rootfs under proot", "done"], + ), + scan(&m, "Wallpapers", &["shuffle wallpapers with preloading"]), + ]; + let (total, ranked) = rank(&m, &scans, 0.5, 10, 3); + assert_eq!(total, 1); + assert_eq!(ranked[0].scan, 0); + assert!(ranked[0].matched.contains(&"androd".to_string())); + assert_eq!(ranked[0].hits[0].index, 0); + assert!(ranked[0].hits[0].snippet.contains("rootfs")); + } + + #[test] + fn min_match_filters_and_terms_can_span_messages() { + let m = Matcher::words("reaper zebra").unwrap(); + let scans = vec![scan(&m, "t", &["the reaper kills sessions", "unrelated"])]; + // One of two terms: 0.5 coverage passes the default cut, not a strict one. + assert_eq!(rank(&m, &scans, 0.5, 10, 3).0, 1); + assert_eq!(rank(&m, &scans, 0.9, 10, 3).0, 0); + } + + #[test] + fn rare_terms_outweigh_common_ones() { + let m = Matcher::words("session unicorn").unwrap(); + let scans = vec![ + scan(&m, "a", &["session session"]), + scan(&m, "b", &["session"]), + scan(&m, "c", &["session unicorn"]), + ]; + let (_, ranked) = rank(&m, &scans, 0.0, 10, 1); + assert_eq!(ranked[0].scan, 2); + } + + #[test] + fn title_only_matches_count() { + let m = Matcher::words("wallpaper").unwrap(); + let scans = vec![scan(&m, "Wallpaper shuffle", &["nothing here"])]; + let (_, ranked) = rank(&m, &scans, 0.5, 10, 3); + assert!(ranked[0].title_match); + assert!(ranked[0].hits.is_empty()); + } + + #[test] + fn regex_mode_counts_matches_and_honors_case() { + let m = Matcher::regex(r"HELD_\w+", true).unwrap(); + let scans = vec![ + scan( + &m, + "a", + &["const HELD_DONE_SETTLE = 5s", "and HELD_X again"], + ), + scan(&m, "b", &["held_done_settle lowercase"]), + ]; + let (total, ranked) = rank(&m, &scans, 0.0, 10, 3); + assert_eq!(total, 1); + assert_eq!(ranked[0].scan, 0); + assert_eq!(ranked[0].hits.len(), 2); + assert!(Matcher::regex("(", false).is_err()); + } + + #[test] + fn snippets_stay_on_char_boundaries_and_one_line() { + let text = format!("{}needle\n\n{}", "é".repeat(300), "ü".repeat(300)); + let start = text.find("needle").unwrap(); + let s = snippet(&text, start, start + 6); + assert!(s.starts_with('…') && s.ends_with('…')); + assert!(s.contains("needle") && !s.contains('\n')); + assert_eq!(snippet("short needle", 6, 12), "short needle"); + } +} diff --git a/crates/mcp/src/tools.rs b/crates/mcp/src/tools.rs index 4cfbd8bf9a..d7905d4940 100644 --- a/crates/mcp/src/tools.rs +++ b/crates/mcp/src/tools.rs @@ -9,19 +9,25 @@ use std::sync::Arc; use std::time::{Duration, SystemTime, UNIX_EPOCH}; +use futures::StreamExt; use serde::{Deserialize, Serialize}; use serde_json::{Value, json}; -use zeron_doc::SessionCommandPayload; +use zeron_doc::{MessageRole, SessionCommandPayload}; use zeron_proto::{ Chat, ChatConfig, HarnessId, ReasoningLevel, RunRequest, SandboxLevel, Session, SessionStatus, Space, UserInputAnswer, }; +use crate::search::{Matcher, rank}; use crate::transcript::{RenderOptions, RenderedMessage, render_entries}; use crate::zeron::{HarnessInfo, TurnOutcome, Zeron, session_for, short}; /// Default and ceiling for the blocking waits. const MAX_BATCH: usize = 32; +/// Transcripts fetched at once by `search_chats`; each is a short-lived watch. +const SEARCH_CONCURRENCY: usize = 8; +const SEARCH_DEFAULT_SCAN: usize = 200; +const SEARCH_MAX_SCAN: usize = 1000; const DEFAULT_WAIT: Duration = Duration::from_secs(600); const MAX_WAIT: Duration = Duration::from_secs(3600); /// A session row older than this is not trusted to still be working @@ -139,6 +145,31 @@ fn catalog() -> Vec { "include_tools": { "type": "boolean", "default": true } })), }, + ToolDef { + name: "search_chats", + description: "Search the content of every chat's transcript (messages, tool lines, titles), not just titles. mode 'words' (default) is forgiving: it ranks chats by how many of the query's words they contain, tolerates typos and word fragments, and weighs rare words above common ones, so a loose description like 'proot rootfs android mount' finds the chat. mode 'regex' is ripgrep-style (case-insensitive unless case_sensitive). Returns the best chats with snippets around the hits; pass a hit's offsetFromNewest to read_chat (offset, minus a few for context) to read the surrounding conversation. Scans the newest max_chats chats matching the filters; the reply says if it stopped short. Your own chat is skipped unless include_self.", + input_schema: json!({ + "type": "object", + "properties": { + "query": { "type": "string", "description": "Words to look for, or a regex in regex mode." }, + "mode": { "type": "string", "enum": ["words", "regex"], "default": "words" }, + "project": { "type": "string", "description": "Only chats in this project (id, path, or name)." }, + "device": { "type": "string", "description": "Only chats hosted on this device (id or name)." }, + "parent": { "type": "string", "description": "Only chats created by this chat (id, prefix, or title)." }, + "include_archived": { "type": "boolean", "default": false }, + "include_self": { "type": "boolean", "default": false, "description": "Also search the chat you are speaking from." }, + "role": { "type": "string", "enum": ["user", "assistant"], "description": "Only search messages from this side." }, + "include_tools": { "type": "boolean", "default": true, "description": "Search the one-line tool ledger (file paths, commands)." }, + "include_reasoning": { "type": "boolean", "default": false }, + "case_sensitive": { "type": "boolean", "default": false, "description": "Regex mode only." }, + "min_match": { "type": "number", "minimum": 0, "maximum": 1, "default": 0.5, "description": "Words mode: the share of the query (rare words weigh more) a chat must cover. Lower it for broader recall." }, + "limit": { "type": "integer", "minimum": 1, "maximum": 100, "default": 10, "description": "Chats to return." }, + "snippets": { "type": "integer", "minimum": 1, "maximum": 20, "default": 3, "description": "Matching messages to show per chat." }, + "max_chats": { "type": "integer", "minimum": 1, "maximum": SEARCH_MAX_SCAN, "default": SEARCH_DEFAULT_SCAN, "description": "How many chats (newest activity first) to scan." } + }, + "required": ["query"] + }), + }, ToolDef { name: "send_message", description: "Send a message to a chat. Messages are attributed to your chat. mode 'auto' starts a turn when the chat is idle, steers a running turn through its live mailbox (at the next supported input boundary without interrupting the agent). Use mode 'queue' only to explicitly hold a message for later. With wait=true, blocks until the turn finishes and returns the assistant's reply. For parallel work use send_messages, or send to all chats with wait=false before waiting.", @@ -273,6 +304,36 @@ struct ReadChatArgs { include_tools: Option, } +#[derive(Deserialize)] +struct SearchArgs { + query: String, + mode: Option, + project: Option, + device: Option, + parent: Option, + #[serde(default)] + include_archived: bool, + #[serde(default)] + include_self: bool, + role: Option, + include_tools: Option, + #[serde(default)] + include_reasoning: bool, + #[serde(default)] + case_sensitive: bool, + min_match: Option, + limit: Option, + snippets: Option, + max_chats: Option, +} + +/// What the tool keeps of a chat between scanning and reporting. +struct ScannedChat { + chat: Chat, + /// (id, role, ISO time) of every message, indexed like the transcript. + messages: Vec<(String, &'static str, String)>, +} + #[derive(Deserialize)] struct SendArgs { chat: String, @@ -381,6 +442,14 @@ fn summarize_chat(chat: &Chat, spaces: &[Space], sessions: &[Session]) -> Value }) } +fn role_name(role: MessageRole) -> &'static str { + match role { + MessageRole::User => "user", + MessageRole::Assistant => "assistant", + MessageRole::System => "system", + } +} + fn last_pending_input(messages: &[RenderedMessage]) -> Option { messages.iter().rev().find_map(|m| m.pending_input.clone()) } @@ -415,6 +484,7 @@ impl Tools { "create_chats" => self.batch(parse(args)?, true).await, "send_messages" => self.batch(parse(args)?, false).await, "read_chat" => self.read_chat(parse(args)?).await, + "search_chats" => self.search_chats(parse(args)?).await, "send_message" => self.send_message(parse(args)?).await, "wait_for_turn" => self.wait_for_turn(parse(args)?).await, "interrupt_chat" => self.interrupt_chat(parse(args)?).await, @@ -544,27 +614,47 @@ impl Tools { })) } - async fn list_chats(&self, args: ListChatsArgs) -> anyhow::Result { + /// Chats passing the project/device/archived/parent filters, with the + /// spaces and sessions their summaries need. + async fn filtered_chats( + &self, + project: Option<&str>, + device: Option<&str>, + include_archived: bool, + parent: Option<&str>, + ) -> anyhow::Result<(Vec, Vec, Vec)> { let (mut chats, spaces, sessions) = tokio::try_join!( self.zeron.chats(), self.zeron.spaces(), self.zeron.sessions() )?; - if let Some(project) = args.project.as_deref() { + if let Some(project) = project { let space = self.zeron.resolve_space(project).await?; chats.retain(|c| c.space_id.as_deref() == Some(space.id.as_str())); } - if args.device.is_some() { - let device = self.zeron.resolve_device_id(args.device.as_deref()).await?; + if device.is_some() { + let device = self.zeron.resolve_device_id(device).await?; chats.retain(|c| c.device_id == device); } - if !args.include_archived { + if !include_archived { chats.retain(|c| !c.archived); } - if let Some(parent) = args.parent.as_deref() { + if let Some(parent) = parent { let parent = self.zeron.resolve_chat(parent).await?; chats.retain(|c| c.parent_chat_id.as_deref() == Some(parent.id.as_str())); } + Ok((chats, spaces, sessions)) + } + + async fn list_chats(&self, args: ListChatsArgs) -> anyhow::Result { + let (mut chats, spaces, sessions) = self + .filtered_chats( + args.project.as_deref(), + args.device.as_deref(), + args.include_archived, + args.parent.as_deref(), + ) + .await?; chats.sort_by(|a, b| { let a_at = a.last_message_at.unwrap_or(a.created_at); let b_at = b.last_message_at.unwrap_or(b.created_at); @@ -794,6 +884,122 @@ impl Tools { })) } + async fn search_chats(&self, args: SearchArgs) -> anyhow::Result { + let query = args.query.trim(); + if query.is_empty() { + anyhow::bail!("query is empty"); + } + let mode = args.mode.as_deref().unwrap_or("words"); + let matcher = match mode { + "words" => Matcher::words(query), + "regex" => Matcher::regex(query, args.case_sensitive), + other => anyhow::bail!("unknown mode {other:?} (use words or regex)"), + } + .map_err(anyhow::Error::msg)?; + let role = match args.role.as_deref() { + None => None, + Some(r @ ("user" | "assistant")) => Some(r), + Some(other) => anyhow::bail!("unknown role {other:?} (use user or assistant)"), + }; + + let (mut chats, spaces, sessions) = self + .filtered_chats( + args.project.as_deref(), + args.device.as_deref(), + args.include_archived, + args.parent.as_deref(), + ) + .await?; + if !args.include_self + && let Some(origin) = self.zeron.origin().chat_id.as_deref() + { + chats.retain(|c| c.id != origin); + } + chats.sort_by_key(|c| std::cmp::Reverse(c.last_message_at.unwrap_or(c.created_at))); + let candidates = chats.len(); + chats.truncate( + args.max_chats + .unwrap_or(SEARCH_DEFAULT_SCAN) + .clamp(1, SEARCH_MAX_SCAN), + ); + let attempted = chats.len(); + + let options = RenderOptions { + include_reasoning: args.include_reasoning, + include_tools: args.include_tools.unwrap_or(true), + }; + let (matcher_ref, role) = (&matcher, role); + let fetched: Vec> = futures::stream::iter( + chats, + ) + .map(|chat| async move { + let entries = match self.zeron.transcript(&chat.id).await { + Ok(entries) => entries, + Err(err) => { + tracing::debug!(chat = %chat.id, "search: transcript unavailable: {err}"); + return None; + } + }; + let rendered = render_entries(&entries, options); + let scan = matcher_ref.scan( + chat.title.as_deref(), + rendered + .iter() + .enumerate() + .filter(|(_, m)| role.is_none_or(|r| role_name(m.role) == r)), + ); + let messages = rendered + .into_iter() + .map(|m| (m.id, role_name(m.role), m.created_at_iso)) + .collect(); + Some((ScannedChat { chat, messages }, scan)) + }) + .buffer_unordered(SEARCH_CONCURRENCY) + .collect() + .await; + let (kept, scans): (Vec<_>, Vec<_>) = fetched.into_iter().flatten().unzip(); + + let limit = args.limit.unwrap_or(10).clamp(1, 100); + let snippets = args.snippets.unwrap_or(3).clamp(1, 20); + let min_match = args.min_match.unwrap_or(0.5).clamp(0.0, 1.0); + let (matched_chats, ranked) = rank(&matcher, &scans, min_match, limit, snippets); + let results: Vec = ranked + .iter() + .map(|r| { + let scanned = &kept[r.scan]; + let total = scanned.messages.len(); + json!({ + "chat": summarize_chat(&scanned.chat, &spaces, &sessions), + "score": (f64::from(r.score) * 100.0).round() / 100.0, + "titleMatch": r.title_match, + "matchedTerms": r.matched, + "messageCount": total, + "matches": r.hits.iter().map(|h| { + let (id, role, at) = &scanned.messages[h.index]; + json!({ + "messageId": id, + "role": role, + "at": at, + "offsetFromNewest": total - 1 - h.index, + "matchedTerms": h.matched, + "snippet": h.snippet, + }) + }).collect::>(), + }) + }) + .collect(); + Ok(json!({ + "query": query, + "mode": mode, + "candidateChats": candidates, + "scannedChats": kept.len(), + "unreadableChats": attempted - kept.len(), + "truncated": attempted < candidates, + "matchedChats": matched_chats, + "results": results, + })) + } + async fn send_message(&self, args: SendArgs) -> anyhow::Result { let text = args.text.trim(); if text.is_empty() { @@ -1102,6 +1308,8 @@ mod tests { writes: Mutex>, dispatch_barrier: Option, beta_parent: Option, + /// Replaces the assistant reply in Beta's transcript. + beta_reply: Option, } fn stream(item: Value) -> RpcReply { @@ -1148,13 +1356,19 @@ mod tests { methods::LIST_MODELS => RpcReply::Value(json!([ { "id": "opus", "label": "Opus" }, { "id": "sonnet", "label": "Sonnet" } ])), - methods::WATCH_DOC_MESSAGES => stream(json!({ "reset": [ - { "id": "u1", "role": "user", "createdAt": 1, "deviceId": "dev-local", - "parts": [{ "kind": "text", "id": "t", "text": "hi" }] }, - { "id": "a1", "role": "assistant", "createdAt": 2, "deviceId": "dev-local", - "status": "complete", - "parts": [{ "kind": "text", "id": "t", "text": "hello back" }] } - ]})), + methods::WATCH_DOC_MESSAGES => { + let reply = match (&self.beta_reply, params["chatId"].as_str()) { + (Some(text), Some("chat-beta-2")) => text.as_str(), + _ => "hello back", + }; + stream(json!({ "reset": [ + { "id": "u1", "role": "user", "createdAt": 1, "deviceId": "dev-local", + "parts": [{ "kind": "text", "id": "t", "text": "hi" }] }, + { "id": "a1", "role": "assistant", "createdAt": 2, "deviceId": "dev-local", + "status": "complete", + "parts": [{ "kind": "text", "id": "t", "text": reply }] } + ]})) + } methods::MUTATE | methods::QUEUE_COMMAND | methods::QUEUE_MESSAGE => { self.writes .lock() @@ -1219,6 +1433,103 @@ mod tests { assert_eq!(read["messages"][0]["id"], "a1"); } + #[tokio::test] + async fn search_chats_finds_content_loosely_and_by_regex() { + let world = Arc::new(World { + beta_reply: Some("The idle reaper kills persistent sessions after 20 minutes".into()), + ..World::default() + }); + let tools = tools(world, Origin::default()); + + // Typo, fragment, and a word that is not in the title: still found. + let found = tools + .call("search_chats", json!({ "query": "persistant reaper" })) + .await + .unwrap(); + assert_eq!(found["scannedChats"], 2); + assert_eq!(found["matchedChats"], 1); + let hit = &found["results"][0]; + assert_eq!(hit["chat"]["title"], "Beta"); + let m = &hit["matches"][0]; + assert_eq!(m["role"], "assistant"); + assert_eq!(m["offsetFromNewest"], 0); + assert!(m["snippet"].as_str().unwrap().contains("idle reaper")); + + // The offset drops straight into read_chat. + let read = tools + .call( + "read_chat", + json!({ "chat": "Beta", "limit": 1, "offset": m["offsetFromNewest"] }), + ) + .await + .unwrap(); + assert_eq!(read["messages"][0]["id"], m["messageId"]); + + let regex = tools + .call( + "search_chats", + json!({ "query": r"\d+ minutes", "mode": "regex" }), + ) + .await + .unwrap(); + assert_eq!(regex["results"][0]["chat"]["title"], "Beta"); + let strict = tools + .call( + "search_chats", + json!({ "query": "REAPER", "mode": "regex", "case_sensitive": true }), + ) + .await + .unwrap(); + assert_eq!(strict["matchedChats"], 0); + + let users = tools + .call("search_chats", json!({ "query": "reaper", "role": "user" })) + .await + .unwrap(); + assert_eq!(users["matchedChats"], 0); + + assert!( + tools + .call("search_chats", json!({ "query": " " })) + .await + .is_err() + ); + assert!( + tools + .call("search_chats", json!({ "query": "(", "mode": "regex" })) + .await + .is_err() + ); + } + + #[tokio::test] + async fn search_skips_the_calling_chat_and_reports_truncation() { + let world = Arc::new(World::default()); + let origin = Origin { + chat_id: Some("chat-alpha-1".into()), + ..Origin::default() + }; + let tools = tools(world, origin); + let found = tools + .call("search_chats", json!({ "query": "hello", "max_chats": 1 })) + .await + .unwrap(); + // Alpha is the caller; Beta is the only candidate and fits the cap. + assert_eq!(found["candidateChats"], 1); + assert_eq!(found["truncated"], false); + assert_eq!(found["results"][0]["chat"]["title"], "Beta"); + + let all = tools + .call( + "search_chats", + json!({ "query": "hello", "include_self": true, "max_chats": 1 }), + ) + .await + .unwrap(); + assert_eq!(all["candidateChats"], 2); + assert_eq!(all["truncated"], true); + } + #[tokio::test] async fn send_attributes_and_refuses_self() { let world = Arc::new(World::default()); diff --git a/docs/mcp.md b/docs/mcp.md index 335c155c59..12bacd9c66 100644 --- a/docs/mcp.md +++ b/docs/mcp.md @@ -110,12 +110,26 @@ name (default: the local engine's device). | `create_chats` | Concurrent `create_chat` requests with per-request results | | `send_messages` | Concurrent `send_message` requests with per-request results | | `read_chat` | `WatchDocMessages` opening `reset` frame, rendered | +| `search_chats` | `WatchChats` + per-chat `WatchDocMessages`, scored locally | | `send_message` | `QueueCommand` Run / Steer, or `QueueMessage` | | `wait_for_turn` | `WatchSessions` until the chat settles | | `interrupt_chat` | `QueueCommand` Interrupt | | `respond_to_input` | `QueueCommand` RespondInput | | `archive_chat` | `Mutate setChatArchived` | +`search_chats` searches transcript content, not titles alone. There is no +engine-side index: it reads each candidate chat's transcript the way +`read_chat` does (8 at a time, newest activity first, capped by `max_chats`, +default 200) and scores it in the MCP process. `words` mode tokenizes the query +and scores every message per term — exact word 1.0, substring 0.8, typo within +the length-scaled edit budget 0.6/0.5 — weighting each term by rarity across +the scanned chats; a chat is listed when it covers `min_match` (default 0.5) of +the query, with terms allowed to land in different messages. `regex` mode is +ripgrep-style. Hits carry a snippet and `offsetFromNewest`, which is the +`read_chat` offset that lands on that message. The calling chat is skipped +unless `include_self`; the reply reports `truncated` and `unreadableChats` so a +partial scan is never mistaken for a complete one. + Watch streams are the engine's only read surface (there is no one-shot "get transcript" RPC); a snapshot is "subscribe, take the first item, drop" — drop cancels server-side, exactly what the sidebar does on attach.