Skip to main content

kaizen/shell/
search.rs

1// SPDX-License-Identifier: AGPL-3.0-or-later
2//! `kaizen sessions search` and `kaizen search reindex`.
3
4use crate::core::config::{self, try_team_salt};
5use crate::core::event::Event;
6use crate::search::extract::{redacted_event_text, snippet};
7use crate::search::{SearchHit, SearchQuery};
8use crate::shell::cli::{open_workspace_read_store, workspace_path};
9use crate::shell::fmt::fmt_ts;
10use crate::store::Store;
11use crate::store::event_index::{is_valid_slug, paths_from_event_payload, skills_from_event_json};
12use anyhow::{Context, Result};
13use std::path::Path;
14
15pub fn cmd_sessions_search(
16    workspace: Option<&Path>,
17    query: &str,
18    since: Option<&str>,
19    agent: Option<&str>,
20    kind: Option<&str>,
21    limit: usize,
22) -> Result<()> {
23    print!(
24        "{}",
25        sessions_search_text(workspace, query, since, agent, kind, limit)?
26    );
27    Ok(())
28}
29
30pub fn sessions_search_text(
31    workspace: Option<&Path>,
32    query: &str,
33    since: Option<&str>,
34    agent: Option<&str>,
35    kind: Option<&str>,
36    limit: usize,
37) -> Result<String> {
38    let (hits, fallback) = sessions_search_hits(workspace, query, since, agent, kind, limit)?;
39    render_hits(&hits, fallback)
40}
41
42pub fn sessions_search_hits(
43    workspace: Option<&Path>,
44    query: &str,
45    since: Option<&str>,
46    agent: Option<&str>,
47    kind: Option<&str>,
48    limit: usize,
49) -> Result<(Vec<SearchHit>, bool)> {
50    let ws = workspace_path(workspace)?;
51    let store = open_workspace_read_store(&ws, false)?;
52    let cfg = config::load(&ws)?;
53    let salt = try_team_salt(&cfg.sync).unwrap_or([0; 32]);
54    let opts = SearchQuery {
55        query: query.to_string(),
56        since_ms: parse_since(since)?,
57        agent: agent.map(str::to_string),
58        kind: kind.map(str::to_string),
59        limit,
60    };
61    let data_dir = crate::core::paths::project_data_dir(&ws)?;
62    if let Some(hits) = exact_tool_hits(&store, &ws, &opts, &salt)? {
63        return Ok((hits, false));
64    }
65    match crate::search::search(&data_dir, &opts, &ws, &salt, |s, q| store.get_event(s, q)) {
66        Ok(hits) => Ok((hits, false)),
67        Err(e) => anyhow::bail!("search index unavailable: {e}; run `kaizen search reindex`"),
68    }
69}
70
71fn exact_tool_hits(
72    store: &Store,
73    workspace: &Path,
74    opts: &SearchQuery,
75    salt: &[u8; 32],
76) -> Result<Option<Vec<SearchHit>>> {
77    let Some(tool) = exact_tool_query(opts) else {
78        return Ok(None);
79    };
80    let ws = workspace.to_string_lossy();
81    let rows =
82        store.search_tool_events(&ws, tool, opts.since_ms, opts.agent.as_deref(), opts.limit)?;
83    let hits = rows
84        .into_iter()
85        .map(|(agent, event)| tool_hit(agent, event, workspace, salt, &opts.query))
86        .collect::<Vec<_>>();
87    Ok((!hits.is_empty()).then_some(hits))
88}
89
90fn exact_tool_query(opts: &SearchQuery) -> Option<&str> {
91    if opts.limit == 0 || opts.kind.is_some() || !is_valid_slug(&opts.query) {
92        return None;
93    }
94    Some(opts.query.as_str())
95}
96
97fn tool_hit(
98    agent: String,
99    event: Event,
100    workspace: &Path,
101    salt: &[u8; 32],
102    query: &str,
103) -> SearchHit {
104    let text = redacted_event_text(&event, workspace, salt);
105    SearchHit {
106        session_id: event.session_id.clone(),
107        seq: event.seq,
108        ts_ms: event.ts_ms,
109        agent,
110        kind: crate::search::kind_label(&event.kind)
111            .unwrap_or("unknown")
112            .to_string(),
113        score: 1.0,
114        snippet: snippet(&text, query),
115        paths: paths_from_event_payload(&event.payload),
116        skills: skills_from_event_json(&event.payload),
117        tokens_total: crate::search::tokens_total(&event),
118    }
119}
120
121pub fn cmd_search_reindex(workspace: Option<&Path>) -> Result<()> {
122    let ws = workspace_path(workspace)?;
123    let store = Store::open(&crate::core::workspace::db_path(&ws)?)?;
124    let cfg = config::load(&ws)?;
125    let ws_str = ws.to_string_lossy().to_string();
126    let sessions = store.list_sessions(&ws_str)?;
127    let events = store.workspace_events(&ws_str)?;
128    let data_dir = crate::core::paths::project_data_dir(&ws)?;
129    let stats = crate::search::reindex_workspace(&data_dir, &ws, &sessions, events, &cfg)
130        .context("reindex search")?;
131    println!(
132        "search reindex: {} events seen, {} docs indexed",
133        stats.events_seen, stats.docs_indexed
134    );
135    Ok(())
136}
137
138fn render_hits(hits: &[SearchHit], fallback: bool) -> Result<String> {
139    use std::fmt::Write;
140    let mut out = String::new();
141    if fallback {
142        writeln!(
143            out,
144            "warning: search index unavailable; falling back to event scan"
145        )?;
146    }
147    writeln!(out, "{:<40} {:<19} {:>7} SNIPPET", "SESSION", "TS", "SCORE")?;
148    for h in hits {
149        writeln!(
150            out,
151            "{:<40} {:<19} {:>7.3} {}",
152            h.session_id,
153            fmt_ts(h.ts_ms),
154            h.score,
155            h.snippet
156        )?;
157    }
158    Ok(out)
159}
160
161fn parse_since(raw: Option<&str>) -> Result<Option<u64>> {
162    let Some(raw) = raw else { return Ok(None) };
163    let days = raw.trim_end_matches('d').parse::<u64>()?;
164    let now = std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH)?;
165    Ok(Some(
166        (now.as_millis() as u64).saturating_sub(days.saturating_mul(86_400_000)),
167    ))
168}