ai_session_search/
message_search_batches.rs

1// SPDX-FileCopyrightText: 2026 Andrew Hundt
2// SPDX-License-Identifier: Apache-2.0
3
4//! Bounded-retention exhaustive traversal for canonical message search.
5//!
6//! [`SessionSearch::message_search_batches`](crate::service::SessionSearch::message_search_batches)
7//! remains the public entry point. This module owns the dedicated read-only worker, snapshot,
8//! cancellation, batch handoff, and terminal metadata lifecycle.
9
10use std::num::NonZeroUsize;
11use std::sync::atomic::{AtomicBool, Ordering};
12use std::sync::{mpsc, Arc};
13
14use anyhow::{anyhow, bail, Result};
15
16use crate::config::Config;
17use crate::db::{Db, ReadSnapshotCleanupError};
18use crate::message_search::{
19    MessageSearchIncludedData, MessageSearchOrigins, MessageSearchPageDocument,
20    MessageSearchReceiptDocument, MessageSearchRequest, MessageSearchResultDocument,
21    MessageSearchRuntimeDiagnostics, PageInfo, ResolvedMessageSearchRequest, SearchSurface,
22};
23use crate::models::{MessageHit, SearchExplain};
24use crate::search_scope::EffectiveAccessScope;
25use crate::service::{MessageSearchBatchControl, MessageService};
26
27/// One owned, fully enriched result batch. The aligned context window at index `i` belongs to
28/// result `i`; included session data is a mergeable delta for sessions first encountered in this
29/// batch and is not repeated in later batches from the same request.
30#[derive(Debug)]
31pub struct MessageSearchBatch {
32    pub(crate) results: Vec<crate::message_search::MessageSearchHit>,
33    pub(crate) context_windows: Vec<Vec<MessageHit>>,
34    pub(crate) included: MessageSearchIncludedData,
35}
36
37impl MessageSearchBatch {
38    pub fn results(&self) -> &[crate::message_search::MessageSearchHit] {
39        &self.results
40    }
41
42    pub fn context_windows(&self) -> &[Vec<MessageHit>] {
43        &self.context_windows
44    }
45
46    pub const fn included(&self) -> &MessageSearchIncludedData {
47        &self.included
48    }
49
50    /// Borrow one canonical semantic result without building an adapter-owned dictionary.
51    pub fn result_document<'a>(
52        &'a self,
53        request: &'a ResolvedMessageSearchRequest,
54        index: usize,
55    ) -> Option<MessageSearchResultDocument<'a>> {
56        self.results.get(index).map(|hit| {
57            MessageSearchResultDocument::from_batch_parts(
58                hit,
59                request,
60                self.context_windows.get(index).map(Vec::as_slice),
61            )
62        })
63    }
64
65    pub fn into_parts(
66        self,
67    ) -> (
68        Vec<crate::message_search::MessageSearchHit>,
69        Vec<Vec<MessageHit>>,
70        MessageSearchIncludedData,
71    ) {
72        (self.results, self.context_windows, self.included)
73    }
74}
75
76/// Terminal metadata produced only after a batch request naturally exhausts its result set.
77#[derive(Debug)]
78pub struct MessageSearchCompletion {
79    page: PageInfo,
80    planner: Option<SearchExplain>,
81    origins: Option<MessageSearchOrigins>,
82    ordered_digest: Option<String>,
83}
84
85impl MessageSearchCompletion {
86    pub const fn page(&self) -> PageInfo {
87        self.page
88    }
89
90    pub const fn search_explanation(&self) -> Option<&SearchExplain> {
91        self.planner.as_ref()
92    }
93
94    pub const fn parameter_origins(&self) -> Option<&MessageSearchOrigins> {
95        self.origins.as_ref()
96    }
97
98    pub fn ordered_digest(&self) -> Option<&str> {
99        self.ordered_digest.as_deref()
100    }
101
102    pub const fn page_document(&self) -> MessageSearchPageDocument {
103        MessageSearchPageDocument::from_page(self.page)
104    }
105
106    pub fn receipt_document<'a>(
107        &'a self,
108        request: &ResolvedMessageSearchRequest,
109    ) -> Option<MessageSearchReceiptDocument<'a>> {
110        MessageSearchReceiptDocument::from_parts(
111            request.receipt_level(),
112            self.search_explanation(),
113            self.parameter_origins(),
114            self.ordered_digest(),
115        )
116    }
117}
118
119enum MessageSearchBatchEvent {
120    Batch(MessageSearchBatch),
121    Complete(Box<MessageSearchCompletion>),
122    Failed(String),
123}
124
125enum MessageSearchBatchState {
126    Open,
127    Exhausted(Box<MessageSearchCompletion>),
128    Closed,
129    Failed(String),
130}
131
132#[derive(Debug)]
133struct MessageSearchCancelled;
134
135impl std::fmt::Display for MessageSearchCancelled {
136    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
137        formatter.write_str("message-search batches were cancelled")
138    }
139}
140
141impl std::error::Error for MessageSearchCancelled {}
142
143pub(crate) fn ensure_message_search_active(cancellation: Option<&AtomicBool>) -> Result<()> {
144    if cancellation.is_some_and(|flag| flag.load(Ordering::Acquire)) {
145        return Err(MessageSearchCancelled.into());
146    }
147    Ok(())
148}
149
150fn is_expected_message_search_cancellation(error: &anyhow::Error) -> bool {
151    !error.is::<ReadSnapshotCleanupError>()
152        && (error.is::<MessageSearchCancelled>()
153            || error.chain().any(|cause| {
154                matches!(
155                    cause.downcast_ref::<rusqlite::Error>(),
156                    Some(rusqlite::Error::SqliteFailure(inner, _))
157                        if inner.code == rusqlite::ErrorCode::OperationInterrupted
158                )
159            }))
160}
161
162/// Bounded-retention exhaustive message-search batches.
163///
164/// The ordinary [`MessageService::search`] API remains the easiest path for callers that want one
165/// materialized list. This owner runs the one-snapshot visitor on a dedicated read-only connection
166/// and transfers at most one enriched batch through a rendezvous channel. `close` interrupts an
167/// active SQLite call, releases the snapshot without draining unread rows, and joins the producer.
168/// Internally, the caller's current batch and the producer's next handoff can coexist; callers may
169/// also retain any number of already returned owned batches, so their memory is caller-controlled.
170///
171/// # Complexity
172///
173/// Let `C` be SQLite candidate work, `H` returned hits, `B` `batch_rows`, `T_B` selected text and
174/// included metadata bytes in one batch, `W_B` context-window rows for one batch, `U` distinct
175/// sessions whose metadata has been emitted, and `I_U` their total ID bytes. SQLite owns filtering,
176/// indexed candidate selection, ordering, offsets, context joins, and session-metadata joins. The
177/// producer steps one result statement, so retrieval is `O(C + H)` rather than repeated offset
178/// scans. With `P = ceil(H / B)`, context uses at most `P` set queries when requested; each
179/// requested metadata form uses at most `min(P, U)` set queries, `O(H log U)` deduplication work,
180/// and emits each session's metadata bytes once. Internal live memory is
181/// `O(B + T_B + W_B + U + I_U)` plus one worker thread and one read-only SQLite connection;
182/// memory for batches retained by the caller is additional. Closing does not scan unread hits,
183/// though its latency includes the active SQLite progress-handler interval or one selected
184/// field's current Rust enrichment step.
185pub struct MessageSearchBatches {
186    request: ResolvedMessageSearchRequest,
187    runtime_diagnostics: Option<MessageSearchRuntimeDiagnostics>,
188    start: Option<mpsc::SyncSender<()>>,
189    events: Option<mpsc::Receiver<MessageSearchBatchEvent>>,
190    cancellation: Arc<AtomicBool>,
191    interrupt: Option<rusqlite::InterruptHandle>,
192    worker: Option<std::thread::JoinHandle<std::result::Result<(), String>>>,
193    state: MessageSearchBatchState,
194}
195
196impl MessageSearchBatches {
197    /// The request after purpose, operation, surface, and typed defaults were resolved.
198    pub const fn request(&self) -> &ResolvedMessageSearchRequest {
199        &self.request
200    }
201
202    /// Request-wide runtime diagnostics, computed once when explicitly requested.
203    ///
204    /// Per-batch [`MessageSearchBatch::included`] contains only session-scoped metadata deltas, so
205    /// adapters can emit these invariant diagnostics once in their initial metadata record.
206    pub const fn runtime_diagnostics(&self) -> Option<&MessageSearchRuntimeDiagnostics> {
207        self.runtime_diagnostics.as_ref()
208    }
209
210    pub fn next_batch(&mut self) -> Result<Option<MessageSearchBatch>> {
211        match &self.state {
212            MessageSearchBatchState::Exhausted(_) => return Ok(None),
213            MessageSearchBatchState::Closed => {
214                bail!(
215                    "message-search batches are closed; start a new batch request to restart the search"
216                )
217            }
218            MessageSearchBatchState::Failed(error) => {
219                bail!("message-search batches previously failed: {error}")
220            }
221            MessageSearchBatchState::Open => {}
222        }
223        // Construction resolves and snapshots the request but does not start result SQL. The first
224        // read supplies demand, so an immediate close releases the snapshot without racing an FTS
225        // virtual-table constructor or doing work the caller never requested.
226        if let Some(start) = self.start.take() {
227            let _ = start.send(());
228        }
229        let event = match self
230            .events
231            .as_ref()
232            .ok_or_else(|| anyhow!("message-search batches lost their producer channel"))?
233            .recv()
234        {
235            Ok(event) => event,
236            Err(_) => {
237                self.events.take();
238                let error = self.join_worker()?.unwrap_or_else(|| {
239                    "message-search batch producer exited without a terminal state".to_string()
240                });
241                self.state = MessageSearchBatchState::Failed(error.clone());
242                bail!("message-search batches failed: {error}");
243            }
244        };
245        match event {
246            MessageSearchBatchEvent::Batch(batch) => Ok(Some(batch)),
247            MessageSearchBatchEvent::Complete(completion) => {
248                self.state = MessageSearchBatchState::Exhausted(completion);
249                self.events.take();
250                if let Some(error) = self.join_worker()? {
251                    self.state = MessageSearchBatchState::Failed(error.clone());
252                    bail!("message-search batches failed after terminal metadata: {error}");
253                }
254                Ok(None)
255            }
256            MessageSearchBatchEvent::Failed(error) => {
257                self.state = MessageSearchBatchState::Failed(error.clone());
258                self.events.take();
259                if let Some(worker_error) = self.join_worker()? {
260                    debug_assert_eq!(worker_error, error);
261                }
262                bail!("message-search batches failed: {error}")
263            }
264        }
265    }
266
267    /// Return terminal metadata without reading unread batches.
268    pub fn completion(&self) -> Result<&MessageSearchCompletion> {
269        match &self.state {
270            MessageSearchBatchState::Exhausted(completion) => Ok(completion),
271            MessageSearchBatchState::Open => bail!(
272                "message-search batches have unread results; call next_batch() until it returns None, or call close() to stop without terminal metadata"
273            ),
274            MessageSearchBatchState::Closed => bail!(
275                "message-search batches were closed before natural exhaustion; no terminal metadata exists"
276            ),
277            MessageSearchBatchState::Failed(error) => {
278                bail!("message-search batches failed before completion: {error}")
279            }
280        }
281    }
282
283    /// Stop without draining unread results. Repeated calls are safe.
284    pub fn close(&mut self) -> Result<()> {
285        if matches!(
286            self.state,
287            MessageSearchBatchState::Closed | MessageSearchBatchState::Exhausted(_)
288        ) {
289            return match self.join_worker()? {
290                Some(error) => {
291                    bail!("message-search batch producer failed during cleanup: {error}")
292                }
293                None => Ok(()),
294            };
295        }
296        self.cancellation.store(true, Ordering::Release);
297        self.start.take();
298        if let Some(interrupt) = self.interrupt.take() {
299            interrupt.interrupt();
300        }
301        self.events.take();
302        let worker_error = self.join_worker()?;
303        self.state = MessageSearchBatchState::Closed;
304        match worker_error {
305            Some(error) => {
306                bail!("message-search batch producer failed while closing: {error}")
307            }
308            None => Ok(()),
309        }
310    }
311
312    fn join_worker(&mut self) -> Result<Option<String>> {
313        self.interrupt.take();
314        if let Some(worker) = self.worker.take() {
315            return match worker.join() {
316                Ok(Ok(())) => Ok(None),
317                Ok(Err(error)) => Ok(Some(error)),
318                Err(_) => Err(anyhow!(
319                    "message-search batch producer panicked during cleanup"
320                )),
321            };
322        }
323        Ok(None)
324    }
325
326    pub(crate) fn spawn(
327        config: Config,
328        access: EffectiveAccessScope,
329        surface: SearchSurface,
330        request: MessageSearchRequest,
331        batch_rows: NonZeroUsize,
332    ) -> Result<Self> {
333        Self::spawn_with_before_traversal(config, access, surface, request, batch_rows, || {})
334    }
335
336    pub(crate) fn spawn_with_before_traversal(
337        config: Config,
338        access: EffectiveAccessScope,
339        surface: SearchSurface,
340        request: MessageSearchRequest,
341        batch_rows: NonZeroUsize,
342        before_traversal: impl FnOnce() + Send + 'static,
343    ) -> Result<Self> {
344        // Request-only eligibility is checked before allocating channels or starting a thread, so
345        // callers receive the typed MessageSearchError and invalid input owns no resources. The
346        // worker repeats the rule on the resolved plan as a drift assertion after config policy.
347        MessageService::validate_batch_request(surface, &request)?;
348        let (ready_sender, ready_receiver) = mpsc::sync_channel::<
349            std::result::Result<
350                (
351                    ResolvedMessageSearchRequest,
352                    Option<MessageSearchRuntimeDiagnostics>,
353                    rusqlite::InterruptHandle,
354                ),
355                String,
356            >,
357        >(1);
358        let (start_sender, start_receiver) = mpsc::sync_channel(0);
359        // A rendezvous channel provides backpressure without retaining a queued result batch.
360        // The producer may hold one batch while the consumer owns and may retain prior batches.
361        // This bounds internal queued work, not memory deliberately retained by the caller.
362        let (event_sender, event_receiver) = mpsc::sync_channel(0);
363        let cancellation = Arc::new(AtomicBool::new(false));
364        let worker_cancellation = Arc::clone(&cancellation);
365        let worker = std::thread::Builder::new()
366            .name("aise-message-search-batches".into())
367            .spawn(move || {
368                let mut ready_sent = false;
369                let result = (|| -> Result<()> {
370                    let worker_threads = NonZeroUsize::new(config.resolve_threads())
371                        .expect("Config::resolve_threads always returns at least one");
372                    let mut db = Db::open_existing_read_only_with_threads(
373                        &config.db_path(),
374                        config.index.busy_timeout_ms,
375                        worker_threads,
376                    )?;
377                    db.set_access_scope(access);
378                    anyhow::ensure!(
379                        db.schema_version()? == crate::db::SCHEMA_VERSION,
380                        "exhaustive batched message search requires database schema {}; run `aise reindex --full`, then retry",
381                        crate::db::SCHEMA_VERSION
382                    );
383                    db.interrupt_while(Arc::clone(&worker_cancellation))?;
384                    db.with_read_snapshot(|| {
385                        let service = MessageService::new(&config, &db, surface);
386                        let plan = service.plan(request)?;
387                        MessageService::validate_batch_plan(&plan)?;
388                        let resolved_request = ResolvedMessageSearchRequest::from_plan(&plan)?;
389                        let runtime_diagnostics =
390                            service.runtime_diagnostics_for_plan(&plan)?;
391                        let interrupt = db.interrupt_handle();
392                        ready_sender
393                            .send(Ok((resolved_request, runtime_diagnostics, interrupt)))
394                            .map_err(|_| {
395                                anyhow!("message-search batches were dropped during startup")
396                            })?;
397                        ready_sent = true;
398                        before_traversal();
399                        start_receiver
400                            .recv()
401                            .map_err(|_| anyhow::Error::new(MessageSearchCancelled))?;
402                        ensure_message_search_active(Some(&worker_cancellation))?;
403                        let outcome = service.visit_search_plan_batches(
404                            plan,
405                            batch_rows,
406                            Some(&worker_cancellation),
407                            |batch| {
408                                ensure_message_search_active(Some(&worker_cancellation))?;
409                                match event_sender.send(MessageSearchBatchEvent::Batch(batch)) {
410                                    Ok(()) => Ok(MessageSearchBatchControl::Continue),
411                                    Err(_) => Ok(MessageSearchBatchControl::Stop),
412                                }
413                            },
414                        )?;
415                        if outcome.exhausted && !worker_cancellation.load(Ordering::Acquire) {
416                            let completion = MessageSearchCompletion {
417                                page: outcome
418                                    .page
419                                    .expect("naturally exhausted traversal has page metadata"),
420                                planner: outcome.planner,
421                                origins: outcome.origins,
422                                ordered_digest: outcome.ordered_digest,
423                            };
424                            let _ = event_sender
425                                .send(MessageSearchBatchEvent::Complete(Box::new(completion)));
426                        }
427                        Ok(())
428                    })
429                })();
430                match result {
431                    Ok(()) => Ok(()),
432                    Err(error) => {
433                        let rendered = format!("{error:#}");
434                        if !ready_sent {
435                            let _ = ready_sender.send(Err(rendered));
436                            Ok(())
437                        } else if worker_cancellation.load(Ordering::Acquire)
438                            && is_expected_message_search_cancellation(&error)
439                        {
440                            Ok(())
441                        } else {
442                            let _ = event_sender
443                                .send(MessageSearchBatchEvent::Failed(rendered.clone()));
444                            Err(rendered)
445                        }
446                    }
447                }
448            })
449            .map_err(|error| {
450                anyhow!("failed to start message-search batch producer: {error}")
451            })?;
452        let (request, runtime_diagnostics, interrupt) = match ready_receiver.recv() {
453            Ok(Ok(ready)) => ready,
454            Ok(Err(error)) => {
455                match worker.join() {
456                    Ok(Ok(())) => {}
457                    Ok(Err(worker_error)) => {
458                        bail!(
459                            "failed to open message-search batches: {error}; producer cleanup also failed: {worker_error}"
460                        )
461                    }
462                    Err(_) => {
463                        bail!(
464                            "failed to open message-search batches: {error}; producer panicked during startup cleanup"
465                        )
466                    }
467                }
468                bail!("failed to open message-search batches: {error}");
469            }
470            Err(_) => match worker.join() {
471                Ok(Ok(())) => {
472                    bail!("message-search batch producer exited during startup")
473                }
474                Ok(Err(error)) => {
475                    bail!("message-search batch producer failed during startup: {error}")
476                }
477                Err(_) => {
478                    bail!("message-search batch producer panicked during startup")
479                }
480            },
481        };
482        Ok(Self {
483            request,
484            runtime_diagnostics,
485            start: Some(start_sender),
486            events: Some(event_receiver),
487            cancellation,
488            interrupt: Some(interrupt),
489            worker: Some(worker),
490            state: MessageSearchBatchState::Open,
491        })
492    }
493}
494
495impl Drop for MessageSearchBatches {
496    fn drop(&mut self) {
497        let _ = self.close();
498    }
499}
500
501#[cfg(test)]
502mod tests {
503    use super::*;
504
505    #[test]
506    fn snapshot_cleanup_failure_is_never_treated_as_successful_cancellation() {
507        let error = anyhow::Error::new(ReadSnapshotCleanupError::new(
508            Some(anyhow::Error::new(MessageSearchCancelled)),
509            anyhow!("injected rollback failure"),
510        ));
511
512        assert!(!is_expected_message_search_cancellation(&error));
513        assert!(error
514            .to_string()
515            .contains("read snapshot also failed to release"));
516    }
517}
ai_session_search/message_search_batches.rs

ai_session_search/
message_search_batches.rs
