aft/commands/semantic_search/
extensions.rs

1use std::path::{Path, PathBuf};
2use std::sync::{Arc, RwLockReadGuard};
3
4use serde::Serialize;
5
6use super::plan_table::{SearchLaneKind, SearchShape};
7use super::{LaneExecution, LaneInput, SearchLane};
8use crate::context::SemanticIndexStatus;
9use crate::parser::SymbolCache;
10use crate::query_shape::{classify, looks_like_regex, QueryKind};
11use crate::search_index::{IndexStatus, SearchIndexSnapshot};
12use crate::semantic_index::SemanticIndex;
13
14/// Original request text retained before lane-specific normalization.
15#[derive(Debug, Clone, PartialEq, Eq)]
16pub struct RawQuery(String);
17
18impl RawQuery {
19    pub fn new(query: impl Into<String>) -> Self {
20        Self(query.into())
21    }
22
23    pub fn original_query(&self) -> &str {
24        &self.0
25    }
26}
27
28/// Byte range of a delimiter's verbatim interior in [`RawQuery`].
29#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
30pub struct Span {
31    pub start: usize,
32    pub end: usize,
33}
34
35impl Span {
36    pub fn extract<'a>(&self, raw_query: &'a str) -> Option<&'a str> {
37        raw_query.get(self.start..self.end)
38    }
39}
40
41/// The five classification facts passed unchanged to lane planning: the first
42/// embedded span, qualifying exact-token count, path presence, log-marker
43/// presence, and identifier-shaped-token presence.
44#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
45pub struct QueryFacts {
46    pub embedded_span: Option<Span>,
47    pub exact_input_tokens: usize,
48    pub has_path_token: bool,
49    pub has_timestamp_or_pid: bool,
50    pub has_identifier_token: bool,
51}
52
53#[derive(Debug, Clone, Copy, PartialEq, Eq)]
54pub struct Token<'a> {
55    pub index: usize,
56    pub text: &'a str,
57}
58
59#[derive(Debug, Clone, PartialEq, Eq)]
60pub struct TokenVariant {
61    pub token_index: usize,
62    pub text: String,
63}
64
65/// Result of waiting for the first search index to become ready, with a time bound.
66#[derive(Debug, Clone, Copy, PartialEq, Eq)]
67pub enum ReadinessWait {
68    Completed,
69    Cancelled,
70}
71
72/// Symbol-cache lifecycle state for the selected root.
73#[derive(Debug, Clone, PartialEq, Eq)]
74pub enum SymbolIndexStatus {
75    Ready,
76    Building,
77    Disabled,
78    Failed(String),
79}
80
81/// A semantic index retained without copying its embedding storage.
82#[derive(Clone)]
83pub struct SemanticSnapshot<'a>(SemanticSnapshotStorage<'a>);
84
85#[derive(Clone)]
86enum SemanticSnapshotStorage<'a> {
87    Guard(Arc<RwLockReadGuard<'a, Option<SemanticIndex>>>),
88    Borrowed(Arc<SemanticIndex>),
89}
90
91impl std::fmt::Debug for SemanticSnapshot<'_> {
92    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
93        formatter
94            .debug_struct("SemanticSnapshot")
95            .field("present", &self.index().is_some())
96            .finish()
97    }
98}
99
100impl<'a> SemanticSnapshot<'a> {
101    pub fn from_guard(guard: RwLockReadGuard<'a, Option<SemanticIndex>>) -> Self {
102        Self(SemanticSnapshotStorage::Guard(Arc::new(guard)))
103    }
104
105    pub fn from_borrowed(index: Arc<SemanticIndex>) -> Self {
106        Self(SemanticSnapshotStorage::Borrowed(index))
107    }
108
109    pub fn index(&self) -> Option<&SemanticIndex> {
110        match &self.0 {
111            SemanticSnapshotStorage::Guard(guard) => guard.as_ref().as_ref(),
112            SemanticSnapshotStorage::Borrowed(index) => Some(index),
113        }
114    }
115}
116
117/// One semantic-index observation and the resource that made it queryable.
118#[derive(Debug, Clone)]
119pub struct SemanticReadiness<'a> {
120    pub status: SemanticIndexStatus,
121    pub snapshot: Option<SemanticSnapshot<'a>>,
122    pub evicted: bool,
123    pub lock_contended: bool,
124}
125
126/// One trigram-index observation and the resource that made it queryable.
127#[derive(Debug, Clone)]
128pub struct TrigramReadiness {
129    pub status: IndexStatus,
130    pub snapshot: Option<Arc<SearchIndexSnapshot>>,
131    pub evicted: bool,
132    pub lock_contended: bool,
133}
134
135/// One symbol-cache observation and the resource that made it queryable.
136#[derive(Clone)]
137pub struct SymbolReadiness {
138    pub status: SymbolIndexStatus,
139    pub snapshot: Option<Arc<SymbolCache>>,
140    pub evicted: bool,
141    pub lock_contended: bool,
142}
143
144impl std::fmt::Debug for SymbolReadiness {
145    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
146        formatter
147            .debug_struct("SymbolReadiness")
148            .field("status", &self.status)
149            .field("has_snapshot", &self.snapshot.is_some())
150            .field("evicted", &self.evicted)
151            .field("lock_contended", &self.lock_contended)
152            .finish()
153    }
154}
155
156/// The three real runtime states captured in one admission sample.
157#[derive(Debug, Clone)]
158pub struct ReadinessObservation<'a> {
159    pub semantic: SemanticReadiness<'a>,
160    pub trigram: TrigramReadiness,
161    pub symbol: SymbolReadiness,
162}
163
164/// Supplies root-scoped runtime state to the extension seam.
165///
166/// Implementations return retained handles so a selected plan keeps using the
167/// exact resources it admitted even if the context evicts its live pointers.
168pub trait ReadinessSource {
169    fn sample<'a>(&'a self) -> ReadinessObservation<'a>;
170    fn bounded_first_search_wait(&self) -> ReadinessWait;
171}
172
173/// Input accepted by [`Root::new`]. The fixed form keeps source compatibility
174/// for older callers while public search samples from a runtime source.
175#[doc(hidden)]
176pub enum RootReadiness<'a> {
177    Source(&'a dyn ReadinessSource),
178    Fixed(Readiness<'a>),
179}
180
181#[doc(hidden)]
182pub trait IntoRootReadiness<'a> {
183    fn into_root_readiness(self) -> RootReadiness<'a>;
184}
185
186impl<'a> IntoRootReadiness<'a> for &'a dyn ReadinessSource {
187    fn into_root_readiness(self) -> RootReadiness<'a> {
188        RootReadiness::Source(self)
189    }
190}
191
192impl<'a> IntoRootReadiness<'a> for Readiness<'a> {
193    fn into_root_readiness(self) -> RootReadiness<'a> {
194        RootReadiness::Fixed(self)
195    }
196}
197
198/// Root-scoped resource source sampled before a lane plan is built.
199pub struct Root<'a> {
200    path: PathBuf,
201    readiness: RootReadiness<'a>,
202}
203
204impl std::fmt::Debug for Root<'_> {
205    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
206        formatter
207            .debug_struct("Root")
208            .field("path", &self.path)
209            .finish()
210    }
211}
212
213impl<'a> Root<'a> {
214    pub fn new(path: impl Into<PathBuf>, readiness: impl IntoRootReadiness<'a>) -> Self {
215        Self {
216            path: path.into(),
217            readiness: readiness.into_root_readiness(),
218        }
219    }
220
221    pub fn path(&self) -> &Path {
222        &self.path
223    }
224
225    pub fn source(&self) -> Option<&'a dyn ReadinessSource> {
226        match &self.readiness {
227            RootReadiness::Source(source) => Some(*source),
228            RootReadiness::Fixed(_) => None,
229        }
230    }
231
232    pub fn fixed_readiness(&self) -> Option<&Readiness<'a>> {
233        match &self.readiness {
234            RootReadiness::Source(_) => None,
235            RootReadiness::Fixed(readiness) => Some(readiness),
236        }
237    }
238}
239
240#[derive(Clone, Default)]
241pub struct RetainedReadinessSnapshots<'a> {
242    semantic: Option<SemanticSnapshot<'a>>,
243    trigram: Option<Arc<SearchIndexSnapshot>>,
244    symbol: Option<Arc<SymbolCache>>,
245}
246
247impl std::fmt::Debug for RetainedReadinessSnapshots<'_> {
248    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
249        formatter
250            .debug_struct("RetainedReadinessSnapshots")
251            .field("semantic", &self.semantic.is_some())
252            .field("trigram", &self.trigram.is_some())
253            .field("symbol", &self.symbol.is_some())
254            .finish()
255    }
256}
257
258impl<'a> RetainedReadinessSnapshots<'a> {
259    pub fn new(
260        semantic: Option<SemanticSnapshot<'a>>,
261        trigram: Option<Arc<SearchIndexSnapshot>>,
262        symbol: Option<Arc<SymbolCache>>,
263    ) -> Self {
264        Self {
265            semantic,
266            trigram,
267            symbol,
268        }
269    }
270
271    pub fn semantic(&self) -> Option<&SemanticIndex> {
272        self.semantic.as_ref().and_then(SemanticSnapshot::index)
273    }
274
275    pub fn trigram(&self) -> Option<&SearchIndexSnapshot> {
276        self.trigram.as_deref()
277    }
278
279    pub fn symbol(&self) -> Option<&SymbolCache> {
280        self.symbol.as_deref()
281    }
282}
283
284#[derive(Debug, Clone, Serialize)]
285pub struct Readiness<'a> {
286    pub symbol_index: bool,
287    pub lexical_index: bool,
288    pub semantic_index: bool,
289    #[serde(skip_serializing_if = "Vec::is_empty")]
290    pub reasons: Vec<String>,
291    #[serde(skip)]
292    retained: RetainedReadinessSnapshots<'a>,
293    #[serde(skip)]
294    cancelled: bool,
295}
296
297impl PartialEq for Readiness<'_> {
298    fn eq(&self, other: &Self) -> bool {
299        self.symbol_index == other.symbol_index
300            && self.lexical_index == other.lexical_index
301            && self.semantic_index == other.semantic_index
302            && self.reasons == other.reasons
303            && self.cancelled == other.cancelled
304    }
305}
306
307impl Eq for Readiness<'_> {}
308
309impl<'a> Readiness<'a> {
310    pub fn new(symbol_index: bool, lexical_index: bool, semantic_index: bool) -> Self {
311        let mut reasons = Vec::new();
312        if !symbol_index {
313            reasons.push("symbol_index_unavailable".to_string());
314        }
315        if !lexical_index {
316            reasons.push("lexical_index_unavailable".to_string());
317        }
318        if !semantic_index {
319            reasons.push("semantic_index_unavailable".to_string());
320        }
321        Self::observed(
322            symbol_index,
323            lexical_index,
324            semantic_index,
325            reasons,
326            RetainedReadinessSnapshots::default(),
327            false,
328        )
329    }
330
331    pub fn observed(
332        symbol_index: bool,
333        lexical_index: bool,
334        semantic_index: bool,
335        reasons: Vec<String>,
336        retained: RetainedReadinessSnapshots<'a>,
337        cancelled: bool,
338    ) -> Self {
339        Self {
340            symbol_index,
341            lexical_index,
342            semantic_index,
343            reasons,
344            retained,
345            cancelled,
346        }
347    }
348
349    pub fn retained(&self) -> &RetainedReadinessSnapshots<'a> {
350        &self.retained
351    }
352
353    pub fn cancelled(&self) -> bool {
354        self.cancelled
355    }
356}
357
358#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
359#[serde(rename_all = "snake_case")]
360pub enum ExactMode {
361    Ready,
362    Fallback,
363    #[serde(rename = "n/a")]
364    NotApplicable,
365}
366
367#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
368pub struct LanePlan<'a> {
369    pub shape: SearchShape,
370    pub query_facts: QueryFacts,
371    #[serde(skip_serializing_if = "Option::is_none")]
372    pub exact_input: Option<String>,
373    pub exact_mode: ExactMode,
374    #[serde(rename = "lanes_run")]
375    pub selected_lanes: Vec<SearchLaneKind>,
376    pub executed_callbacks: Vec<SearchLaneKind>,
377    pub readiness: Readiness<'a>,
378    pub variants: Vec<String>,
379}
380
381impl LanePlan<'_> {
382    pub fn contains(&self, lane: SearchLaneKind) -> bool {
383        self.selected_lanes.contains(&lane)
384    }
385}
386
387/// A-side extension seam. Later ranking campaigns can override one hook at a
388/// time while the base engine remains independently buildable.
389pub trait SearchExtensions: Send + Sync {
390    fn classify(&self, raw_query: &RawQuery) -> (SearchShape, QueryFacts) {
391        classify_raw_query(raw_query)
392    }
393
394    fn variants(&self, _token: Token<'_>) -> Vec<TokenVariant> {
395        Vec::new()
396    }
397
398    fn sample_readiness<'a>(&self, root: &Root<'a>) -> Readiness<'a> {
399        crate::search_b2::readiness::sample(root)
400    }
401
402    fn plan<'a>(
403        &self,
404        shape: &SearchShape,
405        facts: &QueryFacts,
406        readiness: &Readiness<'a>,
407    ) -> LanePlan<'a> {
408        default_lane_plan(*shape, facts, readiness.clone())
409    }
410
411    fn execute_lane(&self, lane: &dyn SearchLane, input: &LaneInput<'_>) -> LaneExecution {
412        lane.execute(input)
413    }
414}
415
416#[derive(Debug, Default, Clone, Copy)]
417pub struct DefaultSearchExtensions;
418
419impl SearchExtensions for DefaultSearchExtensions {}
420
421/// Preserves the pre-B2 classifier while emitting the new explicit facts tuple.
422pub fn classify_raw_query(raw_query: &RawQuery) -> (SearchShape, QueryFacts) {
423    let query = raw_query.original_query().trim();
424    let shape = if looks_like_regex(query) {
425        SearchShape::Regex
426    } else if is_quoted(query) {
427        SearchShape::CodeLiteral
428    } else {
429        let query_shape = classify(query);
430        if query_shape.kind == QueryKind::ErrorCode || looks_like_log_excerpt(query) {
431            SearchShape::LogExcerpt
432        } else if query_shape.kind == QueryKind::Path {
433            SearchShape::Path
434        } else if query_shape.kind == QueryKind::Identifier {
435            SearchShape::Identifier
436        } else if query.split_whitespace().count() <= 2 {
437            SearchShape::Short
438        } else {
439            SearchShape::NaturalLanguage
440        }
441    };
442    let (_, facts) = crate::search_b2::router::classify(raw_query);
443    (shape, facts)
444}
445
446fn default_lane_plan<'a>(
447    shape: SearchShape,
448    facts: &QueryFacts,
449    readiness: Readiness<'a>,
450) -> LanePlan<'a> {
451    let mut selected_lanes = match shape {
452        SearchShape::Identifier => vec![
453            SearchLaneKind::Symbol,
454            SearchLaneKind::Exact,
455            SearchLaneKind::Lexical,
456            SearchLaneKind::Variants,
457            SearchLaneKind::Semantic,
458        ],
459        SearchShape::CodeLiteral => vec![SearchLaneKind::Exact, SearchLaneKind::Lexical],
460        SearchShape::Short => vec![
461            SearchLaneKind::Exact,
462            SearchLaneKind::Lexical,
463            SearchLaneKind::Semantic,
464        ],
465        SearchShape::NaturalLanguage => vec![
466            SearchLaneKind::Exact,
467            SearchLaneKind::Lexical,
468            SearchLaneKind::Variants,
469            SearchLaneKind::Semantic,
470        ],
471        SearchShape::LogExcerpt => vec![
472            SearchLaneKind::Anchored,
473            SearchLaneKind::Exact,
474            SearchLaneKind::Lexical,
475        ],
476        SearchShape::Path => vec![
477            SearchLaneKind::PathLookup,
478            SearchLaneKind::Exact,
479            SearchLaneKind::Lexical,
480            SearchLaneKind::FallbackWalk,
481        ],
482        SearchShape::Regex => vec![SearchLaneKind::FallbackWalk],
483    };
484
485    if !readiness.symbol_index {
486        selected_lanes.retain(|lane| *lane != SearchLaneKind::Symbol);
487    }
488    if !readiness.lexical_index {
489        selected_lanes.retain(|lane| {
490            !matches!(
491                lane,
492                SearchLaneKind::Exact | SearchLaneKind::Anchored | SearchLaneKind::Lexical
493            )
494        });
495    }
496    if !readiness.semantic_index {
497        selected_lanes.retain(|lane| *lane != SearchLaneKind::Semantic);
498    }
499    if !readiness.reasons.is_empty() {
500        selected_lanes.push(SearchLaneKind::ReadinessDisclosure);
501    }
502
503    let exact_mode = if !selected_lanes.contains(&SearchLaneKind::Exact) {
504        ExactMode::NotApplicable
505    } else if readiness.lexical_index && facts.exact_input_tokens > 0 {
506        ExactMode::Ready
507    } else {
508        ExactMode::Fallback
509    };
510    LanePlan {
511        shape,
512        query_facts: facts.clone(),
513        exact_input: None,
514        exact_mode,
515        executed_callbacks: selected_lanes.clone(),
516        selected_lanes,
517        readiness,
518        variants: Vec::new(),
519    }
520}
521
522fn is_quoted(query: &str) -> bool {
523    query.len() >= 2
524        && ((query.starts_with('"') && query.ends_with('"'))
525            || (query.starts_with('\'') && query.ends_with('\'')))
526}
527
528fn looks_like_log_excerpt(query: &str) -> bool {
529    let upper = query.to_ascii_uppercase();
530    [" ERROR ", " WARN ", " INFO ", " DEBUG ", " TRACE "]
531        .iter()
532        .any(|marker| upper.contains(marker))
533        || query.contains("::") && query.contains('[') && query.contains(']')
534}
aft/commands/semantic_search/extensions.rs

aft/commands/semantic_search/
extensions.rs
