Find in thread (⌘F) and message search in the command palette

crates/engine/src/lib.rs

@@ -51,6 +51,7 @@ pub mod spaces;
 pub mod terminals;
 pub mod titles;
 mod transcript_history;
+pub mod transcript_search;
 pub mod uploads;
 pub mod voice;
 pub mod workspace_files;
@@ -165,6 +166,8 @@ pub struct EngineCore {
     pub viewport_focus: desktop_notify::ViewportFocus,
     /// Local→synced profile import (account-scoped runtimes only).
     pub local_import: Option<local_import::LocalImporter>,
+    /// `None` when the index cannot open; search is then unavailable.
+    pub transcript_search: Option<transcript_search::TranscriptSearch>,
     workspace_scope: WorkspaceScope,
     /// Auth service (attached by [`Engine::run`]; a lazy dev-mode instance otherwise).
     auth: std::sync::Mutex<Option<Auth>>,
@@ -245,6 +248,7 @@ impl EngineCore {
         registry.load_prefs(data_dir);
         let store = Arc::new(DocsStore::open(profile.store_root())?);
         let store_for_import = store.clone();
+        let store_for_search = store.clone();
         let journal = Arc::new(RunJournal::open(profile.store_root().join("journals"))?);
         let sessions = SessionsEngine::new(device_id.clone(), journal, registry.clone());
         let doc_host = DocHost::new(
@@ -268,6 +272,13 @@ impl EngineCore {
                 local_only: matches!(profile.scope(), WorkspaceScope::Local),
             },
         )?;
+        let transcript_search = transcript_search::TranscriptSearch::open(
+            profile.store_root(),
+            store_for_search,
+            workspace.clone(),
+        )
+        .inspect_err(|error| tracing::warn!(%error, "transcript search unavailable"))
+        .ok();
         doc_host.set_workspace(workspace.clone());
         doc_host.set_sessions(sessions.clone());
         sessions.set_doc_host(doc_host.clone());
@@ -380,6 +391,7 @@ impl EngineCore {
             harness_updates,
             device_id,
             local_import,
+            transcript_search,
             workspace_scope: profile.scope(),
             auth: std::sync::Mutex::new(None),
             links: std::sync::Mutex::new(None),
@@ -535,6 +547,9 @@ impl EngineCore {
         if let Some(broker) = self.credentials() {
             rpc = rpc.with_credentials(broker);
         }
+        if let Some(search) = self.transcript_search.clone() {
+            rpc = rpc.with_transcript_search(search);
+        }
         Arc::new(rpc)
     }
 

crates/engine/src/rpc.rs

@@ -799,6 +799,7 @@ pub struct EngineRpc {
     credentials: Option<crate::credentials::CredentialBroker>,
     /// Viewports' window focus, for desktop banners.
     viewport_focus: crate::desktop_notify::ViewportFocus,
+    transcript_search: Option<crate::transcript_search::TranscriptSearch>,
     engine_info: EngineInfo,
 }
 
@@ -858,6 +859,7 @@ impl EngineRpc {
             local_import: None,
             credentials: None,
             viewport_focus: Default::default(),
+            transcript_search: None,
             engine_info,
         }
     }
@@ -926,6 +928,14 @@ impl EngineRpc {
         self
     }
 
+    pub fn with_transcript_search(
+        mut self,
+        search: crate::transcript_search::TranscriptSearch,
+    ) -> Self {
+        self.transcript_search = Some(search);
+        self
+    }
+
     fn auth(&self) -> Result<&Auth, RpcError> {
         self.auth
             .as_ref()
@@ -3242,6 +3252,18 @@ impl RpcService for EngineRpc {
             methods::WATCH_CHATS => {
                 Ok(RpcReply::Stream(watch_stream(self.workspace.watch_chats())))
             }
+            methods::SEARCH_TRANSCRIPTS => {
+                let request: zeron_proto::SearchTranscriptsRequest = parse_params(params)?;
+                let search = self
+                    .transcript_search
+                    .as_ref()
+                    .ok_or_else(|| RpcError::Failed("transcript search unavailable".into()))?;
+                let hits = search
+                    .search(request.query, request.limit)
+                    .await
+                    .map_err(|e| RpcError::Failed(e.to_string()))?;
+                RpcReply::value(&hits)
+            }
             methods::WATCH_SIDEBAR_PREFERENCES => Ok(RpcReply::Stream(watch_stream(
                 self.workspace.watch_sidebar_preferences(),
             ))),

crates/engine/src/transcript_search.rs

@@ -0,0 +1,267 @@
+//! Full-text search over transcripts stored on this engine.
+//!
+//! The index is derived data. `transcript-search.sqlite3` holds one FTS5 row
+//! per chat and the snapshot `saved_at` it was built from. Refresh diffs the
+//! index against `DocsStore` snapshots of the workspace's live chats, so
+//! anything that writes a snapshot gets indexed with no hook. That covers live
+//! persistence, chat2 seeding, and profile import. Deleting the file only
+//! costs a rebuild.
+
+use std::collections::{HashMap, HashSet};
+use std::path::Path;
+use std::sync::{Arc, Mutex, PoisonError};
+use std::time::{Duration, Instant};
+
+use rusqlite::{Connection, params};
+use zeron_doc::{MessagePart, MessageRole, SessionDoc};
+use zeron_proto::TranscriptSearchHit;
+use zeron_sync::DocsStore;
+
+use crate::EngineError;
+use crate::workspace_host::WorkspaceHost;
+
+const SCHEMA_VERSION: i64 = 2;
+/// A message's FTS rowid is `chat row * MESSAGE_SLOTS + message index`, so a
+/// chat's messages are one rowid range.
+// ponytail: messages past a million in one chat go unindexed.
+const MESSAGE_SLOTS: i64 = 1_000_000;
+const REFRESH_BUDGET: Duration = Duration::from_secs(1);
+const SNIPPET_TOKENS: i64 = 16;
+
+#[derive(Clone)]
+pub struct TranscriptSearch {
+    inner: Arc<Inner>,
+}
+
+struct Inner {
+    index: Mutex<Connection>,
+    refreshing: Mutex<()>,
+    store: Arc<DocsStore>,
+    workspace: WorkspaceHost,
+}
+
+impl TranscriptSearch {
+    pub fn open(
+        dir: &Path,
+        store: Arc<DocsStore>,
+        workspace: WorkspaceHost,
+    ) -> Result<Self, EngineError> {
+        let conn = Connection::open(dir.join("transcript-search.sqlite3")).map_err(sql)?;
+        conn.pragma_update(None, "journal_mode", "WAL")
+            .map_err(sql)?;
+        conn.pragma_update(None, "synchronous", "NORMAL")
+            .map_err(sql)?;
+        let version: i64 = conn
+            .pragma_query_value(None, "user_version", |row| row.get(0))
+            .map_err(sql)?;
+        if version != SCHEMA_VERSION {
+            conn.execute_batch(&format!(
+                "DROP TABLE IF EXISTS chats;
+                 DROP TABLE IF EXISTS transcripts;
+                 CREATE TABLE chats (
+                    id       INTEGER PRIMARY KEY,
+                    chat_id  TEXT NOT NULL UNIQUE,
+                    saved_at INTEGER NOT NULL
+                 ) STRICT;
+                 CREATE VIRTUAL TABLE transcripts USING fts5(
+                    body, from_user UNINDEXED,
+                    tokenize = 'unicode61 remove_diacritics 2'
+                 );
+                 PRAGMA user_version = {SCHEMA_VERSION};"
+            ))
+            .map_err(sql)?;
+        }
+        Ok(Self {
+            inner: Arc::new(Inner {
+                index: Mutex::new(conn),
+                refreshing: Mutex::new(()),
+                store,
+                workspace,
+            }),
+        })
+    }
+
+    /// An empty query only refreshes, so callers can warm the index.
+    pub async fn search(
+        &self,
+        query: String,
+        limit: u16,
+    ) -> Result<Vec<TranscriptSearchHit>, EngineError> {
+        let inner = self.inner.clone();
+        tokio::task::spawn_blocking(move || {
+            if let Err(error) = inner.refresh() {
+                tracing::warn!(%error, "transcript index refresh failed");
+            }
+            match match_expression(&query) {
+                Some(expr) => inner.query(&expr, limit),
+                None => Ok(Vec::new()),
+            }
+        })
+        .await
+        .map_err(|e| EngineError::Other(e.to_string()))?
+    }
+}
+
+impl Inner {
+    /// Reindex chats whose snapshot changed, newest first, within
+    /// [`REFRESH_BUDGET`]. A refresh already in flight makes this a no-op.
+    fn refresh(&self) -> Result<(), EngineError> {
+        let Ok(_refreshing) = self.refreshing.try_lock() else {
+            return Ok(());
+        };
+        let started = Instant::now();
+        let live: HashSet<String> = self
+            .workspace
+            .read_chats()?
+            .into_iter()
+            .map(|chat| chat.id)
+            .collect();
+        let mut wanted: Vec<(String, i64)> = self
+            .store
+            .snapshot_saved_at()?
+            .into_iter()
+            .filter(|(id, _)| live.contains(id))
+            .collect();
+        let indexed: HashMap<String, (i64, i64)> = {
+            let conn = self.conn();
+            let mut stmt = conn
+                .prepare("SELECT chat_id, id, saved_at FROM chats")
+                .map_err(sql)?;
+            stmt.query_map([], |row| Ok((row.get(0)?, (row.get(1)?, row.get(2)?))))
+                .map_err(sql)?
+                .collect::<Result<_, _>>()
+                .map_err(sql)?
+        };
+        let wanted_ids: HashSet<&str> = wanted.iter().map(|(id, _)| id.as_str()).collect();
+        let gone: Vec<i64> = indexed
+            .iter()
+            .filter(|(chat_id, _)| !wanted_ids.contains(chat_id.as_str()))
+            .map(|(_, (row, _))| *row)
+            .collect();
+        if !gone.is_empty() {
+            let mut conn = self.conn();
+            let tx = conn.transaction().map_err(sql)?;
+            for row in gone {
+                delete_messages(&tx, row)?;
+                tx.execute("DELETE FROM chats WHERE id = ?1", params![row])
+                    .map_err(sql)?;
+            }
+            tx.commit().map_err(sql)?;
+        }
+        wanted.retain(|(id, saved_at)| indexed.get(id).map(|(_, at)| at) != Some(saved_at));
+        wanted.sort_by_key(|(_, saved_at)| std::cmp::Reverse(*saved_at));
+        for (chat_id, saved_at) in wanted {
+            if started.elapsed() > REFRESH_BUDGET {
+                break;
+            }
+            let messages = match self.store.load_snapshot(&chat_id) {
+                Ok(Some(bytes)) => transcript_messages(&bytes).unwrap_or_else(|error| {
+                    tracing::warn!(chat = %chat_id, %error, "transcript unreadable; indexed empty");
+                    Vec::new()
+                }),
+                Ok(None) => continue,
+                Err(error) => return Err(error.into()),
+            };
+            let mut conn = self.conn();
+            let tx = conn.transaction().map_err(sql)?;
+            let row: i64 = tx
+                .query_row(
+                    "INSERT INTO chats (chat_id, saved_at) VALUES (?1, ?2)
+                     ON CONFLICT(chat_id) DO UPDATE SET saved_at = excluded.saved_at
+                     RETURNING id",
+                    params![chat_id, saved_at],
+                    |row| row.get(0),
+                )
+                .map_err(sql)?;
+            delete_messages(&tx, row)?;
+            for (ix, (from_user, body)) in (0..MESSAGE_SLOTS).zip(messages) {
+                tx.execute(
+                    "INSERT INTO transcripts (rowid, body, from_user) VALUES (?1, ?2, ?3)",
+                    params![row * MESSAGE_SLOTS + ix, body, from_user],
+                )
+                .map_err(sql)?;
+            }
+            tx.commit().map_err(sql)?;
+        }
+        Ok(())
+    }
+
+    fn query(&self, expr: &str, limit: u16) -> Result<Vec<TranscriptSearchHit>, EngineError> {
+        let conn = self.conn();
+        let mut stmt = conn
+            .prepare_cached(
+                "SELECT chats.chat_id, snippet(transcripts, 0, '', '', '…', ?3),
+                        transcripts.from_user
+                 FROM transcripts
+                 JOIN chats ON chats.id = transcripts.rowid / ?4
+                 WHERE transcripts MATCH ?1
+                 ORDER BY rank
+                 LIMIT ?2",
+            )
+            .map_err(sql)?;
+        stmt.query_map(params![expr, limit, SNIPPET_TOKENS, MESSAGE_SLOTS], |row| {
+            Ok(TranscriptSearchHit {
+                chat_id: row.get(0)?,
+                snippet: row.get(1)?,
+                from_user: row.get(2)?,
+            })
+        })
+        .map_err(sql)?
+        .collect::<Result<_, _>>()
+        .map_err(sql)
+    }
+
+    fn conn(&self) -> std::sync::MutexGuard<'_, Connection> {
+        self.index.lock().unwrap_or_else(PoisonError::into_inner)
+    }
+}
+
+fn delete_messages(tx: &rusqlite::Transaction<'_>, chat_row: i64) -> Result<(), EngineError> {
+    tx.execute(
+        "DELETE FROM transcripts WHERE rowid >= ?1 AND rowid < ?2",
+        params![chat_row * MESSAGE_SLOTS, (chat_row + 1) * MESSAGE_SLOTS],
+    )
+    .map_err(sql)?;
+    Ok(())
+}
+
+/// `(from user, text)` per user or assistant message. Reasoning is hidden in
+/// the UI, and tool output would drown real matches.
+fn transcript_messages(snapshot: &[u8]) -> Result<Vec<(bool, String)>, EngineError> {
+    let raw = loro::LoroDoc::new();
+    raw.import(snapshot)
+        .map_err(|e| EngineError::Other(e.to_string()))?;
+    let entries = SessionDoc::from_doc(raw).read_entries()?;
+    let mut messages = Vec::new();
+    for entry in entries {
+        if !matches!(entry.role, MessageRole::User | MessageRole::Assistant) {
+            continue;
+        }
+        let mut body = String::new();
+        for part in entry.parts {
+            if let MessagePart::Text { text, .. } = part {
+                body.push_str(&text);
+                body.push('\n');
+            }
+        }
+        if !body.trim().is_empty() {
+            messages.push((entry.role == MessageRole::User, body));
+        }
+    }
+    Ok(messages)
+}
+
+/// Every word must appear, each as a quoted prefix so FTS5 operators in user
+/// text stay literal. Words without a letter or digit have no tokens to match.
+fn match_expression(query: &str) -> Option<String> {
+    let terms: Vec<String> = query
+        .split_whitespace()
+        .filter(|word| word.chars().any(char::is_alphanumeric))
+        .map(|word| format!("\"{}\"*", word.replace('"', "\"\"")))
+        .collect();
+    (!terms.is_empty()).then(|| terms.join(" "))
+}
+
+fn sql(error: rusqlite::Error) -> EngineError {
+    EngineError::Other(format!("transcript index: {error}"))
+}