nexo_core/agent/
transcripts_index.rs

1//! SQLite FTS5 index over transcript content. Source of truth stays
2//! the JSONL files written by [`super::transcripts::TranscriptWriter`];
3//! this index is derivable and can be rebuilt from disk via
4//! [`TranscriptsIndex::rebuild_for_agent`].
5//!
6//! The index is shared across agents — cross-agent isolation is
7//! enforced by the `agent_id` filter on every query.
8
9use anyhow::Context;
10use sqlx::sqlite::{SqliteConnectOptions, SqlitePoolOptions};
11use sqlx::{Row, SqlitePool};
12use std::path::{Path, PathBuf};
13use std::str::FromStr;
14use uuid::Uuid;
15
16use super::transcripts::{TranscriptLine, TranscriptWriter};
17
18#[derive(Debug, Clone)]
19pub struct IndexedHit {
20    pub session_id: Uuid,
21    pub agent_id: String,
22    pub timestamp_unix: i64,
23    pub role: String,
24    pub source_plugin: String,
25    /// FTS5 `snippet()` excerpt, with `[`/`]` markers around matched
26    /// terms. Capped at ~120 chars.
27    pub snippet: String,
28}
29
30#[derive(Debug, Clone)]
31pub struct TranscriptsIndex {
32    pool: SqlitePool,
33}
34
35impl TranscriptsIndex {
36    /// Open or create the index DB at `path`. Idempotent.
37    pub async fn open(path: &Path) -> anyhow::Result<Self> {
38        if let Some(parent) = path.parent() {
39            tokio::fs::create_dir_all(parent).await.ok();
40        }
41        let url = format!("sqlite://{}", path.display());
42        let opts = SqliteConnectOptions::from_str(&url)
43            .with_context(|| format!("invalid sqlite url for {}", path.display()))?
44            .create_if_missing(true)
45            .journal_mode(sqlx::sqlite::SqliteJournalMode::Wal);
46        let pool = SqlitePoolOptions::new()
47            .max_connections(4)
48            .connect_with(opts)
49            .await
50            .with_context(|| format!("opening transcripts index at {}", path.display()))?;
51        // FTS5 virtual table — agent_id et al. are UNINDEXED so they
52        // ride along but aren't tokenized.
53        sqlx::query(
54            "CREATE VIRTUAL TABLE IF NOT EXISTS transcripts_fts USING fts5(
55                content,
56                agent_id        UNINDEXED,
57                session_id      UNINDEXED,
58                timestamp_unix  UNINDEXED,
59                role            UNINDEXED,
60                source_plugin   UNINDEXED,
61                tokenize = 'unicode61 remove_diacritics 2'
62            )",
63        )
64        .execute(&pool)
65        .await
66        .context("creating transcripts_fts table")?;
67        Ok(Self { pool })
68    }
69
70    pub async fn insert(
71        &self,
72        agent_id: &str,
73        session_id: Uuid,
74        timestamp_unix: i64,
75        role: &str,
76        source_plugin: &str,
77        content: &str,
78    ) -> anyhow::Result<()> {
79        sqlx::query(
80            "INSERT INTO transcripts_fts (content, agent_id, session_id, timestamp_unix, role, source_plugin) \
81             VALUES (?, ?, ?, ?, ?, ?)",
82        )
83        .bind(content)
84        .bind(agent_id)
85        .bind(session_id.to_string())
86        .bind(timestamp_unix)
87        .bind(role)
88        .bind(source_plugin)
89        .execute(&self.pool)
90        .await
91        .context("inserting into transcripts_fts")?;
92        Ok(())
93    }
94
95    /// FTS5 MATCH search filtered by `agent_id`. `query` is escaped as
96    /// a single phrase so quotes / special chars in user input cannot
97    /// inject FTS operators.
98    pub async fn search(
99        &self,
100        agent_id: &str,
101        query: &str,
102        limit: usize,
103    ) -> anyhow::Result<Vec<IndexedHit>> {
104        let limit = limit.clamp(1, 500) as i64;
105        let phrase = escape_fts_phrase(query);
106        let rows = sqlx::query(
107            "SELECT session_id, agent_id, timestamp_unix, role, source_plugin, \
108                    snippet(transcripts_fts, 0, '[', ']', '...', 12) AS snip \
109             FROM transcripts_fts \
110             WHERE agent_id = ? AND content MATCH ? \
111             ORDER BY rank \
112             LIMIT ?",
113        )
114        .bind(agent_id)
115        .bind(&phrase)
116        .bind(limit)
117        .fetch_all(&self.pool)
118        .await
119        .context("querying transcripts_fts")?;
120        let mut out = Vec::with_capacity(rows.len());
121        for row in rows {
122            let sid: String = row.try_get("session_id")?;
123            let session_id = Uuid::parse_str(&sid).unwrap_or_else(|_| Uuid::nil());
124            out.push(IndexedHit {
125                session_id,
126                agent_id: row.try_get("agent_id")?,
127                timestamp_unix: row.try_get("timestamp_unix")?,
128                role: row.try_get("role")?,
129                source_plugin: row.try_get("source_plugin")?,
130                snippet: row.try_get("snip")?,
131            });
132        }
133        Ok(out)
134    }
135
136    pub async fn count_for_agent(&self, agent_id: &str) -> anyhow::Result<u64> {
137        let row = sqlx::query("SELECT COUNT(*) AS n FROM transcripts_fts WHERE agent_id = ?")
138            .bind(agent_id)
139            .fetch_one(&self.pool)
140            .await
141            .context("counting transcripts_fts")?;
142        let n: i64 = row.try_get("n")?;
143        Ok(n.max(0) as u64)
144    }
145
146    /// Wipe and rebuild every row owned by `agent_id` from JSONL files
147    /// found under `transcripts_root`. Returns the number of entries
148    /// indexed.
149    pub async fn rebuild_for_agent(
150        &self,
151        agent_id: &str,
152        transcripts_root: &Path,
153    ) -> anyhow::Result<usize> {
154        // Walk the JSONL files first (read-only; safe outside the
155        // transaction). Then DELETE + INSERT inside a single
156        // transaction so a crash mid-rebuild leaves the index
157        // either fully old or fully new — never half-empty.
158        let writer = TranscriptWriter::new(PathBuf::from(transcripts_root), agent_id.to_string());
159        let mut entries = match tokio::fs::read_dir(transcripts_root).await {
160            Ok(d) => d,
161            Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
162                // Empty root → just clear the index for this agent.
163                sqlx::query("DELETE FROM transcripts_fts WHERE agent_id = ?")
164                    .bind(agent_id)
165                    .execute(&self.pool)
166                    .await
167                    .context("wiping rows for agent")?;
168                return Ok(0);
169            }
170            Err(e) => {
171                return Err(anyhow::anyhow!(
172                    "read transcripts_root {}: {e}",
173                    transcripts_root.display()
174                ))
175            }
176        };
177
178        // Collect everything we want to insert before touching the DB.
179        struct Row {
180            sid: Uuid,
181            ts: i64,
182            role: &'static str,
183            source_plugin: String,
184            content: String,
185        }
186        let mut rows: Vec<Row> = Vec::new();
187        while let Some(entry) = entries.next_entry().await? {
188            let p = entry.path();
189            if p.extension().and_then(|s| s.to_str()) != Some("jsonl") {
190                continue;
191            }
192            let stem = p
193                .file_stem()
194                .and_then(|s| s.to_str())
195                .unwrap_or_default()
196                .to_string();
197            let Ok(sid) = Uuid::parse_str(&stem) else {
198                continue;
199            };
200            let lines = writer.read_session(sid).await.unwrap_or_default();
201            for line in lines {
202                if let TranscriptLine::Entry(e) = line {
203                    rows.push(Row {
204                        sid,
205                        ts: e.timestamp.timestamp(),
206                        role: role_str(e.role),
207                        source_plugin: e.source_plugin,
208                        content: e.content,
209                    });
210                }
211            }
212        }
213
214        let mut tx = self.pool.begin().await.context("begin tx")?;
215        sqlx::query("DELETE FROM transcripts_fts WHERE agent_id = ?")
216            .bind(agent_id)
217            .execute(&mut *tx)
218            .await
219            .context("wiping rows for agent")?;
220        let mut indexed = 0usize;
221        for row in &rows {
222            sqlx::query(
223                "INSERT INTO transcripts_fts (content, agent_id, session_id, timestamp_unix, role, source_plugin) \
224                 VALUES (?, ?, ?, ?, ?, ?)",
225            )
226            .bind(&row.content)
227            .bind(agent_id)
228            .bind(row.sid.to_string())
229            .bind(row.ts)
230            .bind(row.role)
231            .bind(&row.source_plugin)
232            .execute(&mut *tx)
233            .await
234            .context("insert row in rebuild")?;
235            indexed += 1;
236        }
237        tx.commit().await.context("commit rebuild tx")?;
238        Ok(indexed)
239    }
240}
241
242fn role_str(role: super::transcripts::TranscriptRole) -> &'static str {
243    use super::transcripts::TranscriptRole as R;
244    match role {
245        R::User => "user",
246        R::Assistant => "assistant",
247        R::Tool => "tool",
248        R::System => "system",
249    }
250}
251
252/// Escape a user query into a single FTS5 phrase. Embedded `"` are
253/// doubled per FTS5 syntax; the whole result is wrapped in quotes.
254/// This keeps user input from injecting FTS operators (`OR`, `NOT`,
255/// `:`, etc.).
256fn escape_fts_phrase(q: &str) -> String {
257    let mut out = String::with_capacity(q.len() + 2);
258    out.push('"');
259    for ch in q.chars() {
260        if ch == '"' {
261            out.push('"');
262            out.push('"');
263        } else {
264            out.push(ch);
265        }
266    }
267    out.push('"');
268    out
269}
270
271#[cfg(test)]
272mod tests {
273    use super::super::transcripts::{TranscriptEntry, TranscriptRole, TranscriptWriter};
274    use super::*;
275    use chrono::Utc;
276    use tempfile::tempdir;
277
278    async fn fresh() -> TranscriptsIndex {
279        let dir = tempdir().unwrap().keep();
280        let path = dir.join("transcripts.db");
281        TranscriptsIndex::open(&path).await.unwrap()
282    }
283
284    #[tokio::test]
285    async fn open_is_idempotent() {
286        let dir = tempdir().unwrap().keep();
287        let path = dir.join("transcripts.db");
288        TranscriptsIndex::open(&path).await.unwrap();
289        TranscriptsIndex::open(&path).await.unwrap();
290    }
291
292    #[tokio::test]
293    async fn insert_then_search_round_trip() {
294        let idx = fresh().await;
295        let sid = Uuid::new_v4();
296        idx.insert("kate", sid, 1000, "user", "wa", "hola codigo de seguridad")
297            .await
298            .unwrap();
299        idx.insert("kate", sid, 1001, "assistant", "wa", "no comparto eso")
300            .await
301            .unwrap();
302        let hits = idx.search("kate", "codigo", 10).await.unwrap();
303        assert_eq!(hits.len(), 1);
304        assert_eq!(hits[0].session_id, sid);
305        assert!(hits[0].snippet.contains("[codigo]"));
306    }
307
308    #[tokio::test]
309    async fn search_filters_by_agent() {
310        let idx = fresh().await;
311        let s1 = Uuid::new_v4();
312        let s2 = Uuid::new_v4();
313        idx.insert("kate", s1, 1, "user", "wa", "shared word here")
314            .await
315            .unwrap();
316        idx.insert("ana", s2, 2, "user", "wa", "shared word here too")
317            .await
318            .unwrap();
319        let kate_hits = idx.search("kate", "shared", 10).await.unwrap();
320        assert_eq!(kate_hits.len(), 1);
321        assert_eq!(kate_hits[0].agent_id, "kate");
322        let ana_hits = idx.search("ana", "shared", 10).await.unwrap();
323        assert_eq!(ana_hits.len(), 1);
324        assert_eq!(ana_hits[0].agent_id, "ana");
325    }
326
327    #[tokio::test]
328    async fn count_for_agent_returns_inserted_count() {
329        let idx = fresh().await;
330        let sid = Uuid::new_v4();
331        for i in 0..3 {
332            idx.insert("kate", sid, 100 + i, "user", "wa", "row")
333                .await
334                .unwrap();
335        }
336        idx.insert("ana", sid, 200, "user", "wa", "row")
337            .await
338            .unwrap();
339        assert_eq!(idx.count_for_agent("kate").await.unwrap(), 3);
340        assert_eq!(idx.count_for_agent("ana").await.unwrap(), 1);
341    }
342
343    #[tokio::test]
344    async fn search_user_input_with_quotes_is_safe() {
345        let idx = fresh().await;
346        let sid = Uuid::new_v4();
347        idx.insert("kate", sid, 1, "user", "wa", "an OR b")
348            .await
349            .unwrap();
350        // FTS would treat OR as operator without escaping; quoted phrase
351        // matches literal "an OR b" only.
352        let hits = idx.search("kate", "OR", 10).await.unwrap();
353        assert!(hits.iter().all(|h| h.session_id == sid));
354    }
355
356    #[tokio::test]
357    async fn search_user_input_with_fts_operators_is_safe() {
358        // Adversarial inputs: each one would, if naively concatenated
359        // into the MATCH expression, change the query semantics. Phrase
360        // mode (the wrapping `"`) plus our `"` doubling neutralizes
361        // every operator the docs list (NEAR, AND, OR, NOT, `:` field,
362        // `^` prefix, parens). The query is treated as a literal
363        // phrase against `content`, so none of the inputs below should
364        // ever return rows from a different agent or a different
365        // content string.
366        let idx = fresh().await;
367        let sid = Uuid::new_v4();
368        idx.insert("kate", sid, 1, "user", "wa", "title: hello world")
369            .await
370            .unwrap();
371        idx.insert("ana", Uuid::new_v4(), 2, "user", "wa", "agent_id: leak")
372            .await
373            .unwrap();
374        for adversarial in [
375            "field:value",
376            "agent_id:ana",
377            "OR NOT AND NEAR",
378            "\"injection\"",
379            "(hello)",
380            "^prefix",
381        ] {
382            let hits = idx.search("kate", adversarial, 10).await.unwrap();
383            for h in &hits {
384                assert_eq!(
385                    h.agent_id, "kate",
386                    "input `{adversarial}` leaked rows from another agent"
387                );
388            }
389        }
390    }
391
392    #[tokio::test]
393    async fn rebuild_for_agent_indexes_jsonl() {
394        let dir = tempdir().unwrap().keep();
395        let writer = TranscriptWriter::new(&dir, "kate");
396        let sid = Uuid::new_v4();
397        for content in ["hola", "que tal", "todo bien"] {
398            writer
399                .append_entry(
400                    sid,
401                    TranscriptEntry {
402                        timestamp: Utc::now(),
403                        role: TranscriptRole::User,
404                        content: content.to_string(),
405                        message_id: None,
406                        source_plugin: "wa".into(),
407                        sender_id: None,
408                    },
409                )
410                .await
411                .unwrap();
412        }
413        let idx_path = dir.join("idx.db");
414        let idx = TranscriptsIndex::open(&idx_path).await.unwrap();
415        let n = idx.rebuild_for_agent("kate", &dir).await.unwrap();
416        assert_eq!(n, 3);
417        assert_eq!(idx.count_for_agent("kate").await.unwrap(), 3);
418        let hits = idx.search("kate", "tal", 10).await.unwrap();
419        assert_eq!(hits.len(), 1);
420    }
421}
nexo_core/agent/transcripts_index.rs

nexo_core/agent/
transcripts_index.rs
