swink-agent-memory 0.14.3

//! Tantivy-backed full-text index for cross-session search.
//!
//! The index is stored in `<sessions_dir>/.search_index/` alongside the JSONL
//! session files.  It is rebuilt (or opened) lazily on the first `search()`
//! call, replaced wholesale per session via [`TantivyIndex::index_session`]
//! (rewrite paths), and extended incrementally via
//! [`TantivyIndex::append_session_entries`] (append paths).
//!
//! ## Writer lifecycle
//!
//! tantivy permits a single [`IndexWriter`] per index directory (it holds the
//! directory write lock), and creating one allocates the full indexer heap.
//! `TantivyIndex` therefore keeps one long-lived writer, created lazily on
//! the first mutation and reused for every subsequent one, instead of paying
//! writer creation on every save/append.  A consequence is that once a
//! `TantivyIndex` has performed a mutation it holds the directory write lock
//! until it is dropped, so only one live `TantivyIndex` may mutate a given
//! index directory at a time (searches are unaffected — they use readers).
//!
//! ## Schema
//!
//! | Field          | Type | Options        | Purpose                                   |
//! |----------------|------|----------------|-------------------------------------------|
//! | `session_id`   | text | STORED, raw    | Filter / reconstruct hit                  |
//! | `session_title`| text | STORED, raw    | Display title                             |
//! | `entry_type`   | text | STORED, raw    | Type discriminator for filter             |
//! | `timestamp`    | u64  | INDEXED, STORED| Timestamp range filter                    |
//! | `body`         | text | TEXT (en_stem) | Tokenised full-text content               |
//! | `entry_json`   | text | STORED         | Raw JSON for roundtrip to `SessionEntry`  |

use std::io;
use std::path::Path;
use std::sync::{Arc, Mutex};

use tantivy::collector::TopDocs;
use tantivy::query::QueryParser;
use tantivy::schema::{
    Field, IndexRecordOption, NumericOptions, Schema, TEXT, TextFieldIndexing, TextOptions, Value,
};
use tantivy::{Index, IndexWriter, ReloadPolicy, TantivyDocument};

use crate::entry::SessionEntry;
use crate::meta::SessionMeta;
use crate::search::{self, SessionHit, SessionSearchOptions, snippet_from_text};

/// Size of the indexer heap in bytes.
const INDEXER_HEAP_BYTES: usize = 15_000_000;

/// Sub-directory name inside the sessions directory.
const INDEX_SUBDIR: &str = ".search_index";

/// All tantivy fields used by [`TantivyIndex`].
#[derive(Clone, Copy)]
struct Fields {
    session_id: Field,
    session_title: Field,
    entry_type: Field,
    timestamp: Field,
    body: Field,
    entry_json: Field,
}

impl Fields {
    fn register(builder: &mut tantivy::schema::SchemaBuilder) -> Self {
        // `raw` tokenizer: no tokenization — stored as a single term.
        let raw = TextOptions::default()
            .set_indexing_options(
                TextFieldIndexing::default()
                    .set_tokenizer("raw")
                    .set_index_option(IndexRecordOption::Basic),
            )
            .set_stored();

        let stored_only = TextOptions::default().set_stored();

        Self {
            session_id: builder.add_text_field("session_id", raw.clone()),
            session_title: builder.add_text_field("session_title", stored_only.clone()),
            entry_type: builder.add_text_field("entry_type", raw),
            timestamp: builder.add_u64_field(
                "timestamp",
                NumericOptions::default().set_indexed().set_stored(),
            ),
            body: builder.add_text_field("body", TEXT),
            entry_json: builder.add_text_field("entry_json", stored_only),
        }
    }
}

/// Tantivy-backed full-text index.
///
/// `TantivyIndex` is internally `Arc<Mutex<_>>` so it can be shared across
/// threads and cloned cheaply.  Mutations reuse one long-lived
/// [`IndexWriter`], so a `TantivyIndex` that has mutated holds the index
/// directory's write lock until dropped (see the module docs).
#[derive(Clone)]
pub struct TantivyIndex {
    inner: Arc<Mutex<Inner>>,
}

struct Inner {
    index: Index,
    fields: Fields,
    /// Long-lived writer, created lazily on the first mutation and reused for
    /// every subsequent one (see the module docs on writer lifecycle).
    writer: Option<IndexWriter>,
    /// Number of full session (re)index operations performed. Test-only
    /// instrumentation proving append paths never fall back to
    /// whole-session rebuilds.
    #[cfg(test)]
    full_reindex_count: usize,
}

impl Inner {
    /// Return the long-lived writer, creating it on first use.
    fn writer(&mut self) -> io::Result<&mut IndexWriter> {
        if self.writer.is_none() {
            self.writer = Some(
                self.index
                    .writer(INDEXER_HEAP_BYTES)
                    .map_err(tantivy_to_io)?,
            );
        }
        Ok(self.writer.as_mut().expect("just created"))
    }

    /// Commit pending writer operations.
    ///
    /// On failure the writer is dropped so the next mutation starts from a
    /// fresh writer instead of reusing one in an unknown state.
    fn commit(&mut self) -> io::Result<()> {
        let writer = self.writer()?;
        if let Err(err) = writer.commit() {
            self.writer = None;
            return Err(tantivy_to_io(err));
        }
        Ok(())
    }
}

impl TantivyIndex {
    /// Open an existing index or create a new one in `<sessions_dir>/.search_index/`.
    pub fn open_or_create(sessions_dir: &Path) -> io::Result<Self> {
        let index_dir = sessions_dir.join(INDEX_SUBDIR);
        std::fs::create_dir_all(&index_dir)?;

        let (index, fields) = open_or_create_index(&index_dir)?;
        Ok(Self {
            inner: Arc::new(Mutex::new(Inner {
                index,
                fields,
                writer: None,
                #[cfg(test)]
                full_reindex_count: 0,
            })),
        })
    }

    /// Index (or re-index) one session.
    ///
    /// Deletes any previously stored documents for this session ID, then
    /// inserts one document per `SessionEntry`.  This is O(session); append
    /// paths should prefer [`TantivyIndex::append_session_entries`].
    pub fn index_session(&self, meta: &SessionMeta, entries: &[SessionEntry]) -> io::Result<()> {
        let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        #[cfg(test)]
        {
            inner.full_reindex_count += 1;
        }
        let fields = inner.fields;
        let writer = inner.writer()?;
        delete_session_docs(writer, &fields, &meta.id);
        add_entry_docs(writer, &fields, meta, entries)?;
        inner.commit()
    }

    /// Add documents for newly appended entries without touching the
    /// session's existing documents.
    ///
    /// Used by append paths, where everything previously indexed for the
    /// session is still valid and only the new `entries` need indexing —
    /// O(new entries) instead of the O(session) cost of
    /// [`TantivyIndex::index_session`].  Rewrite paths must keep using
    /// [`TantivyIndex::index_session`] for its replacement semantics.
    pub fn append_session_entries(
        &self,
        meta: &SessionMeta,
        entries: &[SessionEntry],
    ) -> io::Result<()> {
        if entries.is_empty() {
            return Ok(());
        }
        let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        let fields = inner.fields;
        let writer = inner.writer()?;
        add_entry_docs(writer, &fields, meta, entries)?;
        inner.commit()
    }

    /// Delete every document in the index, leaving it empty.
    ///
    /// Used by `rebuild_search_index` to guarantee replacement semantics when
    /// sessions have been removed outside the store API.
    pub fn clear_all(&self) -> io::Result<()> {
        let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        inner
            .writer()?
            .delete_all_documents()
            .map_err(tantivy_to_io)?;
        inner.commit()
    }

    /// Remove all indexed documents for a session.
    pub fn delete_session(&self, session_id: &str) -> io::Result<()> {
        let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        let fields = inner.fields;
        let writer = inner.writer()?;
        delete_session_docs(writer, &fields, session_id);
        inner.commit()
    }

    /// Number of full session (re)index operations performed via
    /// [`TantivyIndex::index_session`] over this index's lifetime.
    #[cfg(test)]
    pub(crate) fn full_reindex_count(&self) -> usize {
        self.inner
            .lock()
            .unwrap_or_else(|e| e.into_inner())
            .full_reindex_count
    }

    /// Search the index and return hits respecting `options`.
    pub fn search(
        &self,
        query: &str,
        options: &SessionSearchOptions,
    ) -> io::Result<Vec<SessionHit>> {
        let limit = options.limit();
        if query.trim().is_empty() || limit == 0 {
            return Ok(Vec::new());
        }

        let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        let reader = inner
            .index
            .reader_builder()
            .reload_policy(ReloadPolicy::Manual)
            .try_into()
            .map_err(tantivy_to_io)?;

        reader.reload().map_err(tantivy_to_io)?;
        let searcher = reader.searcher();

        let query_parser = QueryParser::for_index(&inner.index, vec![inner.fields.body]);
        // Sanitise the query: tantivy's default parser raises on many special chars.
        let sanitised = sanitise_query(query);
        let parsed_query = query_parser
            .parse_query(&sanitised)
            .map_err(|e| io::Error::new(io::ErrorKind::InvalidInput, e.to_string()))?;

        // Fetch more than `limit` to allow post-filtering.
        let fetch = (limit * 4).max(200);
        let top_docs: Vec<(tantivy::Score, tantivy::DocAddress)> = searcher
            .search(&parsed_query, &TopDocs::with_limit(fetch).order_by_score())
            .map_err(tantivy_to_io)?;

        let mut hits: Vec<SessionHit> = Vec::new();
        for (score, doc_address) in top_docs {
            if hits.len() >= limit {
                break;
            }

            let doc: TantivyDocument = searcher.doc(doc_address).map_err(tantivy_to_io)?;

            let session_id = get_text_field(&doc, inner.fields.session_id);
            let session_title = get_text_field(&doc, inner.fields.session_title);
            let entry_json = get_text_field(&doc, inner.fields.entry_json);
            let entry_type_name = get_text_field(&doc, inner.fields.entry_type);
            let timestamp_val: u64 = doc
                .get_first(inner.fields.timestamp)
                .and_then(|v| v.as_u64())
                .unwrap_or(0);

            // Session ID filter.
            if options
                .session_ids
                .as_ref()
                .is_some_and(|ids| !ids.contains(&session_id))
            {
                continue;
            }
            // Entry type filter.
            if options
                .entry_types
                .as_ref()
                .is_some_and(|types| !types.contains(&entry_type_name))
            {
                continue;
            }
            // Timestamp range filter.
            if options
                .start_time
                .is_some_and(|start| timestamp_val < start.timestamp().cast_unsigned())
            {
                continue;
            }
            if options
                .end_time
                .is_some_and(|end| timestamp_val > end.timestamp().cast_unsigned())
            {
                continue;
            }

            // Reconstruct `SessionEntry` from stored JSON.
            let entry: SessionEntry = match serde_json::from_str(&entry_json) {
                Ok(e) => e,
                Err(err) => {
                    tracing::warn!(
                        session_id = %session_id,
                        error = %err,
                        "skipping tantivy hit with undeserializable entry_json"
                    );
                    continue;
                }
            };

            let body = search::searchable_text_pub(&entry);
            let snippet = snippet_from_text(&body, 0);

            // Convert tantivy's f32 score to a comparable usize (×1000 for ordering).
            // score is always non-negative (BM25 score from tantivy).
            #[allow(clippy::cast_sign_loss)]
            let score = (score.max(0.0) * 1000.0) as usize;
            hits.push(SessionHit::new(
                session_id,
                session_title,
                entry,
                score,
                snippet,
            ));
        }

        Ok(hits)
    }
}

// ─── helpers ────────────────────────────────────────────────────────────────

fn open_or_create_index(index_dir: &Path) -> io::Result<(Index, Fields)> {
    let mut schema_builder = Schema::builder();
    let fields = Fields::register(&mut schema_builder);
    let schema = schema_builder.build();

    let dir = tantivy::directory::MmapDirectory::open(index_dir).map_err(tantivy_to_io)?;
    let index = Index::open_or_create(dir, schema).map_err(tantivy_to_io)?;
    Ok((index, fields))
}

fn delete_session_docs(writer: &IndexWriter, fields: &Fields, session_id: &str) {
    let term = tantivy::Term::from_field_text(fields.session_id, session_id);
    writer.delete_term(term);
}

/// Stage one document per non-empty entry (does not commit).
fn add_entry_docs(
    writer: &IndexWriter,
    fields: &Fields,
    meta: &SessionMeta,
    entries: &[SessionEntry],
) -> io::Result<()> {
    for entry in entries {
        let body = search::searchable_text_pub(entry);
        if body.is_empty() {
            continue;
        }
        let entry_json = serde_json::to_string(entry).map_err(io::Error::other)?;
        let timestamp = entry.timestamp().unwrap_or(0);

        let doc = tantivy::doc!(
            fields.session_id => meta.id.as_str(),
            fields.session_title => meta.title.as_str(),
            fields.entry_type => entry.entry_type_name(),
            fields.timestamp => timestamp,
            fields.body => body.as_str(),
            fields.entry_json => entry_json.as_str()
        );
        writer.add_document(doc).map_err(tantivy_to_io)?;
    }
    Ok(())
}

fn get_text_field(doc: &TantivyDocument, field: Field) -> String {
    doc.get_first(field)
        .and_then(|v| v.as_str())
        .unwrap_or("")
        .to_string()
}

/// Remove characters that confuse tantivy's default query parser when users
/// supply arbitrary search queries (e.g. colons, slashes, brackets).
fn sanitise_query(query: &str) -> String {
    query
        .chars()
        .map(|c| {
            if c.is_alphanumeric() || c == ' ' || c == '-' || c == '_' {
                c
            } else {
                ' '
            }
        })
        .collect()
}

fn tantivy_to_io(e: impl std::fmt::Display) -> io::Error {
    io::Error::other(e.to_string())
}